wehub-resource-sync
ba4be087d5
Create PR to main with cherry-pick from release / cherry-pick (push) Failing after 0s
CICD NeMo / pre-flight (push) Failing after 0s
CICD NeMo / configure (push) Has been skipped
Build, validate, and release Neural Modules / pre-flight (push) Failing after 1s
CICD NeMo / code-linting (push) Has been skipped
Build, validate, and release Neural Modules / release (push) Has been skipped
Build, validate, and release Neural Modules / release-summary (push) Has been cancelled
CICD NeMo / cicd-test-container-build (push) Has been cancelled
CICD NeMo / cicd-import-tests (push) Has been cancelled
CICD NeMo / L0_Setup_Test_Data_And_Models (push) Has been cancelled
CICD NeMo / cicd-main-unit-tests (push) Has been cancelled
CICD NeMo / cicd-main-speech (push) Has been cancelled
CICD NeMo / Nemo_CICD_Test (push) Has been cancelled
CICD NeMo / Coverage (e2e) (push) Has been cancelled
CICD NeMo / Coverage (unit-test) (push) Has been cancelled
CodeQL / Analyze (python) (push) Has been cancelled
CICD NeMo / cicd-wait-in-queue (push) Has been cancelled
1082 lines
44 KiB
Python
1082 lines
44 KiB
Python
# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import os
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
from scipy.optimize import linear_sum_assignment as scipy_linear_sum_assignment
|
|
|
|
from nemo.collections.asr.data.audio_to_label import repeat_signal
|
|
from nemo.collections.asr.parts.utils.longform_clustering import LongFormSpeakerClustering
|
|
from nemo.collections.asr.parts.utils.offline_clustering import (
|
|
SpeakerClustering,
|
|
get_scale_interpolated_embs,
|
|
getCosAffinityMatrix,
|
|
getKneighborsConnections,
|
|
split_input_data,
|
|
)
|
|
from nemo.collections.asr.parts.utils.online_clustering import (
|
|
OnlineSpeakerClustering,
|
|
get_closest_embeddings,
|
|
get_merge_quantity,
|
|
get_minimal_indices,
|
|
merge_vectors,
|
|
run_reducer,
|
|
stitch_cluster_labels,
|
|
)
|
|
from nemo.collections.asr.parts.utils.optimization_utils import LinearSumAssignmentSolver
|
|
from nemo.collections.asr.parts.utils.optimization_utils import linear_sum_assignment as nemo_linear_sum_assignment
|
|
from nemo.collections.asr.parts.utils.speaker_utils import (
|
|
OnlineSegmentor,
|
|
check_ranges,
|
|
fl2int,
|
|
get_new_cursor_for_update,
|
|
get_online_segments_from_slices,
|
|
get_online_subsegments_from_buffer,
|
|
get_speech_labels_for_update,
|
|
get_sub_range_list,
|
|
get_subsegments,
|
|
get_subsegments_scriptable,
|
|
get_target_sig,
|
|
int2fl,
|
|
is_overlap,
|
|
merge_float_intervals,
|
|
merge_int_intervals,
|
|
tensor_to_list,
|
|
)
|
|
|
|
|
|
def check_range_values(target, source):
|
|
bool_list = []
|
|
for tgt, src in zip(target, source):
|
|
for x, y in zip(src, tgt):
|
|
bool_list.append(abs(x - y) < 1e-6)
|
|
return all(bool_list)
|
|
|
|
|
|
def check_labels(target, source):
|
|
bool_list = []
|
|
for x, y in zip(target, source):
|
|
bool_list.append(abs(x - y) < 1e-6)
|
|
return all(bool_list)
|
|
|
|
|
|
def matrix(mat, use_tensor=True, dtype=torch.long):
|
|
if use_tensor:
|
|
mat = torch.Tensor(mat).to(dtype)
|
|
else:
|
|
mat = np.array(mat)
|
|
return mat
|
|
|
|
|
|
def generate_orthogonal_embs(total_spks, perturb_sigma, emb_dim):
|
|
"""Generate a set of artificial orthogonal embedding vectors from random numbers"""
|
|
gaus = torch.randn(emb_dim, emb_dim)
|
|
_svd = torch.linalg.svd(gaus)
|
|
orth = _svd[0] @ _svd[2]
|
|
orth_embs = orth[:total_spks]
|
|
# Assert orthogonality
|
|
assert torch.abs(getCosAffinityMatrix(orth_embs) - torch.diag(torch.ones(total_spks))).sum() < 1e-4
|
|
return orth_embs
|
|
|
|
|
|
def generate_toy_data(
|
|
n_spks=2,
|
|
spk_dur=3,
|
|
emb_dim=192,
|
|
perturb_sigma=0.0,
|
|
ms_window=[1.5, 1.0, 0.5],
|
|
ms_shift=[0.75, 0.5, 0.25],
|
|
torch_seed=0,
|
|
):
|
|
torch.manual_seed(torch_seed)
|
|
spk_timestamps = [(spk_dur * k, spk_dur) for k in range(n_spks)]
|
|
emb_list, seg_list = [], []
|
|
multiscale_segment_counts = [0 for _ in range(len(ms_window))]
|
|
ground_truth = []
|
|
random_orthogonal_embs = generate_orthogonal_embs(n_spks, perturb_sigma, emb_dim)
|
|
for scale_idx, (window, shift) in enumerate(zip(ms_window, ms_shift)):
|
|
for spk_idx, (offset, dur) in enumerate(spk_timestamps):
|
|
segments_stt_dur = get_subsegments_scriptable(offset=offset, window=window, shift=shift, duration=dur)
|
|
segments = [[x[0], x[0] + x[1]] for x in segments_stt_dur]
|
|
emb_cent = random_orthogonal_embs[spk_idx, :]
|
|
emb = emb_cent.tile((len(segments), 1)) + 0.1 * torch.rand(len(segments), emb_dim)
|
|
seg_list.extend(segments)
|
|
emb_list.append(emb)
|
|
if emb.shape[0] == 0:
|
|
import ipdb
|
|
|
|
ipdb.set_trace()
|
|
multiscale_segment_counts[scale_idx] += emb.shape[0]
|
|
|
|
if scale_idx == len(multiscale_segment_counts) - 1:
|
|
ground_truth.extend([spk_idx] * emb.shape[0])
|
|
|
|
emb_tensor = torch.concat(emb_list)
|
|
multiscale_segment_counts = torch.tensor(multiscale_segment_counts)
|
|
segm_tensor = torch.tensor(seg_list)
|
|
multiscale_weights = torch.ones(len(ms_window)).unsqueeze(0)
|
|
ground_truth = torch.tensor(ground_truth)
|
|
return emb_tensor, segm_tensor, multiscale_segment_counts, multiscale_weights, spk_timestamps, ground_truth
|
|
|
|
|
|
class TestDiarizationSequneceUtilFunctions:
|
|
"""Tests diarization and speaker-task related utils."""
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize("Y", [[3, 3, 3, 4, 4, 5], [100, 100, 100, 104, 104, 1005]])
|
|
@pytest.mark.parametrize("target", [[0, 0, 0, 1, 1, 2]])
|
|
@pytest.mark.parametrize("offset", [1, 10])
|
|
def test_minimal_index_ex2(self, Y, target, offset):
|
|
Y = torch.tensor(Y)
|
|
target = torch.tensor(target)
|
|
min_Y = get_minimal_indices(Y)
|
|
assert check_labels(target, min_Y)
|
|
min_Y = get_minimal_indices(Y + offset)
|
|
assert check_labels(target, min_Y)
|
|
|
|
@pytest.mark.parametrize("Y", [[4, 0, 0, 5, 4, 5], [14, 12, 12, 19, 14, 19]])
|
|
@pytest.mark.parametrize("target", [[1, 0, 0, 2, 1, 2]])
|
|
@pytest.mark.parametrize("offset", [1, 10])
|
|
def test_minimal_index_ex2(self, Y, target, offset):
|
|
Y = torch.tensor(Y)
|
|
target = torch.tensor(target)
|
|
min_Y = get_minimal_indices(Y)
|
|
assert check_labels(target, min_Y)
|
|
min_Y = get_minimal_indices(Y + offset)
|
|
assert check_labels(target, min_Y)
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize("N", [2, 4, 16, 64])
|
|
def test_minimal_index_same(self, N):
|
|
Y = matrix([0] * N + [1] * N + [2] * N)
|
|
min_Y = get_minimal_indices(Y)
|
|
target = matrix([0] * N + [1] * N + [2] * N)
|
|
assert check_labels(target, min_Y)
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize("N", [2, 4, 16, 64])
|
|
def test_stitch_cluster_labels_label_switch(self, N):
|
|
Y_old = matrix([0] * N)
|
|
Y_new = matrix([0] * N) + 1
|
|
target = matrix([0] * N)
|
|
result = stitch_cluster_labels(Y_old, Y_new)
|
|
assert check_labels(target, result)
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize("N", [2, 4, 16, 64])
|
|
def test_stitch_cluster_labels_label_many_to_one(self, N):
|
|
Y_old = matrix(np.arange(N).tolist())
|
|
Y_new = matrix([0] * N)
|
|
target = matrix([0] * N)
|
|
result = stitch_cluster_labels(Y_old, Y_new)
|
|
assert check_labels(target, result)
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize("N", [2, 4, 16, 64])
|
|
def test_stitch_cluster_labels_label_one_to_many(self, N):
|
|
Y_old = matrix(np.arange(N).tolist())
|
|
Y_new = matrix([k for k in range(N)])
|
|
target = matrix([k for k in range(N)])
|
|
result = stitch_cluster_labels(Y_old, Y_new)
|
|
assert check_labels(target, result)
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize("N", [2, 4, 16, 64])
|
|
def test_stitch_cluster_labels_one_label_replaced(self, N):
|
|
Y_old = matrix([0] * N + [1] * N + [2] * N)
|
|
Y_new = matrix([1] * N + [2] * N + [3] * N)
|
|
target = matrix([0] * N + [1] * N + [2] * N)
|
|
result = stitch_cluster_labels(Y_old, Y_new)
|
|
assert check_labels(target, result)
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize("N", [2, 4, 16, 64])
|
|
def test_stitch_cluster_labels_confusion_error(self, N):
|
|
Y_old = matrix([0] * N + [1] * (N - 1) + [2] * (N + 1))
|
|
Y_new = matrix([1] * N + [2] * N + [3] * N)
|
|
target = matrix([0] * N + [1] * N + [2] * N)
|
|
result = stitch_cluster_labels(Y_old, Y_new)
|
|
assert check_labels(target, result)
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize("N", [2, 256])
|
|
def test_stitch_cluster_labels_speaker_more_speakers(self, N):
|
|
Y_old = matrix([0] * N + [1] * (N - 1) + [2] * (N + 1) + [0, 0, 0])
|
|
Y_new = matrix([1] * N + [0] * N + [2] * N + [4, 5, 6])
|
|
target = matrix([0] * N + [1] * N + [2] * N + [3, 4, 5])
|
|
result = stitch_cluster_labels(Y_old, Y_new)
|
|
assert check_labels(target, result)
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize("N", [2, 256])
|
|
def test_stitch_cluster_labels_speaker_longer_sequence(self, N):
|
|
Y_old = matrix([0] * N + [1] * N + [2] * N + [0, 0, 0] * N)
|
|
Y_new = matrix([1] * N + [2] * N + [0] * N + [1, 2, 3, 1, 2, 3] * N)
|
|
target = matrix([0] * N + [1] * N + [2] * N + [0, 1, 3, 0, 1, 3] * N)
|
|
result = stitch_cluster_labels(Y_old, Y_new)
|
|
assert check_labels(target, result)
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize("n_spks", [2, 3, 4, 5])
|
|
@pytest.mark.parametrize("merge_quantity", [2, 3])
|
|
def test_embedding_merger(self, n_spks, merge_quantity):
|
|
em, ts, mc, mw, spk_ts, gt = generate_toy_data(n_spks, spk_dur=5, perturb_sigma=10)
|
|
em_s, ts_s = split_input_data(em, ts, mc)
|
|
target_speaker_index = 0
|
|
pre_clus_labels = gt
|
|
ndx = torch.where(pre_clus_labels == target_speaker_index)[0]
|
|
pre_embs = em_s[-1]
|
|
affinity_mat = getCosAffinityMatrix(pre_embs)
|
|
cmat = affinity_mat[:, ndx][ndx, :]
|
|
# Check the dimension of the selected affinity values
|
|
assert cmat.shape[0] == cmat.shape[1] == torch.sum(pre_clus_labels == target_speaker_index).item()
|
|
index_2d, rest_inds = get_closest_embeddings(cmat, merge_quantity)
|
|
# Check the most closest affinity value
|
|
assert torch.max(cmat.sum(0)) == cmat.sum(0)[index_2d[0]]
|
|
spk_cluster_labels, emb_ndx = pre_clus_labels[ndx], pre_embs[ndx]
|
|
merged_embs, merged_clus_labels = merge_vectors(index_2d, emb_ndx, spk_cluster_labels)
|
|
# Check the number of merged embeddings and labels
|
|
assert (torch.sum(gt == target_speaker_index).item() - merge_quantity) == merged_clus_labels.shape[0]
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize("n_spks", [1, 8])
|
|
@pytest.mark.parametrize("spk_dur", [0.2, 0.25, 0.5, 1, 10])
|
|
def test_cosine_affinity_calculation(self, n_spks, spk_dur):
|
|
em, ts, mc, mw, spk_ts, gt = generate_toy_data(n_spks=n_spks, spk_dur=spk_dur)
|
|
em_s, ts_s = split_input_data(em, ts, mc)
|
|
affinity_mat = getCosAffinityMatrix(em_s[-1])
|
|
# affinity_mat should not contain any nan element
|
|
assert torch.any(torch.isnan(affinity_mat)) == False
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize("n_spks", [1, 8])
|
|
@pytest.mark.parametrize("spk_dur", [0.2, 0.25, 0.5, 1, 10])
|
|
def test_cosine_affinity_calculation_scale_interpol(self, n_spks, spk_dur):
|
|
em, ts, mc, mw, spk_ts, gt = generate_toy_data(n_spks=n_spks, spk_dur=spk_dur)
|
|
em_s, ts_s = split_input_data(em, ts, mc)
|
|
embs, _ = get_scale_interpolated_embs(mw, em_s, ts_s)
|
|
affinity_mat = getCosAffinityMatrix(embs)
|
|
# affinity_mat should not contain any nan element
|
|
assert torch.any(torch.isnan(affinity_mat)) == False
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize("n_spks", [4, 5, 6])
|
|
@pytest.mark.parametrize("target_speaker_index", [0, 1, 2])
|
|
@pytest.mark.parametrize("merge_quantity", [2, 3])
|
|
def test_embedding_reducer(self, n_spks, target_speaker_index, merge_quantity):
|
|
em, ts, mc, mw, spk_ts, gt = generate_toy_data(n_spks=n_spks, spk_dur=10)
|
|
em_s, ts_s = split_input_data(em, ts, mc)
|
|
merged_embs, merged_clus_labels, _ = run_reducer(
|
|
pre_embs=em_s[-1],
|
|
target_spk_idx=target_speaker_index,
|
|
merge_quantity=merge_quantity,
|
|
pre_clus_labels=gt,
|
|
)
|
|
assert (torch.sum(gt == target_speaker_index).item() - merge_quantity) == merged_clus_labels.shape[0]
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize("ntbr", [3])
|
|
@pytest.mark.parametrize("pcl", [torch.tensor([0] * 70 + [1] * 32)])
|
|
@pytest.mark.parametrize("mspb", [25])
|
|
def test_merge_scheduler_2clus(self, ntbr, pcl, mspb):
|
|
class_target_vol = get_merge_quantity(
|
|
num_to_be_removed=ntbr,
|
|
pre_clus_labels=pcl,
|
|
min_count_per_cluster=mspb,
|
|
)
|
|
assert all(class_target_vol == torch.tensor([3, 0]))
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize("ntbr", [3])
|
|
@pytest.mark.parametrize("pcl", [torch.tensor([0] * 80 + [1] * 35 + [2] * 32)])
|
|
@pytest.mark.parametrize("mspb", [0, 25])
|
|
def test_merge_scheduler_3clus(self, ntbr, pcl, mspb):
|
|
class_target_vol = get_merge_quantity(
|
|
num_to_be_removed=ntbr,
|
|
pre_clus_labels=pcl,
|
|
min_count_per_cluster=mspb,
|
|
)
|
|
assert all(class_target_vol == torch.tensor([3, 0, 0]))
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize("ntbr", [132 - 45])
|
|
@pytest.mark.parametrize("pcl", [torch.tensor([2] * 70 + [0] * 32 + [1] * 27 + [3] * 3)])
|
|
@pytest.mark.parametrize("mspb", [3, 10])
|
|
def test_merge_scheduler_4clus_shuff(self, ntbr, pcl, mspb):
|
|
class_target_vol = get_merge_quantity(
|
|
num_to_be_removed=ntbr,
|
|
pre_clus_labels=pcl,
|
|
min_count_per_cluster=mspb,
|
|
)
|
|
assert all(class_target_vol == torch.tensor([18, 13, 56, 0]))
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize("ntbr", [3])
|
|
@pytest.mark.parametrize("pcl", [torch.tensor([0] * 5 + [1] * 4 + [2] * 3)])
|
|
@pytest.mark.parametrize("mspb", [0, 2])
|
|
def test_merge_scheduler_3clus(self, ntbr, pcl, mspb):
|
|
class_target_vol = get_merge_quantity(
|
|
num_to_be_removed=ntbr,
|
|
pre_clus_labels=pcl,
|
|
min_count_per_cluster=mspb,
|
|
)
|
|
assert all(class_target_vol == torch.tensor([2, 1, 0]))
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize("ntbr", [2])
|
|
@pytest.mark.parametrize("pcl", [torch.tensor([0] * 7 + [1] * 5 + [2] * 3 + [3] * 5)])
|
|
@pytest.mark.parametrize("mspb", [2])
|
|
def test_merge_scheduler_3clus_repeat(self, ntbr, pcl, mspb):
|
|
class_target_vol = get_merge_quantity(
|
|
num_to_be_removed=ntbr,
|
|
pre_clus_labels=pcl,
|
|
min_count_per_cluster=mspb,
|
|
)
|
|
assert all(class_target_vol == torch.tensor([2, 0, 0, 0]))
|
|
|
|
|
|
class TestClassExport:
|
|
@pytest.mark.unit
|
|
def test_online_segmentor_class_export(self):
|
|
_OnlineSegmentor = torch.jit.script(OnlineSegmentor)
|
|
online_segmentor = _OnlineSegmentor(sample_rate=16000)
|
|
assert isinstance(online_segmentor, OnlineSegmentor)
|
|
|
|
@pytest.mark.unit
|
|
def test_online_segmentor_instance_export(self):
|
|
online_segmentor = OnlineSegmentor(sample_rate=16000)
|
|
online_segmentor = torch.jit.script(online_segmentor)
|
|
isinstance(online_segmentor, torch.jit._script.RecursiveScriptClass)
|
|
|
|
@pytest.mark.unit
|
|
def test_online_speaker_clustering_instance_export(self):
|
|
online_clus = OnlineSpeakerClustering(
|
|
max_num_speakers=8,
|
|
max_rp_threshold=0.15,
|
|
sparse_search_volume=30,
|
|
history_buffer_size=150,
|
|
current_buffer_size=150,
|
|
cuda=True,
|
|
)
|
|
online_clus = torch.jit.script(online_clus)
|
|
isinstance(online_clus, torch.jit._script.RecursiveScriptClass)
|
|
|
|
@pytest.mark.unit
|
|
def test_online_speaker_clustering_instance_export(self):
|
|
offline_speaker_clustering = SpeakerClustering(maj_vote_spk_count=False, min_samples_for_nmesc=0, cuda=True)
|
|
offline_speaker_clustering = torch.jit.script(offline_speaker_clustering)
|
|
isinstance(offline_speaker_clustering, torch.jit._script.RecursiveScriptClass)
|
|
|
|
|
|
class TestGetSubsegments:
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize(
|
|
"offset, window, shift, duration, min_subsegment_duration, decimals, use_asr_style_frame_count, sample_rate, feat_per_sec, expected",
|
|
[
|
|
(12.05, 1.5, 0.75, 2.4, 0.01, 2, False, 16000, 100, [[12.05, 1.5], [12.8, 1.5], [13.55, 0.9]]),
|
|
(0, 1.0, 0.5, 0.4, 0.01, 2, False, 16000, 100, [[0, 0.4]]),
|
|
(0, 2.0, 1.0, 1.5, 0.5, 2, False, 16000, 100, [[0, 1.5]]),
|
|
(
|
|
10,
|
|
1.5,
|
|
0.75,
|
|
4.5,
|
|
0.5,
|
|
2,
|
|
False,
|
|
16000,
|
|
100,
|
|
[[10, 1.5], [10.75, 1.5], [11.5, 1.5], [12.25, 1.5], [13.0, 1.5]],
|
|
),
|
|
(0, 1.5, 0.5, 0.3, 0.01, 2, True, 16000, 100, [[0, 0.3]]),
|
|
],
|
|
)
|
|
def test_get_subsegments(
|
|
self,
|
|
offset,
|
|
window,
|
|
shift,
|
|
duration,
|
|
min_subsegment_duration,
|
|
decimals,
|
|
use_asr_style_frame_count,
|
|
sample_rate,
|
|
feat_per_sec,
|
|
expected,
|
|
):
|
|
|
|
for is_scriptable in [True, False]:
|
|
if is_scriptable:
|
|
result = get_subsegments_scriptable(
|
|
offset=offset,
|
|
window=window,
|
|
shift=shift,
|
|
duration=duration,
|
|
)
|
|
else:
|
|
result = get_subsegments(
|
|
offset=offset,
|
|
window=window,
|
|
shift=shift,
|
|
duration=duration,
|
|
min_subsegment_duration=min_subsegment_duration,
|
|
decimals=decimals,
|
|
use_asr_style_frame_count=use_asr_style_frame_count,
|
|
sample_rate=sample_rate,
|
|
feat_per_sec=feat_per_sec,
|
|
)
|
|
result_round = []
|
|
for subsegment in result:
|
|
result_round.append([round(x, decimals) for x in subsegment])
|
|
assert result_round == expected
|
|
|
|
@pytest.mark.unit
|
|
def test_min_subsegment_duration_filtering(self):
|
|
result = get_subsegments(
|
|
offset=0,
|
|
window=1.5,
|
|
shift=0.5,
|
|
duration=3,
|
|
min_subsegment_duration=2.0,
|
|
decimals=2,
|
|
use_asr_style_frame_count=False,
|
|
)
|
|
expected = [] # Only subsegments meeting the duration filter should remain
|
|
assert result == expected
|
|
|
|
@pytest.mark.unit
|
|
def test_zero_duration(self):
|
|
result = get_subsegments(
|
|
offset=0,
|
|
window=1.0,
|
|
shift=0.5,
|
|
duration=0,
|
|
min_subsegment_duration=0.01,
|
|
decimals=2,
|
|
use_asr_style_frame_count=False,
|
|
)
|
|
assert result == []
|
|
|
|
@pytest.mark.unit
|
|
def test_edge_case_short_slice(self):
|
|
result = get_subsegments(
|
|
offset=0,
|
|
window=0.5,
|
|
shift=0.25, # Shift larger than duration
|
|
duration=0.25,
|
|
min_subsegment_duration=0.01,
|
|
decimals=2,
|
|
use_asr_style_frame_count=False,
|
|
)
|
|
assert result == [[0.0, 0.25]]
|
|
|
|
|
|
class TestDiarizationSegmentationUtils:
|
|
"""
|
|
Test segmentation util functions
|
|
"""
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize(
|
|
"intervals",
|
|
[
|
|
[[1, 4], [2, 6], [8, 10], [15, 18]],
|
|
[[8, 10], [15, 18], [2, 6], [1, 3]],
|
|
[[8, 10], [15, 18], [2, 6], [1, 3], [3, 5]],
|
|
[[8, 10], [8, 8], [15, 18], [2, 6], [1, 6], [2, 4]],
|
|
],
|
|
)
|
|
@pytest.mark.parametrize("target", [[[1, 6], [8, 10], [15, 18]]])
|
|
def test_merge_int_intervals_ex1(self, intervals, target):
|
|
merged = merge_int_intervals(intervals)
|
|
assert check_range_values(target, merged)
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize(
|
|
"intervals",
|
|
[
|
|
[[6, 8], [0, 9], [2, 4], [4, 7]],
|
|
[[0, 9], [6, 8], [4, 7], [2, 4]],
|
|
[[0, 4], [0, 0], [4, 9], [2, 4]],
|
|
[[6, 8], [2, 8], [0, 3], [3, 4], [4, 5], [5, 9]],
|
|
],
|
|
)
|
|
@pytest.mark.parametrize("target", [[[0, 9]]])
|
|
def test_merge_int_intervals_ex2(self, intervals, target):
|
|
merged = merge_int_intervals(intervals)
|
|
assert check_range_values(target, merged)
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize("intervals", [[[0, 1], [1, 9]], [[0, 0], [0, 9]], [[0, 9], [0, 9]]])
|
|
@pytest.mark.parametrize("target", [[[0, 9]]])
|
|
def test_merge_int_intervals_edge_test(self, intervals, target):
|
|
merged = merge_int_intervals(intervals)
|
|
assert check_range_values(target, merged)
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize("rangeA", [[1.0, 2.0]])
|
|
@pytest.mark.parametrize("rangeB", [[0.5, 1.5], [0.9999, 1.0001]])
|
|
def test_is_overlap_true(self, rangeA, rangeB):
|
|
assert is_overlap(rangeA, rangeB)
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize("rangeA", [[1.0, 2.0]])
|
|
@pytest.mark.parametrize("rangeB", [[2.0, 2.5], [-1.0, 1.00]])
|
|
def test_is_overlap_false(self, rangeA, rangeB):
|
|
assert not is_overlap(rangeA, rangeB)
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize("x", [1.0, 2.3456])
|
|
@pytest.mark.parametrize("decimals", [1, 2, 3, 4])
|
|
def test_fl2int(self, x, decimals):
|
|
assert fl2int(x, decimals) == round(x * 10**decimals, 0)
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize("x", [1234])
|
|
@pytest.mark.parametrize(
|
|
"decimals",
|
|
[
|
|
1,
|
|
2,
|
|
3,
|
|
4,
|
|
],
|
|
)
|
|
def test_int2fl(self, x, decimals):
|
|
assert abs(int2fl(x, decimals) - round(x / (10**decimals), decimals)) < (10 ** -(decimals + 1))
|
|
|
|
@pytest.mark.unit
|
|
def test_merge_float_intervals_edge_margin_test(self):
|
|
intervals = [[0.0, 1.0], [1.0, 2.0]]
|
|
|
|
target_0 = [[0.0, 2.0]]
|
|
merged_0 = merge_float_intervals(intervals, margin=0)
|
|
assert check_range_values(target_0, merged_0)
|
|
|
|
target_1 = [[0.0, 1.0], [1.0, 2.0]]
|
|
merged_1 = merge_float_intervals(intervals, margin=1)
|
|
assert check_range_values(target_1, merged_1)
|
|
|
|
target_2 = [[0.0, 1.0], [1.0, 2.0]]
|
|
merged_2 = merge_float_intervals(intervals, margin=2)
|
|
assert check_range_values(target_2, merged_2)
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize(
|
|
"intervals",
|
|
[
|
|
[[0.25, 1.7], [1.5, 3.0], [2.8, 5.0], [5.5, 10.0]],
|
|
[[0.25, 5.0], [5.5, 10.0], [1.5, 3.5]],
|
|
[[5.5, 8.05], [8.0, 10.0], [0.25, 5.0]],
|
|
[[0.25, 3.0], [1.5, 3.0], [5.5, 10.0], [2.8, 5.0]],
|
|
[[0.25, 1.7], [1.5, 3.0], [2.8, 5.0], [5.5, 10.0]],
|
|
],
|
|
)
|
|
@pytest.mark.parametrize("target", [[[0.25, 5.0], [5.5, 10.0]]])
|
|
def test_merge_float_overlaps(self, intervals, target):
|
|
merged = merge_float_intervals(intervals)
|
|
assert check_range_values(target, merged)
|
|
|
|
@pytest.mark.unit
|
|
def test_get_speech_labels_for_update(self):
|
|
frame_start = 3.0
|
|
buffer_end = 6.0
|
|
cumulative_speech_labels = torch.tensor([[0.0000, 3.7600]])
|
|
vad_timestamps = torch.tensor([[0.9600, 4.8400]])
|
|
cursor_for_old_segments = 1.0
|
|
speech_labels_for_update, cumulative_speech_labels = get_speech_labels_for_update(
|
|
frame_start,
|
|
buffer_end,
|
|
cumulative_speech_labels,
|
|
vad_timestamps,
|
|
cursor_for_old_segments,
|
|
)
|
|
assert (speech_labels_for_update - torch.tensor([[1.0000, 3.7600]])).sum() < 1e-8
|
|
assert (cumulative_speech_labels - torch.tensor([[0.9600, 4.8400]])).sum() < 1e-8
|
|
|
|
# Check if the ranges are containing faulty values
|
|
assert check_ranges(speech_labels_for_update)
|
|
assert check_ranges(cumulative_speech_labels)
|
|
|
|
@pytest.mark.unit
|
|
def test_get_online_subsegments_from_buffer(self):
|
|
torch.manual_seed(0)
|
|
sample_rate = 16000
|
|
speech_labels_for_update = torch.Tensor([[0.0000, 3.7600]])
|
|
audio_buffer = torch.randn(5 * sample_rate)
|
|
segment_indexes = []
|
|
window = 2.0
|
|
shift = 1.0
|
|
slice_length = int(window * sample_rate)
|
|
range_target = [[0.0, 2.0], [1.0, 3.0], [2.0, 3.76]]
|
|
sigs_list, sig_rangel_list, sig_indexes = get_online_subsegments_from_buffer(
|
|
buffer_start=0.0,
|
|
buffer_end=5.0,
|
|
sample_rate=sample_rate,
|
|
speech_labels_for_update=speech_labels_for_update,
|
|
audio_buffer=audio_buffer,
|
|
segment_indexes=segment_indexes,
|
|
window=window,
|
|
shift=shift,
|
|
)
|
|
assert check_range_values(target=range_target, source=sig_rangel_list)
|
|
for k, rg in enumerate(sig_rangel_list):
|
|
signal = get_target_sig(audio_buffer, rg[0], rg[1], slice_length, sample_rate)
|
|
if len(signal) < int(window * sample_rate):
|
|
signal = repeat_signal(signal, len(signal), slice_length)
|
|
assert len(signal) == int(slice_length), "Length mismatch"
|
|
assert (np.abs(signal - sigs_list[k])).sum() < 1e-8, "Audio stream mismatch"
|
|
assert (torch.tensor(sig_indexes) - torch.arange(len(range_target))).sum() < 1e-8, "Segment index mismatch"
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize("frame_start", [3.0])
|
|
@pytest.mark.parametrize("segment_range_ts", [[[0.0, 2.0]]])
|
|
@pytest.mark.parametrize("gt_cursor_for_old_segments", [3.0])
|
|
@pytest.mark.parametrize("gt_cursor_index", [1])
|
|
def test_get_new_cursor_for_update_mulsegs_ex1(
|
|
self, frame_start, segment_range_ts, gt_cursor_for_old_segments, gt_cursor_index
|
|
):
|
|
cursor_for_old_segments, cursor_index = get_new_cursor_for_update(frame_start, segment_range_ts)
|
|
assert cursor_for_old_segments == gt_cursor_for_old_segments
|
|
assert cursor_index == gt_cursor_index
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize("target_range", [[1.0, 4.0]])
|
|
@pytest.mark.parametrize(
|
|
"source_range_list", [[[2.0, 3.0], [3.0, 4.0]], [[0.0, 2.0], [3.0, 5.0]], [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]]
|
|
)
|
|
def get_sub_range_list(self, target_range, source_range_list):
|
|
sub_range_list = get_sub_range_list(target_range, source_range_list)
|
|
assert sub_range_list == [[2.0, 3.0], [3.0, 4.0]]
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize("source_range_list", [[[0.0, 2.0]], [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]])
|
|
def test_tensor_to_list(self, source_range_list):
|
|
a_range_tensor = torch.tensor(source_range_list)
|
|
converted_list = tensor_to_list(a_range_tensor)
|
|
assert source_range_list == converted_list
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize(
|
|
"buffer_start, buffer_end, subsegments, ind_offset, window, sample_rate",
|
|
[
|
|
(0.0, 2.0, [[0.5, 1.0], [1.5, 2.0]], 0, 0.1, 16000),
|
|
(0.0, 5.0, [[0.5, 2.5], [2.7, 5.0]], 0, 1.0, 16000),
|
|
],
|
|
)
|
|
def test_get_online_segments_from_slices(
|
|
self, buffer_start, buffer_end, subsegments, ind_offset, window, sample_rate
|
|
):
|
|
sig = torch.randn(int(sample_rate * buffer_end))
|
|
ind_offset, sigs_list, sig_rangel_list, sig_indexes = get_online_segments_from_slices(
|
|
sig, buffer_start, buffer_end, subsegments, ind_offset, window, sample_rate
|
|
)
|
|
assert ind_offset == 2
|
|
assert len(sigs_list) == 2
|
|
assert len(sig_rangel_list) == 2
|
|
assert len(sig_indexes) == 2
|
|
|
|
|
|
class TestClusteringUtilFunctions:
|
|
@pytest.mark.parametrize("p_value", [1, 5, 9])
|
|
@pytest.mark.parametrize("N", [9, 20])
|
|
@pytest.mark.parametrize("mask_method", ['binary', 'sigmoid', 'drop'])
|
|
def test_get_k_neighbors_connections(self, p_value: int, N: int, mask_method: str, seed=0):
|
|
torch.manual_seed(seed)
|
|
random_mat = torch.rand(N, N)
|
|
affinity_mat = 0.5 * (random_mat + random_mat.T)
|
|
affinity_mat = affinity_mat / torch.max(affinity_mat)
|
|
binarized_affinity_mat = getKneighborsConnections(affinity_mat, p_value, mask_method)
|
|
if mask_method == 'binary':
|
|
assert all(binarized_affinity_mat.sum(dim=0) == float(p_value))
|
|
elif mask_method == 'sigmoid':
|
|
assert all(binarized_affinity_mat.sum(dim=0) <= float(p_value))
|
|
elif mask_method == 'drop':
|
|
assert all(binarized_affinity_mat.sum(dim=0) <= float(p_value))
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize("Y_aggr", [torch.tensor([0, 1, 0, 1])])
|
|
@pytest.mark.parametrize("chunk_cluster_count, embeddings_per_chunk", [(2, 50)])
|
|
@pytest.mark.parametrize("window_range_list", [[[0, 1], [1, 2], [2, 3], [3, 4]]])
|
|
@pytest.mark.parametrize(
|
|
"absolute_merge_mapping",
|
|
[[[torch.tensor([]), torch.tensor([0, 2])], [torch.tensor([]), torch.tensor([1, 3])]]],
|
|
)
|
|
@pytest.mark.parametrize("org_len", [4])
|
|
def test_unpack_labels(
|
|
self, Y_aggr, window_range_list, absolute_merge_mapping, chunk_cluster_count, embeddings_per_chunk, org_len
|
|
):
|
|
expected_result = Y_aggr
|
|
longform_speaker_clustering = LongFormSpeakerClustering(cuda=False)
|
|
output = longform_speaker_clustering.unpack_labels(Y_aggr, window_range_list, absolute_merge_mapping, org_len)
|
|
assert torch.equal(output, expected_result)
|
|
|
|
|
|
class TestSpeakerClustering:
|
|
"""
|
|
Test speaker clustering module
|
|
"""
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize("cuda", [True, False])
|
|
def test_offline_clus_script_save_load(self, cuda):
|
|
exported_filename = 'speaker_clustering_script.pt'
|
|
speaker_clustering_python = SpeakerClustering(maj_vote_spk_count=False, cuda=cuda)
|
|
speaker_clustering_scripted_source = torch.jit.script(speaker_clustering_python)
|
|
torch.jit.save(speaker_clustering_scripted_source, exported_filename)
|
|
assert os.path.exists(exported_filename)
|
|
os.remove(exported_filename)
|
|
assert not os.path.exists(exported_filename)
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize("cuda", [True, False])
|
|
def test_online_clus_script_save_load(self, cuda):
|
|
exported_filename = 'speaker_clustering_script.pt'
|
|
speaker_clustering_python = OnlineSpeakerClustering(
|
|
max_num_speakers=8,
|
|
max_rp_threshold=0.15,
|
|
sparse_search_volume=30,
|
|
history_buffer_size=150,
|
|
current_buffer_size=150,
|
|
cuda=cuda,
|
|
)
|
|
speaker_clustering_scripted_source = torch.jit.script(speaker_clustering_python)
|
|
torch.jit.save(speaker_clustering_scripted_source, exported_filename)
|
|
assert os.path.exists(exported_filename)
|
|
os.remove(exported_filename)
|
|
assert not os.path.exists(exported_filename)
|
|
|
|
@pytest.mark.run_only_on('GPU')
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize("n_spks", [1, 2, 3, 4, 5, 6, 7])
|
|
@pytest.mark.parametrize("total_sec, SSV, perturb_sigma, seed", [(30, 10, 0.1, 0)])
|
|
@pytest.mark.parametrize("jit_script", [False, True])
|
|
def test_offline_speaker_clustering(self, n_spks, total_sec, SSV, perturb_sigma, seed, jit_script, cuda=True):
|
|
spk_dur = total_sec / n_spks
|
|
em, ts, mc, mw, spk_ts, gt = generate_toy_data(
|
|
n_spks=n_spks, spk_dur=spk_dur, perturb_sigma=perturb_sigma, torch_seed=seed
|
|
)
|
|
offline_speaker_clustering = SpeakerClustering(maj_vote_spk_count=False, cuda=cuda)
|
|
assert isinstance(offline_speaker_clustering, SpeakerClustering)
|
|
if jit_script:
|
|
offline_speaker_clustering = torch.jit.script(offline_speaker_clustering)
|
|
|
|
Y_out = offline_speaker_clustering.forward_infer(
|
|
embeddings_in_scales=em,
|
|
timestamps_in_scales=ts,
|
|
multiscale_segment_counts=mc,
|
|
multiscale_weights=mw,
|
|
oracle_num_speakers=-1,
|
|
max_num_speakers=8,
|
|
enhanced_count_thres=40,
|
|
sparse_search_volume=SSV,
|
|
max_rp_threshold=0.15,
|
|
fixed_thres=-1.0,
|
|
)
|
|
permuted_Y = stitch_cluster_labels(Y_old=gt, Y_new=Y_out)
|
|
permuted_Y = permuted_Y.to(gt.device)
|
|
# mc[-1] is the number of base scale segments
|
|
assert len(set(permuted_Y.tolist())) == n_spks
|
|
assert Y_out.shape[0] == mc[-1]
|
|
assert all(permuted_Y == gt)
|
|
|
|
@pytest.mark.run_only_on('CPU')
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize("n_spks", [1, 2, 3, 4, 5, 6, 7])
|
|
@pytest.mark.parametrize("total_sec, SSV, perturb_sigma, seed", [(30, 10, 0.1, 0)])
|
|
@pytest.mark.parametrize("jit_script", [False, True])
|
|
def test_offline_speaker_clustering_cpu(self, n_spks, total_sec, SSV, perturb_sigma, seed, jit_script, cuda=False):
|
|
self.test_offline_speaker_clustering(n_spks, total_sec, SSV, perturb_sigma, seed, jit_script, cuda=cuda)
|
|
|
|
@pytest.mark.run_only_on('CPU')
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize("n_spks", [1])
|
|
@pytest.mark.parametrize("spk_dur", [0.25, 0.5, 0.75, 1, 1.5, 2])
|
|
@pytest.mark.parametrize("SSV, enhanced_count_thres, min_samples_for_nmesc", [(5, 40, 6)])
|
|
@pytest.mark.parametrize("seed", [0])
|
|
def test_offline_speaker_clustering_very_short_cpu(
|
|
self,
|
|
n_spks,
|
|
spk_dur,
|
|
SSV,
|
|
enhanced_count_thres,
|
|
min_samples_for_nmesc,
|
|
seed,
|
|
):
|
|
em, ts, mc, mw, spk_ts, gt = generate_toy_data(
|
|
n_spks=n_spks, spk_dur=spk_dur, perturb_sigma=0.1, torch_seed=seed
|
|
)
|
|
offline_speaker_clustering = SpeakerClustering(maj_vote_spk_count=False, min_samples_for_nmesc=0, cuda=False)
|
|
assert isinstance(offline_speaker_clustering, SpeakerClustering)
|
|
Y_out = offline_speaker_clustering.forward_infer(
|
|
embeddings_in_scales=em,
|
|
timestamps_in_scales=ts,
|
|
multiscale_segment_counts=mc,
|
|
multiscale_weights=mw,
|
|
oracle_num_speakers=-1,
|
|
max_num_speakers=8,
|
|
enhanced_count_thres=enhanced_count_thres,
|
|
sparse_search_volume=SSV,
|
|
max_rp_threshold=0.15,
|
|
fixed_thres=-1.0,
|
|
)
|
|
permuted_Y = stitch_cluster_labels(Y_old=gt, Y_new=Y_out)
|
|
permuted_Y = permuted_Y.to(gt.device)
|
|
# mc[-1] is the number of base scale segments
|
|
assert len(set(permuted_Y.tolist())) == n_spks
|
|
assert Y_out.shape[0] == mc[-1]
|
|
assert all(permuted_Y == gt)
|
|
|
|
@pytest.mark.run_only_on('GPU')
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize("spk_dur", [0.25, 0.5, 0.75, 1, 2, 4])
|
|
@pytest.mark.parametrize("n_spks, SSV, enhanced_count_thres, min_samples_for_nmesc", [(1, 5, 40, 6)])
|
|
@pytest.mark.parametrize("seed", [0])
|
|
def test_offline_speaker_clustering_very_short_gpu(
|
|
self,
|
|
n_spks,
|
|
spk_dur,
|
|
SSV,
|
|
enhanced_count_thres,
|
|
min_samples_for_nmesc,
|
|
seed,
|
|
):
|
|
em, ts, mc, mw, spk_ts, gt = generate_toy_data(
|
|
n_spks=n_spks, spk_dur=spk_dur, perturb_sigma=0.1, torch_seed=seed
|
|
)
|
|
offline_speaker_clustering = SpeakerClustering(maj_vote_spk_count=False, min_samples_for_nmesc=0, cuda=True)
|
|
assert isinstance(offline_speaker_clustering, SpeakerClustering)
|
|
Y_out = offline_speaker_clustering.forward_infer(
|
|
embeddings_in_scales=em,
|
|
timestamps_in_scales=ts,
|
|
multiscale_segment_counts=mc,
|
|
multiscale_weights=mw,
|
|
oracle_num_speakers=-1,
|
|
max_num_speakers=8,
|
|
enhanced_count_thres=enhanced_count_thres,
|
|
sparse_search_volume=SSV,
|
|
max_rp_threshold=0.15,
|
|
fixed_thres=-1.0,
|
|
)
|
|
permuted_Y = stitch_cluster_labels(Y_old=gt, Y_new=Y_out)
|
|
permuted_Y = permuted_Y.to(gt.device)
|
|
# mc[-1] is the number of base scale segments
|
|
assert Y_out.shape[0] == mc[-1]
|
|
assert all(permuted_Y == gt)
|
|
|
|
@pytest.mark.run_only_on('CPU')
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize("n_spks, SSV, enhanced_count_thres, min_samples_for_nmesc", [(2, 5, 40, 6)])
|
|
@pytest.mark.parametrize("spk_dur, chunk_cluster_count, embeddings_per_chunk", [(120, 4, 50), (240, 4, 100)])
|
|
@pytest.mark.parametrize("seed", [0])
|
|
@pytest.mark.parametrize("jit_script", [False, True])
|
|
def test_longform_speaker_clustering_cpu(
|
|
self,
|
|
n_spks,
|
|
spk_dur,
|
|
SSV,
|
|
enhanced_count_thres,
|
|
min_samples_for_nmesc,
|
|
chunk_cluster_count,
|
|
embeddings_per_chunk,
|
|
jit_script,
|
|
seed,
|
|
):
|
|
em, ts, mc, mw, spk_ts, gt = generate_toy_data(
|
|
n_spks=n_spks, spk_dur=spk_dur, perturb_sigma=0.1, torch_seed=seed
|
|
)
|
|
longform_speaker_clustering = LongFormSpeakerClustering(cuda=False)
|
|
if jit_script:
|
|
longform_speaker_clustering = torch.jit.script(longform_speaker_clustering)
|
|
else:
|
|
assert isinstance(longform_speaker_clustering, LongFormSpeakerClustering)
|
|
Y_out = longform_speaker_clustering.forward_infer(
|
|
embeddings_in_scales=em,
|
|
timestamps_in_scales=ts,
|
|
multiscale_segment_counts=mc,
|
|
multiscale_weights=mw,
|
|
oracle_num_speakers=-1,
|
|
max_num_speakers=n_spks,
|
|
enhanced_count_thres=enhanced_count_thres,
|
|
sparse_search_volume=SSV,
|
|
max_rp_threshold=0.15,
|
|
fixed_thres=-1.0,
|
|
chunk_cluster_count=chunk_cluster_count,
|
|
embeddings_per_chunk=embeddings_per_chunk,
|
|
)
|
|
permuted_Y = stitch_cluster_labels(Y_old=gt, Y_new=Y_out)
|
|
permuted_Y = permuted_Y.to(gt.device)
|
|
|
|
# mc[-1] is the number of base scale segments
|
|
assert Y_out.shape[0] == mc[-1]
|
|
assert all(permuted_Y == gt)
|
|
|
|
@pytest.mark.run_only_on('GPU')
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize("n_spks, SSV, enhanced_count_thres, min_samples_for_nmesc", [(2, 5, 40, 6)])
|
|
@pytest.mark.parametrize("spk_dur, chunk_cluster_count, embeddings_per_chunk", [(120, 4, 50), (240, 4, 100)])
|
|
@pytest.mark.parametrize("seed", [0])
|
|
@pytest.mark.parametrize("jit_script", [False, True])
|
|
def test_longform_speaker_clustering_gpu(
|
|
self,
|
|
n_spks,
|
|
spk_dur,
|
|
SSV,
|
|
enhanced_count_thres,
|
|
min_samples_for_nmesc,
|
|
chunk_cluster_count,
|
|
embeddings_per_chunk,
|
|
jit_script,
|
|
seed,
|
|
):
|
|
em, ts, mc, mw, spk_ts, gt = generate_toy_data(
|
|
n_spks=n_spks, spk_dur=spk_dur, perturb_sigma=0.1, torch_seed=seed
|
|
)
|
|
longform_speaker_clustering = LongFormSpeakerClustering(cuda=True)
|
|
|
|
if jit_script:
|
|
longform_speaker_clustering = torch.jit.script(longform_speaker_clustering)
|
|
else:
|
|
assert isinstance(longform_speaker_clustering, LongFormSpeakerClustering)
|
|
|
|
Y_out = longform_speaker_clustering.forward_infer(
|
|
embeddings_in_scales=em,
|
|
timestamps_in_scales=ts,
|
|
multiscale_segment_counts=mc,
|
|
multiscale_weights=mw,
|
|
oracle_num_speakers=-1,
|
|
max_num_speakers=n_spks,
|
|
enhanced_count_thres=enhanced_count_thres,
|
|
sparse_search_volume=SSV,
|
|
max_rp_threshold=0.15,
|
|
fixed_thres=-1.0,
|
|
chunk_cluster_count=chunk_cluster_count,
|
|
embeddings_per_chunk=embeddings_per_chunk,
|
|
)
|
|
permuted_Y = stitch_cluster_labels(Y_old=gt, Y_new=Y_out)
|
|
permuted_Y = permuted_Y.to(gt.device)
|
|
|
|
# mc[-1] is the number of base scale segments
|
|
assert Y_out.shape[0] == mc[-1]
|
|
assert all(permuted_Y == gt)
|
|
|
|
@pytest.mark.run_only_on('GPU')
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize("n_spks", [1, 2, 3])
|
|
@pytest.mark.parametrize("total_sec, buffer_size, sigma", [(30, 30, 0.1)])
|
|
@pytest.mark.parametrize("seed", [0])
|
|
@pytest.mark.parametrize("jit_script", [False, True])
|
|
def test_online_speaker_clustering(self, n_spks, total_sec, buffer_size, sigma, seed, jit_script, cuda=True):
|
|
step_per_frame = 2
|
|
spk_dur = total_sec / n_spks
|
|
em, ts, mc, _, _, gt = generate_toy_data(n_spks, spk_dur=spk_dur, perturb_sigma=sigma, torch_seed=seed)
|
|
em_s, ts_s = split_input_data(em, ts, mc)
|
|
|
|
emb_gen = em_s[-1]
|
|
segment_indexes = ts_s[-1]
|
|
if cuda:
|
|
device = torch.cuda.current_device()
|
|
emb_gen, segment_indexes = emb_gen.to(device), segment_indexes.to(device)
|
|
|
|
history_buffer_size = buffer_size
|
|
current_buffer_size = buffer_size
|
|
|
|
online_clus = OnlineSpeakerClustering(
|
|
max_num_speakers=8,
|
|
max_rp_threshold=0.15,
|
|
sparse_search_volume=30,
|
|
history_buffer_size=history_buffer_size,
|
|
current_buffer_size=current_buffer_size,
|
|
cuda=cuda,
|
|
)
|
|
if jit_script:
|
|
online_clus = torch.jit.script(online_clus)
|
|
|
|
n_frames = int(emb_gen.shape[0] / step_per_frame)
|
|
evaluation_list = []
|
|
|
|
# Simulate online speaker clustering
|
|
for frame_index in range(n_frames):
|
|
curr_emb = emb_gen[0 : (frame_index + 1) * step_per_frame]
|
|
base_segment_indexes = torch.arange(curr_emb.shape[0]).to(curr_emb.device)
|
|
# Check history_buffer_size and history labels
|
|
assert (
|
|
online_clus.history_embedding_buffer_emb.shape[0] <= history_buffer_size
|
|
), "History buffer size error"
|
|
assert (
|
|
online_clus.history_embedding_buffer_emb.shape[0]
|
|
== online_clus.history_embedding_buffer_label.shape[0]
|
|
)
|
|
|
|
# Call clustering function
|
|
merged_clus_labels = online_clus.forward_infer(
|
|
curr_emb=curr_emb, base_segment_indexes=base_segment_indexes, frame_index=frame_index, cuda=cuda
|
|
)
|
|
|
|
# Resolve permutations
|
|
assert len(merged_clus_labels) == (frame_index + 1) * step_per_frame
|
|
# Resolve permutation issue by using stitch_cluster_labels function
|
|
merged_clus_labels = merged_clus_labels.cpu()
|
|
merged_clus_labels = stitch_cluster_labels(Y_old=gt[: len(merged_clus_labels)], Y_new=merged_clus_labels)
|
|
evaluation_list.extend(list(merged_clus_labels == gt[: len(merged_clus_labels)]))
|
|
|
|
assert online_clus.is_online
|
|
cumul_label_acc = sum(evaluation_list) / len(evaluation_list)
|
|
assert cumul_label_acc > 0.9
|
|
|
|
@pytest.mark.run_only_on('CPU')
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize("n_spks, total_sec, buffer_size, sigma, seed", [(3, 30, 30, 0.1, 0)])
|
|
@pytest.mark.parametrize("jit_script", [False, True])
|
|
def test_online_speaker_clustering_cpu(self, n_spks, total_sec, buffer_size, sigma, seed, jit_script, cuda=False):
|
|
self.test_online_speaker_clustering(n_spks, total_sec, buffer_size, sigma, seed, jit_script, cuda)
|
|
|
|
|
|
class TestLinearSumAssignmentAlgorithm:
|
|
@pytest.mark.unit
|
|
def test_lsa_solver_export_test(self):
|
|
cost_matrix = torch.randint(0, 10, (3, 3))
|
|
solver = LinearSumAssignmentSolver(cost_matrix)
|
|
solver = torch.jit.script(solver)
|
|
assert isinstance(solver, torch.jit._script.RecursiveScriptClass)
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize(
|
|
"cost_matrix",
|
|
[torch.tensor([[7, 6, 2, 9, 2], [6, 2, 1, 3, 9], [5, 6, 8, 9, 5], [6, 8, 5, 8, 6], [9, 5, 6, 4, 7]])],
|
|
)
|
|
def test_linear_sum_assignment_algorithm_cost_matrix(self, cost_matrix):
|
|
"""
|
|
Test the linear sum assignment algorithm with a cost matrix
|
|
|
|
Compare with the scipy implementation and make sure the final cost is the same.
|
|
NOTE: There could be multiple solutions with the same cost in linear sum assignment problem.
|
|
This test only checks if the cost is the same.
|
|
"""
|
|
row_ind_nm, col_ind_nm = nemo_linear_sum_assignment(cost_matrix)
|
|
row_ind_sc, col_ind_sc = scipy_linear_sum_assignment(cost_matrix.cpu().numpy())
|
|
cost_nm = sum(cost_matrix[row_ind_nm, col_ind_nm])
|
|
cost_sc = sum(cost_matrix[row_ind_sc, col_ind_sc])
|
|
assert cost_nm == cost_sc
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize("seed", [0, 1])
|
|
@pytest.mark.parametrize("mat_size", [1, 2, 4, 8])
|
|
def test_linear_sum_assignment_algorithm_random_matrix(self, seed, mat_size):
|
|
torch.manual_seed(seed)
|
|
cost_matrix = torch.randint(0, 10, (mat_size, mat_size))
|
|
self.test_linear_sum_assignment_algorithm_cost_matrix(cost_matrix)
|