Spaces:
Runtime error
Runtime error
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 | |