diff --git a/funasr/models/campplus/cluster_backend.py b/funasr/models/campplus/cluster_backend.py index 998ccd679..a369d626d 100644 --- a/funasr/models/campplus/cluster_backend.py +++ b/funasr/models/campplus/cluster_backend.py @@ -11,6 +11,7 @@ from sklearn.cluster._kmeans import k_means from sklearn.cluster import HDBSCAN +from sklearn.preprocessing import normalize class SpectralCluster: @@ -191,6 +192,21 @@ def __call__(self, X): return labels +class KMeansCluster: + r"""Cluster a large set of speaker embeddings into a known number of groups.""" + + def __call__(self, X, num_clusters): + """Cluster L2-normalized embeddings with bounded memory usage.""" + normalized_X = normalize(X) + _, labels, _ = k_means( + normalized_X, + num_clusters, + random_state=0, + n_init=10, + ) + return labels + + class ClusterBackend(torch.nn.Module): r"""Perfom clustering for input embeddings and output the labels. Args: @@ -210,6 +226,7 @@ def __init__(self, merge_thr=0.78): self.spectral_cluster = SpectralCluster() self.umap_hdbscan_cluster = UmapHdbscan() + self.kmeans_cluster = KMeansCluster() def forward(self, X, **params): # clustering and return the labels @@ -223,9 +240,10 @@ def forward(self, X, **params): assert len(X.shape) == 2, "modelscope error: the shape of input should be [N, C]" if X.shape[0] < 20: return np.zeros(X.shape[0], dtype="int") - if X.shape[0] < 2048 or k is not None: - # unexpected corner case + if X.shape[0] < 2048: labels = self.spectral_cluster(X, k) + elif k is not None: + labels = self.kmeans_cluster(X, k) else: labels = self.umap_hdbscan_cluster(X) diff --git a/tests/test_cluster_backend.py b/tests/test_cluster_backend.py new file mode 100644 index 000000000..c2a6daeec --- /dev/null +++ b/tests/test_cluster_backend.py @@ -0,0 +1,44 @@ +import numpy as np +import pytest +import torch + +from funasr.models.campplus.cluster_backend import ClusterBackend + + +def test_large_known_speaker_count_uses_fixed_k_clustering(monkeypatch): + backend = ClusterBackend() + embeddings = torch.ones((2048, 2)) + expected = np.arange(embeddings.shape[0]) % 2 + + def fail_spectral(*args, **kwargs): + pytest.fail("large inputs must not use dense spectral clustering") + + def fail_umap(*args, **kwargs): + pytest.fail("a known speaker count should use fixed-K clustering") + + def fixed_k_cluster(actual_embeddings, num_clusters): + assert actual_embeddings is embeddings + assert num_clusters == 2 + return expected + + monkeypatch.setattr(backend, "spectral_cluster", fail_spectral) + monkeypatch.setattr(backend, "umap_hdbscan_cluster", fail_umap) + monkeypatch.setattr(backend, "kmeans_cluster", fixed_k_cluster, raising=False) + + labels = backend(embeddings, oracle_num=2) + + np.testing.assert_array_equal(labels, expected) + + +def test_large_fixed_k_clustering_separates_cosine_clusters(): + backend = ClusterBackend() + embeddings = torch.zeros((2048, 2)) + embeddings[:1024, 0] = 1 + embeddings[1024:, 1] = 1 + + labels = backend(embeddings, oracle_num=2) + + assert np.unique(labels).size == 2 + assert np.unique(labels[:1024]).size == 1 + assert np.unique(labels[1024:]).size == 1 + assert labels[0] != labels[-1]