cvqluu / nn-similarity-diarization Goto Github PK

def __init__(self, input_size=256, hidden_size=256, num_layers=2):
    super(LSTMSimilarity, self).__init__()
    self.lstm = nn.LSTM(input_size,
                        hidden_size,
                        num_layers=num_layers,
                        bidirectional=True,
                        batch_first=True)
    self.fc1 = nn.Linear(hidden_size*2, 64)
    self.nl = nn.ReLU(inplace=True)
    self.fc2 = nn.Linear(64, 1)

def forward(self, x):
    self.lstm.flatten_parameters()
    x, _ = self.lstm(x)
    x = self.fc1(x)
    x = self.nl(x)
    x = self.fc2(x).squeeze(2)
    return x

rttm_cols.append(rttm_cols[3].astype(float) + rttm_cols[4].astype(float))#起始时间+duration=结束时间
recording_ids = sorted(set(segment_cols[1])) # recording_id 如iaaa
events0 = np.array(segment_cols[2:4]).astype(float).transpose() #segment起止时间
events1 = np.vstack([rttm_cols[3].astype(float), rttm_cols[-1]]).transpose() #ref起止时间（groundtruth)

rec_batches = []

for rec_id in tqdm(recording_ids): #tqdm是进度条
    seg_indexes = segment_cols[1] == rec_id  #recording_id iaaa ==优先级高
    rttm_indexes = rttm_cols[1] == rec_id
    ev0 = events0[seg_indexes]   #rec_id对应的segment里每段的起止时间
    ev1 = events1[rttm_indexes]  #rec_id对应的rffm里每段音频的起止时间
    ev1_labels = rttm_cols[7][rttm_indexes] #rec_id对应的rffm里每段音频的speaker
    ev0_labels = assign_overlaps(evnew, ev1, ev1_labels)
    ev0_labels = ['{}_{}'.format(rec_id, l) for l in ev0_labels]  #形成speaker_id,如iaaa_A
    batch = (segment_cols[0][seg_indexes], ev0_labels, vec_paths[seg_indexes], segment_rows[seg_indexes])
    rec_batches.append(batch)

return recording_ids, rec_batches