Class ClusterBackend

funasr/models/campplus/cluster_backend.py:194–268

r"""Perfom clustering for input embeddings and output the labels. Args: model_dir: A model dir. model_config: The model config.

class ClusterBackend(torch.nn.Module):
    r"""Perfom clustering for input embeddings and output the labels.
    Args:
        model_dir: A model dir.
        model_config: The model config.
    """

    def __init__(self, merge_thr=0.78):
        """Initialize ClusterBackend.
        
            Args:
                merge_thr: TODO.
            """
        super().__init__()
        self.model_config = {"merge_thr": merge_thr}
        # self.other_config = kwargs

        self.spectral_cluster = SpectralCluster()
        self.umap_hdbscan_cluster = UmapHdbscan()

    def forward(self, X, **params):
        # clustering and return the labels
        """Forward pass for training.
        
            Args:
                X: TODO.
                **params: Additional keyword arguments.
            """
        k = params["oracle_num"] if "oracle_num" in params else None
        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
            labels = self.spectral_cluster(X, k)
        else:
            labels = self.umap_hdbscan_cluster(X)

        if k is None and "merge_thr" in self.model_config:
            labels = self.merge_by_cos(labels, X, self.model_config["merge_thr"])

        return labels

    def merge_by_cos(self, labels, embs, cos_thr):
        # merge the similar speakers by cosine similarity
        """Merge by cos.
        
            Args:
                labels: TODO.
                embs: TODO.
                cos_thr: TODO.
            """
        assert cos_thr > 0 and cos_thr <= 1
        while True:
            spk_num = labels.max() + 1
            if spk_num == 1:
                break
            spk_center = []