| | |
|---|
| | | import numpy as np |
|---|
| | | import torch |
|---|
| | | import torch.multiprocessing |
|---|
| | | import torch.nn.functional as F |
|---|
| | | from itertools import combinations |
|---|
| | | from itertools import permutations |
|---|
| | | |
|---|
| | | |
|---|
| | | def generate_mapping_dict(max_speaker_num=6, max_olp_speaker_num=3): |
|---|
| | | all_kinds = [] |
|---|
| | | all_kinds.append(0) |
|---|
| | | for i in range(max_olp_speaker_num): |
|---|
| | | selected_num = i + 1 |
|---|
| | | coms = np.array(list(combinations(np.arange(max_speaker_num), selected_num))) |
|---|
| | | for com in coms: |
|---|
| | | tmp = np.zeros(max_speaker_num) |
|---|
| | | tmp[com] = 1 |
|---|
| | | item = int(raw_dec_trans(tmp.reshape(1, -1), max_speaker_num)[0]) |
|---|
| | | all_kinds.append(item) |
|---|
| | | all_kinds_order = sorted(all_kinds) |
|---|
| | | |
|---|
| | | mapping_dict = {} |
|---|
| | | mapping_dict['dec2label'] = {} |
|---|
| | | mapping_dict['label2dec'] = {} |
|---|
| | | for i in range(len(all_kinds_order)): |
|---|
| | | dec = all_kinds_order[i] |
|---|
| | | mapping_dict['dec2label'][dec] = i |
|---|
| | | mapping_dict['label2dec'][i] = dec |
|---|
| | | oov_id = len(all_kinds_order) |
|---|
| | | mapping_dict['oov'] = oov_id |
|---|
| | | return mapping_dict |
|---|
| | | |
|---|
| | | |
|---|
| | | def raw_dec_trans(x, max_speaker_num): |
|---|
| | | num_list = [] |
|---|
| | | for i in range(max_speaker_num): |
|---|
| | | num_list.append(x[:, i]) |
|---|
| | | base = 1 |
|---|
| | | T = x.shape[0] |
|---|
| | | res = np.zeros((T)) |
|---|
| | | for num in num_list: |
|---|
| | | res += num * base |
|---|
| | | base = base * 2 |
|---|
| | | return res |
|---|
| | | |
|---|
| | | |
|---|
| | | def mapping_func(num, mapping_dict): |
|---|
| | | if num in mapping_dict['dec2label'].keys(): |
|---|
| | | label = mapping_dict['dec2label'][num] |
|---|
| | | else: |
|---|
| | | label = mapping_dict['oov'] |
|---|
| | | return label |
|---|
| | | |
|---|
| | | |
|---|
| | | def dec_trans(x, max_speaker_num, mapping_dict): |
|---|
| | | num_list = [] |
|---|
| | | for i in range(max_speaker_num): |
|---|
| | | num_list.append(x[:, i]) |
|---|
| | | base = 1 |
|---|
| | | T = x.shape[0] |
|---|
| | | res = np.zeros((T)) |
|---|
| | | for num in num_list: |
|---|
| | | res += num * base |
|---|
| | | base = base * 2 |
|---|
| | | res = np.array([mapping_func(i, mapping_dict) for i in res]) |
|---|
| | | return res |
|---|
| | | |
|---|
| | | |
|---|
| | | def create_powerlabel(label, mapping_dict, max_speaker_num=6, max_olp_speaker_num=3): |
|---|
| | | T, C = label.shape |
|---|
| | | padding_label = np.zeros((T, max_speaker_num)) |
|---|
| | | padding_label[:, :C] = label |
|---|
| | | out_label = dec_trans(padding_label, max_speaker_num, mapping_dict) |
|---|
| | | out_label = torch.from_numpy(out_label) |
|---|
| | | return out_label |
|---|
| | | |
|---|
| | | |
|---|
| | | def generate_perm_pse(label, n_speaker, mapping_dict, max_speaker_num, max_olp_speaker_num=3): |
|---|
| | | perms = np.array(list(permutations(range(n_speaker)))).astype(np.float32) |
|---|
| | | perms = torch.from_numpy(perms).to(label.device).to(torch.int64) |
|---|
| | | perm_labels = [label[:, perm] for perm in perms] |
|---|
| | | perm_pse_labels = [create_powerlabel(perm_label.cpu().numpy(), mapping_dict, max_speaker_num). |
|---|
| | | to(perm_label.device, non_blocking=True) for perm_label in perm_labels] |
|---|
| | | return perm_labels, perm_pse_labels |
|---|
| | | |
|---|
| | | |
|---|
| | | def generate_min_pse(label, n_speaker, mapping_dict, max_speaker_num, pse_logit, max_olp_speaker_num=3): |
|---|
| | | perm_labels, perm_pse_labels = generate_perm_pse(label, n_speaker, mapping_dict, max_speaker_num, |
|---|
| | | max_olp_speaker_num=max_olp_speaker_num) |
|---|
| | | losses = [F.cross_entropy(input=pse_logit, target=perm_pse_label.to(torch.long)) * len(pse_logit) |
|---|
| | | for perm_pse_label in perm_pse_labels] |
|---|
| | | loss = torch.stack(losses) |
|---|
| | | min_index = torch.argmin(loss) |
|---|
| | | selected_perm_label, selected_pse_label = perm_labels[min_index], perm_pse_labels[min_index] |
|---|
| | | return selected_perm_label, selected_pse_label |
|---|
