Initial commit: FunASR Speech Recognition Toolkit

Add complete FunASR codebase including models, runtime, and documentation.

funasr/models/campplus/cluster_backend.py

@@ -0,0 +1,268 @@
+#!/usr/bin/env python3
+# -*- encoding: utf-8 -*-
+# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved.
+#  MIT License  (https://opensource.org/licenses/MIT)
+# Modified from 3D-Speaker (https://github.com/alibaba-damo-academy/3D-Speaker)
+
+import scipy
+import torch
+import sklearn
+import numpy as np
+
+from sklearn.cluster._kmeans import k_means
+from sklearn.cluster import HDBSCAN
+
+
+class SpectralCluster:
+    r"""A spectral clustering mehtod using unnormalized Laplacian of affinity matrix.
+    This implementation is adapted from https://github.com/speechbrain/speechbrain.
+    """
+
+    def __init__(self, min_num_spks=1, max_num_spks=15, pval=0.022):
+        """Initialize SpectralCluster.
+        
+            Args:
+                min_num_spks: TODO.
+                max_num_spks: TODO.
+                pval: TODO.
+            """
+        self.min_num_spks = min_num_spks
+        self.max_num_spks = max_num_spks
+        self.pval = pval
+
+    def __call__(self, X, oracle_num=None):
+        # Similarity matrix computation
+        """Internal: call  .
+        
+            Args:
+                X: TODO.
+                oracle_num: TODO.
+            """
+        sim_mat = self.get_sim_mat(X)
+
+        # Refining similarity matrix with pval
+        prunned_sim_mat = self.p_pruning(sim_mat)
+
+        # Symmetrization
+        sym_prund_sim_mat = 0.5 * (prunned_sim_mat + prunned_sim_mat.T)
+
+        # Laplacian calculation
+        laplacian = self.get_laplacian(sym_prund_sim_mat)
+
+        # Get Spectral Embeddings
+        emb, num_of_spk = self.get_spec_embs(laplacian, oracle_num)
+
+        # Perform clustering
+        labels = self.cluster_embs(emb, num_of_spk)
+
+        return labels
+
+    def get_sim_mat(self, X):
+        # Cosine similarities
+        """Get sim mat.
+        
+            Args:
+                X: TODO.
+            """
+        M = sklearn.metrics.pairwise.cosine_similarity(X, X)
+        return M
+
+    def p_pruning(self, A):
+        """P pruning.
+        
+            Args:
+                A: TODO.
+            """
+        if A.shape[0] * self.pval < 6:
+            pval = 6.0 / A.shape[0]
+        else:
+            pval = self.pval
+
+        n_elems = int((1 - pval) * A.shape[0])
+
+        # For each row in a affinity matrix
+        for i in range(A.shape[0]):
+            low_indexes = np.argsort(A[i, :])
+            low_indexes = low_indexes[0:n_elems]
+
+            # Replace smaller similarity values by 0s
+            A[i, low_indexes] = 0
+        return A
+
+    def get_laplacian(self, M):
+        """Get laplacian.
+        
+            Args:
+                M: TODO.
+            """
+        M[np.diag_indices(M.shape[0])] = 0
+        D = np.sum(np.abs(M), axis=1)
+        D = np.diag(D)
+        L = D - M
+        return L
+
+    def get_spec_embs(self, L, k_oracle=None):
+        """Get spec embs.
+        
+            Args:
+                L: TODO.
+                k_oracle: TODO.
+            """
+        lambdas, eig_vecs = scipy.linalg.eigh(L)
+
+        if k_oracle is not None:
+            num_of_spk = k_oracle
+        else:
+            lambda_gap_list = self.getEigenGaps(
+                lambdas[self.min_num_spks - 1 : self.max_num_spks + 1]
+            )
+            num_of_spk = np.argmax(lambda_gap_list) + self.min_num_spks
+
+        emb = eig_vecs[:, :num_of_spk]
+        return emb, num_of_spk
+
+    def cluster_embs(self, emb, k):
+        """Cluster embs.
+        
+            Args:
+                emb: TODO.
+                k: TODO.
+            """
+        _, labels, _ = k_means(emb, k)
+        return labels
+
+    def getEigenGaps(self, eig_vals):
+        """Geteigengaps.
+        
+            Args:
+                eig_vals: TODO.
+            """
+        eig_vals_gap_list = []
+        for i in range(len(eig_vals) - 1):
+            gap = float(eig_vals[i + 1]) - float(eig_vals[i])
+            eig_vals_gap_list.append(gap)
+        return eig_vals_gap_list
+
+
+class UmapHdbscan:
+    r"""
+    Reference:
+    - Siqi Zheng, Hongbin Suo. Reformulating Speaker Diarization as Community Detection With
+      Emphasis On Topological Structure. ICASSP2022
+    """
+
+    def __init__(
+        self, n_neighbors=20, n_components=60, min_samples=10, min_cluster_size=10, metric="cosine"
+    ):
+        """Initialize UmapHdbscan.
+        
+            Args:
+                n_neighbors: TODO.
+                n_components: TODO.
+                min_samples: TODO.
+                min_cluster_size: Size/dimension parameter.
+                metric: TODO.
+            """
+        self.n_neighbors = n_neighbors
+        self.n_components = n_components
+        self.min_samples = min_samples
+        self.min_cluster_size = min_cluster_size
+        self.metric = metric
+
+    def __call__(self, X):
+        """Internal: call  .
+        
+            Args:
+                X: TODO.
+            """
+        import umap.umap_ as umap
+
+        umap_X = umap.UMAP(
+            n_neighbors=self.n_neighbors,
+            min_dist=0.0,
+            n_components=min(self.n_components, X.shape[0] - 2),
+            metric=self.metric,
+        ).fit_transform(X)
+        labels = HDBSCAN(
+            min_samples=self.min_samples,
+            min_cluster_size=self.min_cluster_size,
+            allow_single_cluster=True,
+        ).fit_predict(umap_X)
+        return labels
+
+
+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 = []
+            for i in range(spk_num):
+                spk_emb = embs[labels == i].mean(0)
+                spk_center.append(spk_emb)
+            assert len(spk_center) > 0
+            spk_center = np.stack(spk_center, axis=0)
+            norm_spk_center = spk_center / np.linalg.norm(spk_center, axis=1, keepdims=True)
+            affinity = np.matmul(norm_spk_center, norm_spk_center.T)
+            affinity = np.triu(affinity, 1)
+            spks = np.unravel_index(np.argmax(affinity), affinity.shape)
+            if affinity[spks] < cos_thr:
+                break
+            for i in range(len(labels)):
+                if labels[i] == spks[1]:
+                    labels[i] = spks[0]
+                elif labels[i] > spks[1]:
+                    labels[i] -= 1
+        return labels

funasr/models/campplus/components.py

@@ -0,0 +1,450 @@
+#!/usr/bin/env python3
+# -*- encoding: utf-8 -*-
+# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved.
+#  MIT License  (https://opensource.org/licenses/MIT)
+# Modified from 3D-Speaker (https://github.com/alibaba-damo-academy/3D-Speaker)
+
+import torch
+import torch.nn.functional as F
+import torch.utils.checkpoint as cp
+
+
+class BasicResBlock(torch.nn.Module):
+    expansion = 1
+
+    def __init__(self, in_planes, planes, stride=1):
+        """Initialize BasicResBlock.
+        
+            Args:
+                in_planes: TODO.
+                planes: TODO.
+                stride: TODO.
+            """
+        super(BasicResBlock, self).__init__()
+        self.conv1 = torch.nn.Conv2d(
+            in_planes, planes, kernel_size=3, stride=(stride, 1), padding=1, bias=False
+        )
+        self.bn1 = torch.nn.BatchNorm2d(planes)
+        self.conv2 = torch.nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn2 = torch.nn.BatchNorm2d(planes)
+
+        self.shortcut = torch.nn.Sequential()
+        if stride != 1 or in_planes != self.expansion * planes:
+            self.shortcut = torch.nn.Sequential(
+                torch.nn.Conv2d(
+                    in_planes,
+                    self.expansion * planes,
+                    kernel_size=1,
+                    stride=(stride, 1),
+                    bias=False,
+                ),
+                torch.nn.BatchNorm2d(self.expansion * planes),
+            )
+
+    def forward(self, x):
+        """Forward pass for training.
+        
+            Args:
+                x: TODO.
+            """
+        out = F.relu(self.bn1(self.conv1(x)))
+        out = self.bn2(self.conv2(out))
+        out += self.shortcut(x)
+        out = F.relu(out)
+        return out
+
+
+class FCM(torch.nn.Module):
+    def __init__(self, block=BasicResBlock, num_blocks=[2, 2], m_channels=32, feat_dim=80):
+        """Initialize FCM.
+        
+            Args:
+                block: TODO.
+                num_blocks: TODO.
+                m_channels: TODO.
+                feat_dim: Size/dimension parameter.
+            """
+        super(FCM, self).__init__()
+        self.in_planes = m_channels
+        self.conv1 = torch.nn.Conv2d(1, m_channels, kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn1 = torch.nn.BatchNorm2d(m_channels)
+
+        self.layer1 = self._make_layer(block, m_channels, num_blocks[0], stride=2)
+        self.layer2 = self._make_layer(block, m_channels, num_blocks[0], stride=2)
+
+        self.conv2 = torch.nn.Conv2d(
+            m_channels, m_channels, kernel_size=3, stride=(2, 1), padding=1, bias=False
+        )
+        self.bn2 = torch.nn.BatchNorm2d(m_channels)
+        self.out_channels = m_channels * (feat_dim // 8)
+
+    def _make_layer(self, block, planes, num_blocks, stride):
+        """Internal: make layer.
+        
+            Args:
+                block: TODO.
+                planes: TODO.
+                num_blocks: TODO.
+                stride: TODO.
+            """
+        strides = [stride] + [1] * (num_blocks - 1)
+        layers = []
+        for stride in strides:
+            layers.append(block(self.in_planes, planes, stride))
+            self.in_planes = planes * block.expansion
+        return torch.nn.Sequential(*layers)
+
+    def forward(self, x):
+        """Forward pass for training.
+        
+            Args:
+                x: TODO.
+            """
+        x = x.unsqueeze(1)
+        out = F.relu(self.bn1(self.conv1(x)))
+        out = self.layer1(out)
+        out = self.layer2(out)
+        out = F.relu(self.bn2(self.conv2(out)))
+
+        shape = out.shape
+        out = out.reshape(shape[0], shape[1] * shape[2], shape[3])
+        return out
+
+
+def get_nonlinear(config_str, channels):
+    """Get nonlinear.
+    
+        Args:
+            config_str: TODO.
+            channels: TODO.
+        """
+    nonlinear = torch.nn.Sequential()
+    for name in config_str.split("-"):
+        if name == "relu":
+            nonlinear.add_module("relu", torch.nn.ReLU(inplace=True))
+        elif name == "prelu":
+            nonlinear.add_module("prelu", torch.nn.PReLU(channels))
+        elif name == "batchnorm":
+            nonlinear.add_module("batchnorm", torch.nn.BatchNorm1d(channels))
+        elif name == "batchnorm_":
+            nonlinear.add_module("batchnorm", torch.nn.BatchNorm1d(channels, affine=False))
+        else:
+            raise ValueError("Unexpected module ({}).".format(name))
+    return nonlinear
+
+
+def statistics_pooling(x, dim=-1, keepdim=False, unbiased=True, eps=1e-2):
+    """Statistics pooling.
+    
+        Args:
+            x: TODO.
+            dim: TODO.
+            keepdim: TODO.
+            unbiased: TODO.
+            eps: TODO.
+        """
+    mean = x.mean(dim=dim)
+    std = x.std(dim=dim, unbiased=unbiased)
+    stats = torch.cat([mean, std], dim=-1)
+    if keepdim:
+        stats = stats.unsqueeze(dim=dim)
+    return stats
+
+
+class StatsPool(torch.nn.Module):
+    def forward(self, x):
+        """Forward pass for training.
+        
+            Args:
+                x: TODO.
+            """
+        return statistics_pooling(x)
+
+
+class TDNNLayer(torch.nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride=1,
+        padding=0,
+        dilation=1,
+        bias=False,
+        config_str="batchnorm-relu",
+    ):
+        """Initialize TDNNLayer.
+        
+            Args:
+                in_channels: TODO.
+                out_channels: TODO.
+                kernel_size: Size/dimension parameter.
+                stride: TODO.
+                padding: TODO.
+                dilation: TODO.
+                bias: TODO.
+                config_str: TODO.
+            """
+        super(TDNNLayer, self).__init__()
+        if padding < 0:
+            assert (
+                kernel_size % 2 == 1
+            ), "Expect equal paddings, but got even kernel size ({})".format(kernel_size)
+            padding = (kernel_size - 1) // 2 * dilation
+        self.linear = torch.nn.Conv1d(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            bias=bias,
+        )
+        self.nonlinear = get_nonlinear(config_str, out_channels)
+
+    def forward(self, x):
+        """Forward pass for training.
+        
+            Args:
+                x: TODO.
+            """
+        x = self.linear(x)
+        x = self.nonlinear(x)
+        return x
+
+
+class CAMLayer(torch.nn.Module):
+    def __init__(
+        self, bn_channels, out_channels, kernel_size, stride, padding, dilation, bias, reduction=2
+    ):
+        """Initialize CAMLayer.
+        
+            Args:
+                bn_channels: TODO.
+                out_channels: TODO.
+                kernel_size: Size/dimension parameter.
+                stride: TODO.
+                padding: TODO.
+                dilation: TODO.
+                bias: TODO.
+                reduction: TODO.
+            """
+        super(CAMLayer, self).__init__()
+        self.linear_local = torch.nn.Conv1d(
+            bn_channels,
+            out_channels,
+            kernel_size,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            bias=bias,
+        )
+        self.linear1 = torch.nn.Conv1d(bn_channels, bn_channels // reduction, 1)
+        self.relu = torch.nn.ReLU(inplace=True)
+        self.linear2 = torch.nn.Conv1d(bn_channels // reduction, out_channels, 1)
+        self.sigmoid = torch.nn.Sigmoid()
+
+    def forward(self, x):
+        """Forward pass for training.
+        
+            Args:
+                x: TODO.
+            """
+        y = self.linear_local(x)
+        context = x.mean(-1, keepdim=True) + self.seg_pooling(x)
+        context = self.relu(self.linear1(context))
+        m = self.sigmoid(self.linear2(context))
+        return y * m
+
+    def seg_pooling(self, x, seg_len=100, stype="avg"):
+        """Seg pooling.
+        
+            Args:
+                x: TODO.
+                seg_len: TODO.
+                stype: TODO.
+            """
+        if stype == "avg":
+            seg = F.avg_pool1d(x, kernel_size=seg_len, stride=seg_len, ceil_mode=True)
+        elif stype == "max":
+            seg = F.max_pool1d(x, kernel_size=seg_len, stride=seg_len, ceil_mode=True)
+        else:
+            raise ValueError("Wrong segment pooling type.")
+        shape = seg.shape
+        seg = seg.unsqueeze(-1).expand(*shape, seg_len).reshape(*shape[:-1], -1)
+        seg = seg[..., : x.shape[-1]]
+        return seg
+
+
+class CAMDenseTDNNLayer(torch.nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        bn_channels,
+        kernel_size,
+        stride=1,
+        dilation=1,
+        bias=False,
+        config_str="batchnorm-relu",
+        memory_efficient=False,
+    ):
+        """Initialize CAMDenseTDNNLayer.
+        
+            Args:
+                in_channels: TODO.
+                out_channels: TODO.
+                bn_channels: TODO.
+                kernel_size: Size/dimension parameter.
+                stride: TODO.
+                dilation: TODO.
+                bias: TODO.
+                config_str: TODO.
+                memory_efficient: TODO.
+            """
+        super(CAMDenseTDNNLayer, self).__init__()
+        assert kernel_size % 2 == 1, "Expect equal paddings, but got even kernel size ({})".format(
+            kernel_size
+        )
+        padding = (kernel_size - 1) // 2 * dilation
+        self.memory_efficient = memory_efficient
+        self.nonlinear1 = get_nonlinear(config_str, in_channels)
+        self.linear1 = torch.nn.Conv1d(in_channels, bn_channels, 1, bias=False)
+        self.nonlinear2 = get_nonlinear(config_str, bn_channels)
+        self.cam_layer = CAMLayer(
+            bn_channels,
+            out_channels,
+            kernel_size,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            bias=bias,
+        )
+
+    def bn_function(self, x):
+        """Bn function.
+        
+            Args:
+                x: TODO.
+            """
+        return self.linear1(self.nonlinear1(x))
+
+    def forward(self, x):
+        """Forward pass for training.
+        
+            Args:
+                x: TODO.
+            """
+        if self.training and self.memory_efficient:
+            x = cp.checkpoint(self.bn_function, x)
+        else:
+            x = self.bn_function(x)
+        x = self.cam_layer(self.nonlinear2(x))
+        return x
+
+
+class CAMDenseTDNNBlock(torch.nn.ModuleList):
+    def __init__(
+        self,
+        num_layers,
+        in_channels,
+        out_channels,
+        bn_channels,
+        kernel_size,
+        stride=1,
+        dilation=1,
+        bias=False,
+        config_str="batchnorm-relu",
+        memory_efficient=False,
+    ):
+        """Initialize CAMDenseTDNNBlock.
+        
+            Args:
+                num_layers: TODO.
+                in_channels: TODO.
+                out_channels: TODO.
+                bn_channels: TODO.
+                kernel_size: Size/dimension parameter.
+                stride: TODO.
+                dilation: TODO.
+                bias: TODO.
+                config_str: TODO.
+                memory_efficient: TODO.
+            """
+        super(CAMDenseTDNNBlock, self).__init__()
+        for i in range(num_layers):
+            layer = CAMDenseTDNNLayer(
+                in_channels=in_channels + i * out_channels,
+                out_channels=out_channels,
+                bn_channels=bn_channels,
+                kernel_size=kernel_size,
+                stride=stride,
+                dilation=dilation,
+                bias=bias,
+                config_str=config_str,
+                memory_efficient=memory_efficient,
+            )
+            self.add_module("tdnnd%d" % (i + 1), layer)
+
+    def forward(self, x):
+        """Forward pass for training.
+        
+            Args:
+                x: TODO.
+            """
+        for layer in self:
+            x = torch.cat([x, layer(x)], dim=1)
+        return x
+
+
+class TransitLayer(torch.nn.Module):
+    def __init__(self, in_channels, out_channels, bias=True, config_str="batchnorm-relu"):
+        """Initialize TransitLayer.
+        
+            Args:
+                in_channels: TODO.
+                out_channels: TODO.
+                bias: TODO.
+                config_str: TODO.
+            """
+        super(TransitLayer, self).__init__()
+        self.nonlinear = get_nonlinear(config_str, in_channels)
+        self.linear = torch.nn.Conv1d(in_channels, out_channels, 1, bias=bias)
+
+    def forward(self, x):
+        """Forward pass for training.
+        
+            Args:
+                x: TODO.
+            """
+        x = self.nonlinear(x)
+        x = self.linear(x)
+        return x
+
+
+class DenseLayer(torch.nn.Module):
+    def __init__(self, in_channels, out_channels, bias=False, config_str="batchnorm-relu"):
+        """Initialize DenseLayer.
+        
+            Args:
+                in_channels: TODO.
+                out_channels: TODO.
+                bias: TODO.
+                config_str: TODO.
+            """
+        super(DenseLayer, self).__init__()
+        self.linear = torch.nn.Conv1d(in_channels, out_channels, 1, bias=bias)
+        self.nonlinear = get_nonlinear(config_str, out_channels)
+
+    def forward(self, x):
+        """Forward pass for training.
+        
+            Args:
+                x: TODO.
+            """
+        if len(x.shape) == 2:
+            x = self.linear(x.unsqueeze(dim=-1)).squeeze(dim=-1)
+        else:
+            x = self.linear(x)
+        x = self.nonlinear(x)
+        return x

funasr/models/campplus/model.py

@@ -0,0 +1,195 @@
+#!/usr/bin/env python3
+# -*- encoding: utf-8 -*-
+# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved.
+#  MIT License  (https://opensource.org/licenses/MIT)
+# Modified from 3D-Speaker (https://github.com/alibaba-damo-academy/3D-Speaker)
+
+import time
+import torch
+import numpy as np
+from collections import OrderedDict
+from contextlib import contextmanager
+from distutils.version import LooseVersion
+
+from funasr.register import tables
+from funasr.models.campplus.utils import extract_feature
+from funasr.utils.load_utils import load_audio_text_image_video
+from funasr.models.campplus.components import (
+    DenseLayer,
+    StatsPool,
+    TDNNLayer,
+    CAMDenseTDNNBlock,
+    TransitLayer,
+    get_nonlinear,
+    FCM,
+)
+
+
+if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"):
+    from torch.cuda.amp import autocast
+else:
+    # Nothing to do if torch<1.6.0
+    @contextmanager
+    def autocast(enabled=True):
+        """Autocast.
+        
+            Args:
+                enabled: TODO.
+            """
+        yield
+
+
+@tables.register("model_classes", "CAMPPlus")
+class CAMPPlus(torch.nn.Module):
+    """CAM++ Speaker Verification Model.
+
+    Extracts fixed-dimensional speaker embeddings from variable-length audio.
+    Used for speaker verification and speaker diarization pipelines.
+
+    Output: 192-dimensional speaker embedding per utterance.
+    """
+
+    def __init__(
+        self,
+        feat_dim=80,
+        embedding_size=192,
+        growth_rate=32,
+        bn_size=4,
+        init_channels=128,
+        config_str="batchnorm-relu",
+        memory_efficient=True,
+        output_level="segment",
+        **kwargs,
+    ):
+        """Initialize CAMPPlus.
+        
+            Args:
+                feat_dim: Size/dimension parameter.
+                embedding_size: Size/dimension parameter.
+                growth_rate: TODO.
+                bn_size: Size/dimension parameter.
+                init_channels: TODO.
+                config_str: TODO.
+                memory_efficient: TODO.
+                output_level: TODO.
+                **kwargs: Additional keyword arguments.
+            """
+        super().__init__()
+
+        self.head = FCM(feat_dim=feat_dim)
+        channels = self.head.out_channels
+        self.output_level = output_level
+
+        self.xvector = torch.nn.Sequential(
+            OrderedDict(
+                [
+                    (
+                        "tdnn",
+                        TDNNLayer(
+                            channels,
+                            init_channels,
+                            5,
+                            stride=2,
+                            dilation=1,
+                            padding=-1,
+                            config_str=config_str,
+                        ),
+                    ),
+                ]
+            )
+        )
+        channels = init_channels
+        for i, (num_layers, kernel_size, dilation) in enumerate(
+            zip((12, 24, 16), (3, 3, 3), (1, 2, 2))
+        ):
+            block = CAMDenseTDNNBlock(
+                num_layers=num_layers,
+                in_channels=channels,
+                out_channels=growth_rate,
+                bn_channels=bn_size * growth_rate,
+                kernel_size=kernel_size,
+                dilation=dilation,
+                config_str=config_str,
+                memory_efficient=memory_efficient,
+            )
+            self.xvector.add_module("block%d" % (i + 1), block)
+            channels = channels + num_layers * growth_rate
+            self.xvector.add_module(
+                "transit%d" % (i + 1),
+                TransitLayer(channels, channels // 2, bias=False, config_str=config_str),
+            )
+            channels //= 2
+
+        self.xvector.add_module("out_nonlinear", get_nonlinear(config_str, channels))
+
+        if self.output_level == "segment":
+            self.xvector.add_module("stats", StatsPool())
+            self.xvector.add_module(
+                "dense", DenseLayer(channels * 2, embedding_size, config_str="batchnorm_")
+            )
+        else:
+            assert (
+                self.output_level == "frame"
+            ), "`output_level` should be set to 'segment' or 'frame'. "
+
+        for m in self.modules():
+            if isinstance(m, (torch.nn.Conv1d, torch.nn.Linear)):
+                torch.nn.init.kaiming_normal_(m.weight.data)
+                if m.bias is not None:
+                    torch.nn.init.zeros_(m.bias)
+
+    def forward(self, x):
+        """Extract speaker embedding from fbank features.
+
+        Args:
+            x (Tensor): Input fbank features, shape (batch, time, feat_dim).
+
+        Returns:
+            Tensor: Speaker embedding, shape (batch, embedding_size) for segment level,
+                or (batch, time, channels) for frame level.
+        """
+        x = x.permute(0, 2, 1)  # (B,T,F) => (B,F,T)
+        x = self.head(x)
+        x = self.xvector(x)
+        if self.output_level == "frame":
+            x = x.transpose(1, 2)
+        return x
+
+    def inference(
+        self,
+        data_in,
+        data_lengths=None,
+        key: list = None,
+        tokenizer=None,
+        frontend=None,
+        **kwargs,
+    ):
+        """Run speaker embedding extraction on audio input.
+
+        Args:
+            data_in: Audio input (file path, numpy array, or list).
+            data_lengths: Not used.
+            key (list): Sample identifiers.
+            tokenizer: Not used.
+            frontend: Not used.
+            **kwargs: Must include 'device' (str) and optional 'fs' (int, default 16000).
+
+        Returns:
+            tuple: (results, meta_data) where results is
+                [{"spk_embedding": Tensor of shape (1, 192)}]
+        """
+        # extract fbank feats
+        meta_data = {}
+        time1 = time.perf_counter()
+        audio_sample_list = load_audio_text_image_video(
+            data_in, fs=16000, audio_fs=kwargs.get("fs", 16000), data_type="sound"
+        )
+        time2 = time.perf_counter()
+        meta_data["load_data"] = f"{time2 - time1:0.3f}"
+        speech, speech_lengths, speech_times = extract_feature(audio_sample_list)
+        speech = speech.to(device=kwargs["device"])
+        time3 = time.perf_counter()
+        meta_data["extract_feat"] = f"{time3 - time2:0.3f}"
+        meta_data["batch_data_time"] = np.array(speech_times).sum().item() / 16000.0
+        results = [{"spk_embedding": self.forward(speech.to(torch.float32))}]
+        return results, meta_data

funasr/models/campplus/template.yaml

@@ -0,0 +1,23 @@
+# This is an example that demonstrates how to configure a model file.
+# You can modify the configuration according to your own requirements.
+
+# to print the register_table:
+# from funasr.register import tables
+# tables.print()
+
+# network architecture
+model: CAMPPlus
+model_conf:
+    feat_dim: 80
+    embedding_size: 192
+    growth_rate: 32
+    bn_size: 4
+    init_channels: 128
+    config_str: 'batchnorm-relu'
+    memory_efficient: True
+    output_level: 'segment'
+
+# frontend related
+frontend: WavFrontend
+frontend_conf:
+    fs: 16000

funasr/models/campplus/utils.py

@@ -0,0 +1,649 @@
+#!/usr/bin/env python3
+# -*- encoding: utf-8 -*-
+# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved.
+#  MIT License  (https://opensource.org/licenses/MIT)
+# Modified from 3D-Speaker (https://github.com/alibaba-damo-academy/3D-Speaker)
+
+import io
+import os
+import torch
+import requests
+import tempfile
+import contextlib
+import numpy as np
+import librosa as sf
+from typing import Union
+from pathlib import Path
+from typing import Generator, Union
+from abc import ABCMeta, abstractmethod
+import torchaudio.compliance.kaldi as Kaldi
+
+from funasr.models.transformer.utils.nets_utils import pad_list
+
+
+def check_audio_list(audio: list):
+    """Check audio list.
+    
+        Args:
+            audio: TODO.
+        """
+    audio_dur = 0
+    for i in range(len(audio)):
+        seg = audio[i]
+        assert seg[1] >= seg[0], "modelscope error: Wrong time stamps."
+        assert isinstance(seg[2], np.ndarray), "modelscope error: Wrong data type."
+        assert (
+            int(seg[1] * 16000) - int(seg[0] * 16000) == seg[2].shape[0]
+        ), "modelscope error: audio data in list is inconsistent with time length."
+        if i > 0:
+            assert seg[0] >= audio[i - 1][1], "modelscope error: Wrong time stamps."
+        audio_dur += seg[1] - seg[0]
+    return audio_dur
+    # assert audio_dur > 5, 'modelscope error: The effective audio duration is too short.'
+
+
+def sv_preprocess(inputs: Union[np.ndarray, list]):
+    """Sv preprocess.
+    
+        Args:
+            inputs: TODO.
+        """
+    output = []
+    for i in range(len(inputs)):
+        if isinstance(inputs[i], str):
+            file_bytes = File.read(inputs[i])
+            data, fs = sf.load(io.BytesIO(file_bytes), dtype="float32")
+            if len(data.shape) == 2:
+                data = data[:, 0]
+            data = torch.from_numpy(data).unsqueeze(0)
+            data = data.squeeze(0)
+        elif isinstance(inputs[i], np.ndarray):
+            assert len(inputs[i].shape) == 1, "modelscope error: Input array should be [N, T]"
+            data = inputs[i]
+            if data.dtype in ["int16", "int32", "int64"]:
+                data = (data / (1 << 15)).astype("float32")
+            else:
+                data = data.astype("float32")
+            data = torch.from_numpy(data)
+        else:
+            raise ValueError(
+                "modelscope error: The input type is restricted to audio address and nump array."
+            )
+        output.append(data)
+    return output
+
+
+def sv_chunk(vad_segments: list, fs=16000) -> list:
+    """Sv chunk.
+    
+        Args:
+            vad_segments: TODO.
+            fs: TODO.
+        """
+    config = {
+        "seg_dur": 1.5,
+        "seg_shift": 0.75,
+    }
+
+    def seg_chunk(seg_data):
+        """Seg chunk.
+        
+            Args:
+                seg_data: TODO.
+            """
+        seg_st = seg_data[0]
+        data = seg_data[2]
+        chunk_len = int(config["seg_dur"] * fs)
+        chunk_shift = int(config["seg_shift"] * fs)
+        last_chunk_ed = 0
+        seg_res = []
+        for chunk_st in range(0, data.shape[0], chunk_shift):
+            chunk_ed = min(chunk_st + chunk_len, data.shape[0])
+            if chunk_ed <= last_chunk_ed:
+                break
+            last_chunk_ed = chunk_ed
+            chunk_st = max(0, chunk_ed - chunk_len)
+            chunk_data = data[chunk_st:chunk_ed]
+            if chunk_data.shape[0] < chunk_len:
+                chunk_data = np.pad(chunk_data, (0, chunk_len - chunk_data.shape[0]), "constant")
+            seg_res.append([chunk_st / fs + seg_st, chunk_ed / fs + seg_st, chunk_data])
+        return seg_res
+
+    segs = []
+    for i, s in enumerate(vad_segments):
+        segs.extend(seg_chunk(s))
+
+    return segs
+
+
+def extract_feature(audio):
+    """Extract feature.
+    
+        Args:
+            audio: TODO.
+        """
+    features = []
+    feature_times = []
+    feature_lengths = []
+    for au in audio:
+        feature = Kaldi.fbank(au.unsqueeze(0), num_mel_bins=80)
+        feature = feature - feature.mean(dim=0, keepdim=True)
+        features.append(feature)
+        feature_times.append(au.shape[0])
+        feature_lengths.append(feature.shape[0])
+    # padding for batch inference
+    features_padded = pad_list(features, pad_value=0)
+    # features = torch.cat(features)
+    return features_padded, feature_lengths, feature_times
+
+
+def postprocess(
+    segments: list,
+    vad_segments: list,
+    labels: np.ndarray,
+    embeddings: np.ndarray,
+    return_spk_center: bool = False,
+) -> Union[list, tuple]:
+    """Postprocess.
+    
+        Args:
+            segments: TODO.
+            vad_segments: TODO.
+            labels: TODO.
+            embeddings: TODO.
+        """
+    assert len(segments) == len(labels)
+    labels = correct_labels(labels)
+    distribute_res = []
+    for i in range(len(segments)):
+        distribute_res.append([segments[i][0], segments[i][1], labels[i]])
+    # merge the same speakers chronologically
+    distribute_res = merge_seque(distribute_res)
+
+    def is_overlapped(t1, t2):
+        """Is overlapped.
+        
+            Args:
+                t1: TODO.
+                t2: TODO.
+            """
+        if t1 > t2 + 1e-4:
+            return True
+        return False
+
+    # distribute the overlap region
+    for i in range(1, len(distribute_res)):
+        if is_overlapped(distribute_res[i - 1][1], distribute_res[i][0]):
+            p = (distribute_res[i][0] + distribute_res[i - 1][1]) / 2
+            distribute_res[i][0] = p
+            distribute_res[i - 1][1] = p
+
+    # smooth the result
+    distribute_res = smooth(distribute_res)
+
+    if return_spk_center:
+        # spk_embs[i] is the centroid (mean of clustered chunk embeddings) for
+        # corrected speaker label i, aligned with the `spk` ids in sentence_info.
+        # Computed lazily: only when the caller requests speaker centers.
+        spk_embs = np.stack(
+            [embeddings[labels == i].mean(0) for i in range(labels.max() + 1)]
+        )
+        return distribute_res, spk_embs
+    return distribute_res
+
+
+def correct_labels(labels):
+    """Correct labels.
+    
+        Args:
+            labels: TODO.
+        """
+    labels_id = 0
+    id2id = {}
+    new_labels = []
+    for i in labels:
+        if i not in id2id:
+            id2id[i] = labels_id
+            labels_id += 1
+        new_labels.append(id2id[i])
+    return np.array(new_labels)
+
+
+def merge_seque(distribute_res):
+    """Merge seque.
+    
+        Args:
+            distribute_res: TODO.
+        """
+    res = [distribute_res[0]]
+    for i in range(1, len(distribute_res)):
+        if distribute_res[i][2] != res[-1][2] or distribute_res[i][0] > res[-1][1]:
+            res.append(distribute_res[i])
+        else:
+            res[-1][1] = distribute_res[i][1]
+    return res
+
+
+def smooth(res, mindur=0.7):
+    # if only one segment, return directly
+    """Smooth.
+    
+        Args:
+            res: TODO.
+            mindur: TODO.
+        """
+    if len(res) < 2:
+        return res
+    # short segments are assigned to nearest speakers.
+    for i in range(len(res)):
+        res[i][0] = round(res[i][0], 2)
+        res[i][1] = round(res[i][1], 2)
+        if res[i][1] - res[i][0] < mindur:
+            if i == 0:
+                res[i][2] = res[i + 1][2]
+            elif i == len(res) - 1:
+                res[i][2] = res[i - 1][2]
+            elif res[i][0] - res[i - 1][1] <= res[i + 1][0] - res[i][1]:
+                res[i][2] = res[i - 1][2]
+            else:
+                res[i][2] = res[i + 1][2]
+    # merge the speakers
+    res = merge_seque(res)
+
+    return res
+
+
+def distribute_spk(sentence_list, sd_time_list):
+    """Distribute spk.
+    
+        Args:
+            sentence_list: TODO.
+            sd_time_list: TODO.
+        """
+    sd_time_list = [(spk_st * 1000, spk_ed * 1000, spk) for spk_st, spk_ed, spk in sd_time_list]
+    for d in sentence_list:
+        sentence_start = d['start']
+        sentence_end = d['end']
+        sentence_spk = 0
+        max_overlap = 0
+        for spk_st, spk_ed, spk in sd_time_list:
+            overlap = max(min(sentence_end, spk_ed) - max(sentence_start, spk_st), 0)
+            if overlap > max_overlap:
+                max_overlap = overlap
+                sentence_spk = spk
+            if overlap > 0 and sentence_spk == spk:
+                max_overlap += overlap
+        d['spk'] = int(sentence_spk)
+    return sentence_list
+
+
+class Storage(metaclass=ABCMeta):
+    """Abstract class of storage.
+
+    All backends need to implement two apis: ``read()`` and ``read_text()``.
+    ``read()`` reads the file as a byte stream and ``read_text()`` reads
+    the file as texts.
+    """
+
+    @abstractmethod
+    def read(self, filepath: str):
+        """Read.
+        
+            Args:
+                filepath: TODO.
+            """
+        pass
+
+    @abstractmethod
+    def read_text(self, filepath: str):
+        """Read text.
+        
+            Args:
+                filepath: TODO.
+            """
+        pass
+
+    @abstractmethod
+    def write(self, obj: bytes, filepath: Union[str, Path]) -> None:
+        """Write.
+        
+            Args:
+                obj: TODO.
+                filepath: TODO.
+            """
+        pass
+
+    @abstractmethod
+    def write_text(self, obj: str, filepath: Union[str, Path], encoding: str = "utf-8") -> None:
+        """Write text.
+        
+            Args:
+                obj: TODO.
+                filepath: TODO.
+                encoding: TODO.
+            """
+        pass
+
+
+class LocalStorage(Storage):
+    """Local hard disk storage"""
+
+    def read(self, filepath: Union[str, Path]) -> bytes:
+        """Read data from a given ``filepath`` with 'rb' mode.
+
+        Args:
+            filepath (str or Path): Path to read data.
+
+        Returns:
+            bytes: Expected bytes object.
+        """
+        with open(filepath, "rb") as f:
+            content = f.read()
+        return content
+
+    def read_text(self, filepath: Union[str, Path], encoding: str = "utf-8") -> str:
+        """Read data from a given ``filepath`` with 'r' mode.
+
+        Args:
+            filepath (str or Path): Path to read data.
+            encoding (str): The encoding format used to open the ``filepath``.
+                Default: 'utf-8'.
+
+        Returns:
+            str: Expected text reading from ``filepath``.
+        """
+        with open(filepath, "r", encoding=encoding) as f:
+            value_buf = f.read()
+        return value_buf
+
+    def write(self, obj: bytes, filepath: Union[str, Path]) -> None:
+        """Write data to a given ``filepath`` with 'wb' mode.
+
+        Note:
+            ``write`` will create a directory if the directory of ``filepath``
+            does not exist.
+
+        Args:
+            obj (bytes): Data to be written.
+            filepath (str or Path): Path to write data.
+        """
+        dirname = os.path.dirname(filepath)
+        if dirname and not os.path.exists(dirname):
+            os.makedirs(dirname, exist_ok=True)
+
+        with open(filepath, "wb") as f:
+            f.write(obj)
+
+    def write_text(self, obj: str, filepath: Union[str, Path], encoding: str = "utf-8") -> None:
+        """Write data to a given ``filepath`` with 'w' mode.
+
+        Note:
+            ``write_text`` will create a directory if the directory of
+            ``filepath`` does not exist.
+
+        Args:
+            obj (str): Data to be written.
+            filepath (str or Path): Path to write data.
+            encoding (str): The encoding format used to open the ``filepath``.
+                Default: 'utf-8'.
+        """
+        dirname = os.path.dirname(filepath)
+        if dirname and not os.path.exists(dirname):
+            os.makedirs(dirname, exist_ok=True)
+
+        with open(filepath, "w", encoding=encoding) as f:
+            f.write(obj)
+
+    @contextlib.contextmanager
+    def as_local_path(self, filepath: Union[str, Path]) -> Generator[Union[str, Path], None, None]:
+        """Only for unified API and do nothing."""
+        yield filepath
+
+
+class HTTPStorage(Storage):
+    """HTTP and HTTPS storage."""
+
+    def read(self, url):
+        # TODO @wenmeng.zwm add progress bar if file is too large
+        """Read.
+        
+            Args:
+                url: TODO.
+            """
+        r = requests.get(url)
+        r.raise_for_status()
+        return r.content
+
+    def read_text(self, url):
+        """Read text.
+        
+            Args:
+                url: TODO.
+            """
+        r = requests.get(url)
+        r.raise_for_status()
+        return r.text
+
+    @contextlib.contextmanager
+    def as_local_path(self, filepath: str) -> Generator[Union[str, Path], None, None]:
+        """Download a file from ``filepath``.
+
+        ``as_local_path`` is decorated by :meth:`contextlib.contextmanager`. It
+        can be called with ``with`` statement, and when exists from the
+        ``with`` statement, the temporary path will be released.
+
+        Args:
+            filepath (str): Download a file from ``filepath``.
+
+        Examples:
+            >>> storage = HTTPStorage()
+            >>> # After existing from the ``with`` clause,
+            >>> # the path will be removed
+            >>> with storage.get_local_path('http://path/to/file') as path:
+            ...     # do something here
+        """
+        try:
+            f = tempfile.NamedTemporaryFile(delete=False)
+            f.write(self.read(filepath))
+            f.close()
+            yield f.name
+        finally:
+            os.remove(f.name)
+
+    def write(self, obj: bytes, url: Union[str, Path]) -> None:
+        """Write.
+        
+            Args:
+                obj: TODO.
+                url: TODO.
+            """
+        raise NotImplementedError("write is not supported by HTTP Storage")
+
+    def write_text(self, obj: str, url: Union[str, Path], encoding: str = "utf-8") -> None:
+        """Write text.
+        
+            Args:
+                obj: TODO.
+                url: TODO.
+                encoding: TODO.
+            """
+        raise NotImplementedError("write_text is not supported by HTTP Storage")
+
+
+class OSSStorage(Storage):
+    """OSS storage."""
+
+    def __init__(self, oss_config_file=None):
+        # read from config file or env var
+        """Initialize OSSStorage.
+        
+            Args:
+                oss_config_file: TODO.
+            """
+        raise NotImplementedError("OSSStorage.__init__ to be implemented in the future")
+
+    def read(self, filepath):
+        """Read.
+        
+            Args:
+                filepath: TODO.
+            """
+        raise NotImplementedError("OSSStorage.read to be implemented in the future")
+
+    def read_text(self, filepath, encoding="utf-8"):
+        """Read text.
+        
+            Args:
+                filepath: TODO.
+                encoding: TODO.
+            """
+        raise NotImplementedError("OSSStorage.read_text to be implemented in the future")
+
+    @contextlib.contextmanager
+    def as_local_path(self, filepath: str) -> Generator[Union[str, Path], None, None]:
+        """Download a file from ``filepath``.
+
+        ``as_local_path`` is decorated by :meth:`contextlib.contextmanager`. It
+        can be called with ``with`` statement, and when exists from the
+        ``with`` statement, the temporary path will be released.
+
+        Args:
+            filepath (str): Download a file from ``filepath``.
+
+        Examples:
+            >>> storage = OSSStorage()
+            >>> # After existing from the ``with`` clause,
+            >>> # the path will be removed
+            >>> with storage.get_local_path('http://path/to/file') as path:
+            ...     # do something here
+        """
+        try:
+            f = tempfile.NamedTemporaryFile(delete=False)
+            f.write(self.read(filepath))
+            f.close()
+            yield f.name
+        finally:
+            os.remove(f.name)
+
+    def write(self, obj: bytes, filepath: Union[str, Path]) -> None:
+        """Write.
+        
+            Args:
+                obj: TODO.
+                filepath: TODO.
+            """
+        raise NotImplementedError("OSSStorage.write to be implemented in the future")
+
+    def write_text(self, obj: str, filepath: Union[str, Path], encoding: str = "utf-8") -> None:
+        """Write text.
+        
+            Args:
+                obj: TODO.
+                filepath: TODO.
+                encoding: TODO.
+            """
+        raise NotImplementedError("OSSStorage.write_text to be implemented in the future")
+
+
+G_STORAGES = {}
+
+
+class File(object):
+    _prefix_to_storage: dict = {
+        "oss": OSSStorage,
+        "http": HTTPStorage,
+        "https": HTTPStorage,
+        "local": LocalStorage,
+    }
+
+    @staticmethod
+    def _get_storage(uri):
+        """Internal: get storage.
+        
+            Args:
+                uri: TODO.
+            """
+        assert isinstance(uri, str), f"uri should be str type, but got {type(uri)}"
+
+        if "://" not in uri:
+            # local path
+            storage_type = "local"
+        else:
+            prefix, _ = uri.split("://")
+            storage_type = prefix
+
+        assert storage_type in File._prefix_to_storage, (
+            f"Unsupported uri {uri}, valid prefixs: " f"{list(File._prefix_to_storage.keys())}"
+        )
+
+        if storage_type not in G_STORAGES:
+            G_STORAGES[storage_type] = File._prefix_to_storage[storage_type]()
+
+        return G_STORAGES[storage_type]
+
+    @staticmethod
+    def read(uri: str) -> bytes:
+        """Read data from a given ``filepath`` with 'rb' mode.
+
+        Args:
+            filepath (str or Path): Path to read data.
+
+        Returns:
+            bytes: Expected bytes object.
+        """
+        storage = File._get_storage(uri)
+        return storage.read(uri)
+
+    @staticmethod
+    def read_text(uri: Union[str, Path], encoding: str = "utf-8") -> str:
+        """Read data from a given ``filepath`` with 'r' mode.
+
+        Args:
+            filepath (str or Path): Path to read data.
+            encoding (str): The encoding format used to open the ``filepath``.
+                Default: 'utf-8'.
+
+        Returns:
+            str: Expected text reading from ``filepath``.
+        """
+        storage = File._get_storage(uri)
+        return storage.read_text(uri)
+
+    @staticmethod
+    def write(obj: bytes, uri: Union[str, Path]) -> None:
+        """Write data to a given ``filepath`` with 'wb' mode.
+
+        Note:
+            ``write`` will create a directory if the directory of ``filepath``
+            does not exist.
+
+        Args:
+            obj (bytes): Data to be written.
+            filepath (str or Path): Path to write data.
+        """
+        storage = File._get_storage(uri)
+        return storage.write(obj, uri)
+
+    @staticmethod
+    def write_text(obj: str, uri: str, encoding: str = "utf-8") -> None:
+        """Write data to a given ``filepath`` with 'w' mode.
+
+        Note:
+            ``write_text`` will create a directory if the directory of
+            ``filepath`` does not exist.
+
+        Args:
+            obj (str): Data to be written.
+            filepath (str or Path): Path to write data.
+            encoding (str): The encoding format used to open the ``filepath``.
+                Default: 'utf-8'.
+        """
+        storage = File._get_storage(uri)
+        return storage.write_text(obj, uri)
+
+    @contextlib.contextmanager
+    def as_local_path(uri: str) -> Generator[Union[str, Path], None, None]:
+        """Only for unified API and do nothing."""
+        storage = File._get_storage(uri)
+        with storage.as_local_path(uri) as local_path:
+            yield local_path