add_eres2net_speaker_diarization

Link: https://code.alibaba-inc.com/Ali-MaaS/MaaS-lib/codereview/14412468
* add_eres2net_speaker_diarization

* fix lint issue

* update code

* update pipeline

modelscope/metainfo.py

@@ -200,6 +200,7 @@ class Models(object):
     eres2net_sv = 'eres2net-sv'
     eres2net_aug_sv = 'eres2net-aug-sv'
     scl_sd = 'scl-sd'
+    scl_sd_xvector = 'scl-sd-xvector'
     campplus_lre = 'cam++-lre'
     eres2net_lre = 'eres2net-lre'
     cluster_backend = 'cluster-backend'

modelscope/models/audio/sv/ERes2Net.py

@@ -8,6 +8,7 @@ import math
 import os
 from typing import Any, Dict, Union
 
+import numpy as np
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -323,13 +324,18 @@ class SpeakerVerificationERes2Net(TorchModel):
         self.embedding_model.eval()
 
     def forward(self, audio):
-        assert len(audio.shape) == 2 and audio.shape[
-            0] == 1, 'modelscope error: the shape of input audio to model needs to be [1, T]'
-        # audio shape: [1, T]
+        if isinstance(audio, np.ndarray):
+            audio = torch.from_numpy(audio)
+        if len(audio.shape) == 1:
+            audio = audio.unsqueeze(0)
+        assert len(
+            audio.shape
+        ) == 2, 'modelscope error: the shape of input audio to model needs to be [N, T]'
+        # audio shape: [N, T]
         feature = self.__extract_feature(audio)
         embedding = self.embedding_model(feature)
 
-        return embedding
+        return embedding.detach().cpu()
 
     def __extract_feature(self, audio):
         feature = Kaldi.fbank(audio, num_mel_bins=self.feature_dim)

modelscope/models/audio/sv/ERes2Net_aug.py

@@ -8,6 +8,7 @@ import math
 import os
 from typing import Any, Dict, Union
 
+import numpy as np
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -316,13 +317,18 @@ class SpeakerVerificationERes2Net(TorchModel):
         self.embedding_model.eval()
 
     def forward(self, audio):
-        assert len(audio.shape) == 2 and audio.shape[
-            0] == 1, 'modelscope error: the shape of input audio to model needs to be [1, T]'
-        # audio shape: [1, T]
+        if isinstance(audio, np.ndarray):
+            audio = torch.from_numpy(audio)
+        if len(audio.shape) == 1:
+            audio = audio.unsqueeze(0)
+        assert len(
+            audio.shape
+        ) == 2, 'modelscope error: the shape of input audio to model needs to be [N, T]'
+        # audio shape: [N, T]
         feature = self.__extract_feature(audio)
         embedding = self.embedding_model(feature)
 
-        return embedding
+        return embedding.detach().cpu()
 
     def __extract_feature(self, audio):
         feature = Kaldi.fbank(audio, num_mel_bins=self.feature_dim)

modelscope/models/audio/sv/TDNN.py

@@ -0,0 +1,303 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class Conv1d_O(nn.Module):
+
+    def __init__(
+        self,
+        out_channels,
+        kernel_size,
+        input_shape=None,
+        in_channels=None,
+        stride=1,
+        dilation=1,
+        padding='same',
+        groups=1,
+        bias=True,
+        padding_mode='reflect',
+        skip_transpose=False,
+    ):
+        super().__init__()
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.dilation = dilation
+        self.padding = padding
+        self.padding_mode = padding_mode
+        self.unsqueeze = False
+        self.skip_transpose = skip_transpose
+
+        if input_shape is None and in_channels is None:
+            raise ValueError('Must provide one of input_shape or in_channels')
+
+        if in_channels is None:
+            in_channels = self._check_input_shape(input_shape)
+
+        self.conv = nn.Conv1d(
+            in_channels,
+            out_channels,
+            self.kernel_size,
+            stride=self.stride,
+            dilation=self.dilation,
+            padding=0,
+            groups=groups,
+            bias=bias,
+        )
+
+    def forward(self, x):
+        """Returns the output of the convolution.
+
+        Arguments
+        ---------
+        x : torch.Tensor (batch, time, channel)
+            input to convolve. 2d or 4d tensors are expected.
+        """
+
+        if not self.skip_transpose:
+            x = x.transpose(1, -1)
+
+        if self.unsqueeze:
+            x = x.unsqueeze(1)
+
+        if self.padding == 'same':
+            x = self._manage_padding(x, self.kernel_size, self.dilation,
+                                     self.stride)
+
+        elif self.padding == 'causal':
+            num_pad = (self.kernel_size - 1) * self.dilation
+            x = F.pad(x, (num_pad, 0))
+
+        elif self.padding == 'valid':
+            pass
+
+        else:
+            raise ValueError(
+                "Padding must be 'same', 'valid' or 'causal'. Got "
+                + self.padding)
+
+        wx = self.conv(x)
+
+        if self.unsqueeze:
+            wx = wx.squeeze(1)
+
+        if not self.skip_transpose:
+            wx = wx.transpose(1, -1)
+
+        return wx
+
+    def _manage_padding(
+        self,
+        x,
+        kernel_size: int,
+        dilation: int,
+        stride: int,
+    ):
+        # Detecting input shape
+        L_in = x.shape[-1]
+
+        # Time padding
+        padding = get_padding_elem(L_in, stride, kernel_size, dilation)
+
+        # Applying padding
+        x = F.pad(x, padding, mode=self.padding_mode)
+
+        return x
+
+    def _check_input_shape(self, shape):
+        """Checks the input shape and returns the number of input channels.
+        """
+
+        if len(shape) == 2:
+            self.unsqueeze = True
+            in_channels = 1
+        elif self.skip_transpose:
+            in_channels = shape[1]
+        elif len(shape) == 3:
+            in_channels = shape[2]
+        else:
+            raise ValueError('conv1d expects 2d, 3d inputs. Got '
+                             + str(len(shape)))
+
+        # Kernel size must be odd
+        if self.kernel_size % 2 == 0:
+            raise ValueError(
+                'The field kernel size must be an odd number. Got %s.' %
+                (self.kernel_size))
+        return in_channels
+
+
+# Skip transpose as much as possible for efficiency
+class Conv1d(Conv1d_O):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(skip_transpose=True, *args, **kwargs)
+
+
+def get_padding_elem(L_in: int, stride: int, kernel_size: int, dilation: int):
+    """This function computes the number of elements to add for zero-padding.
+
+    Arguments
+    ---------
+    L_in : int
+    stride: int
+    kernel_size : int
+    dilation : int
+    """
+    if stride > 1:
+        n_steps = math.ceil(((L_in - kernel_size * dilation) / stride) + 1)
+        L_out = stride * (n_steps - 1) + kernel_size * dilation
+        padding = [kernel_size // 2, kernel_size // 2]
+
+    else:
+        L_out = (L_in - dilation * (kernel_size - 1) - 1) // stride + 1
+
+        padding = [(L_in - L_out) // 2, (L_in - L_out) // 2]
+    return padding
+
+
+class BatchNorm1d_O(nn.Module):
+
+    def __init__(
+        self,
+        input_shape=None,
+        input_size=None,
+        eps=1e-05,
+        momentum=0.1,
+        affine=True,
+        track_running_stats=True,
+        combine_batch_time=False,
+        skip_transpose=False,
+    ):
+        super().__init__()
+        self.combine_batch_time = combine_batch_time
+        self.skip_transpose = skip_transpose
+
+        if input_size is None and skip_transpose:
+            input_size = input_shape[1]
+        elif input_size is None:
+            input_size = input_shape[-1]
+
+        self.norm = nn.BatchNorm1d(
+            input_size,
+            eps=eps,
+            momentum=momentum,
+            affine=affine,
+            track_running_stats=track_running_stats,
+        )
+
+    def forward(self, x):
+        """Returns the normalized input tensor.
+
+        Arguments
+        ---------
+        x : torch.Tensor (batch, time, [channels])
+            input to normalize. 2d or 3d tensors are expected in input
+            4d tensors can be used when combine_dims=True.
+        """
+        shape_or = x.shape
+        if self.combine_batch_time:
+            if x.ndim == 3:
+                x = x.reshape(shape_or[0] * shape_or[1], shape_or[2])
+            else:
+                x = x.reshape(shape_or[0] * shape_or[1], shape_or[3],
+                              shape_or[2])
+
+        elif not self.skip_transpose:
+            x = x.transpose(-1, 1)
+
+        x_n = self.norm(x)
+
+        if self.combine_batch_time:
+            x_n = x_n.reshape(shape_or)
+        elif not self.skip_transpose:
+            x_n = x_n.transpose(1, -1)
+
+        return x_n
+
+
+class BatchNorm1d(BatchNorm1d_O):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(skip_transpose=True, *args, **kwargs)
+
+
+class Xvector(torch.nn.Module):
+    """This model extracts X-vectors for speaker recognition and diarization.
+
+    Arguments
+    ---------
+    device : str
+        Device used e.g. "cpu" or "cuda".
+    activation : torch class
+        A class for constructing the activation layers.
+    tdnn_blocks : int
+        Number of time-delay neural (TDNN) layers.
+    tdnn_channels : list of ints
+        Output channels for TDNN layer.
+    tdnn_kernel_sizes : list of ints
+        List of kernel sizes for each TDNN layer.
+    tdnn_dilations : list of ints
+        List of dilations for kernels in each TDNN layer.
+    lin_neurons : int
+        Number of neurons in linear layers.
+
+    Example
+    -------
+    >>> compute_xvect = Xvector('cpu')
+    >>> input_feats = torch.rand([5, 10, 40])
+    >>> outputs = compute_xvect(input_feats)
+    >>> outputs.shape
+    torch.Size([5, 1, 512])
+    """
+
+    def __init__(
+        self,
+        device='cpu',
+        activation=torch.nn.LeakyReLU,
+        tdnn_blocks=5,
+        tdnn_channels=[512, 512, 512, 512, 1500],
+        tdnn_kernel_sizes=[5, 3, 3, 1, 1],
+        tdnn_dilations=[1, 2, 3, 1, 1],
+        lin_neurons=512,
+        in_channels=80,
+    ):
+
+        super().__init__()
+        self.blocks = nn.ModuleList()
+
+        # TDNN layers
+        for block_index in range(tdnn_blocks):
+            out_channels = tdnn_channels[block_index]
+            self.blocks.extend([
+                Conv1d(
+                    in_channels=in_channels,
+                    out_channels=out_channels,
+                    kernel_size=tdnn_kernel_sizes[block_index],
+                    dilation=tdnn_dilations[block_index],
+                ),
+                activation(),
+                BatchNorm1d(input_size=out_channels),
+            ])
+            in_channels = tdnn_channels[block_index]
+
+    def forward(self, x, lens=None):
+        """Returns the x-vectors.
+
+        Arguments
+        ---------
+        x : torch.Tensor
+        """
+
+        x = x.transpose(1, 2)
+
+        for layer in self.blocks:
+            try:
+                x = layer(x, lengths=lens)
+            except TypeError:
+                x = layer(x)
+        x = x.transpose(1, 2)
+        return x

modelscope/models/audio/sv/speaker_change_locator_xvector.py

@@ -0,0 +1,329 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+
+import os
+from collections import OrderedDict
+from typing import Any, Dict, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchaudio.compliance.kaldi as Kaldi
+
+from modelscope.metainfo import Models
+from modelscope.models import MODELS, TorchModel
+from modelscope.models.audio.sv.TDNN import Xvector
+from modelscope.utils.constant import Tasks
+from modelscope.utils.device import create_device
+
+
+class MultiHeadSelfAttention(nn.Module):
+
+    def __init__(self, n_units, h=8, dropout=0.1):
+        super(MultiHeadSelfAttention, self).__init__()
+        self.linearQ = nn.Linear(n_units, n_units)
+        self.linearK = nn.Linear(n_units, n_units)
+        self.linearV = nn.Linear(n_units, n_units)
+        self.linearO = nn.Linear(n_units, n_units)
+        self.d_k = n_units // h
+        self.h = h
+        self.dropout = nn.Dropout(p=dropout)
+        self.att = None
+
+    def forward(self, x, batch_size):
+        # x: (BT, F)
+        q = self.linearQ(x).reshape(batch_size, -1, self.h, self.d_k)
+        k = self.linearK(x).reshape(batch_size, -1, self.h, self.d_k)
+        v = self.linearV(x).reshape(batch_size, -1, self.h, self.d_k)
+        scores = torch.matmul(q.transpose(1, 2), k.permute(
+            0, 2, 3, 1)) / np.sqrt(self.d_k)
+        # scores: (B, h, T, T)
+        self.att = F.softmax(scores, dim=3)
+        p_att = self.dropout(self.att)
+        # v : (B, T, h, d_k)
+        # p_att : (B, h, T, T)
+        x = torch.matmul(p_att, v.transpose(1, 2))
+        # x : (B, h, T, d_k)
+        x = x.transpose(1, 2).reshape(-1, self.h * self.d_k)
+        return self.linearO(x)
+
+
+class PositionwiseFeedForward(nn.Module):
+
+    def __init__(self, n_units, d_units, dropout):
+        super(PositionwiseFeedForward, self).__init__()
+        self.linear1 = nn.Linear(n_units, d_units)
+        self.linear2 = nn.Linear(d_units, n_units)
+        self.dropout = nn.Dropout(p=dropout)
+
+    def forward(self, x):
+        return self.linear2(self.dropout(F.relu(self.linear1(x))))
+
+
+class PosEncoding(nn.Module):
+
+    def __init__(self, max_seq_len, d_word_vec):
+        super(PosEncoding, self).__init__()
+        pos_enc = np.array([[
+            pos / np.power(10000, 2.0 * (j // 2) / d_word_vec)
+            for j in range(d_word_vec)
+        ] for pos in range(max_seq_len)])
+        pos_enc[:, 0::2] = np.sin(pos_enc[:, 0::2])
+        pos_enc[:, 1::2] = np.cos(pos_enc[:, 1::2])
+        pad_row = np.zeros([1, d_word_vec])
+        pos_enc = np.concatenate([pad_row, pos_enc]).astype(np.float32)
+
+        self.pos_enc = torch.nn.Embedding(max_seq_len + 1, d_word_vec)
+        self.pos_enc.weight = torch.nn.Parameter(
+            torch.from_numpy(pos_enc), requires_grad=False)
+
+    def forward(self, input_len):
+        max_len = torch.max(input_len)
+        input_pos = torch.LongTensor([
+            list(range(1, len + 1)) + [0] * (max_len - len)
+            for len in input_len
+        ])
+
+        input_pos = input_pos.to(list(self.pos_enc.parameters())[0].device)
+        return self.pos_enc(input_pos)
+
+
+class TransformerEncoder(nn.Module):
+
+    def __init__(self,
+                 idim,
+                 n_units=256,
+                 n_layers=2,
+                 e_units=512,
+                 h=4,
+                 dropout=0.1):
+        super(TransformerEncoder, self).__init__()
+        self.linear_in = nn.Linear(idim, n_units)
+        self.lnorm_in = nn.LayerNorm(n_units)
+
+        self.n_layers = n_layers
+        self.dropout = nn.Dropout(p=dropout)
+        for i in range(n_layers):
+            setattr(self, '{}{:d}'.format('lnorm1_', i), nn.LayerNorm(n_units))
+            setattr(self, '{}{:d}'.format('self_att_', i),
+                    MultiHeadSelfAttention(n_units, h))
+            setattr(self, '{}{:d}'.format('lnorm2_', i), nn.LayerNorm(n_units))
+            setattr(self, '{}{:d}'.format('ff_', i),
+                    PositionwiseFeedForward(n_units, e_units, dropout))
+        self.lnorm_out = nn.LayerNorm(n_units)
+
+    def forward(self, x):
+        # x: [B, num_anchors, T, n_in]
+        bs, num, tframe, dim = x.size()
+        x = x.reshape(bs * num, tframe, -1)  # [B*num_anchors, T, dim]
+        # x: (B, T, F) ... batch, time, (mel)freq
+        B_size, T_size, _ = x.shape
+        # e: (BT, F)
+        e = self.linear_in(x.reshape(B_size * T_size, -1))
+        # Encoder stack
+        for i in range(self.n_layers):
+            # layer normalization
+            e = getattr(self, '{}{:d}'.format('lnorm1_', i))(e)
+            # self-attention
+            s = getattr(self, '{}{:d}'.format('self_att_', i))(e, x.shape[0])
+            # residual
+            e = e + self.dropout(s)
+            # layer normalization
+            e = getattr(self, '{}{:d}'.format('lnorm2_', i))(e)
+            # positionwise feed-forward
+            s = getattr(self, '{}{:d}'.format('ff_', i))(e)
+            # residual
+            e = e + self.dropout(s)
+        # final layer normalization
+        # output: (BT, F)
+        # output: (B, F, T)
+        output = self.lnorm_out(e).reshape(B_size, T_size, -1)
+        output = output.reshape(bs, num, tframe,
+                                -1)  # [B, num_anchors, T, dim]
+        return output
+
+
+class TransformerEncoder_out(nn.Module):
+
+    def __init__(self,
+                 idim,
+                 n_units=256,
+                 n_layers=2,
+                 e_units=512,
+                 h=4,
+                 dropout=0.1):
+        super(TransformerEncoder_out, self).__init__()
+        self.linear_in = nn.Linear(idim, n_units)
+        self.lnorm_in = nn.LayerNorm(n_units)
+
+        self.n_layers = n_layers
+        self.dropout = nn.Dropout(p=dropout)
+        for i in range(n_layers):
+            setattr(self, '{}{:d}'.format('lnorm1_', i), nn.LayerNorm(n_units))
+            setattr(self, '{}{:d}'.format('self_att_', i),
+                    MultiHeadSelfAttention(n_units, h))
+            setattr(self, '{}{:d}'.format('lnorm2_', i), nn.LayerNorm(n_units))
+            setattr(self, '{}{:d}'.format('ff_', i),
+                    PositionwiseFeedForward(n_units, e_units, dropout))
+        self.lnorm_out = nn.LayerNorm(n_units)
+
+    def forward(self, x):
+        # x: (B, T, F)
+        B_size, T_size, _ = x.shape
+        # e: (BT, F)
+        e = self.linear_in(x.reshape(B_size * T_size, -1))
+        # Encoder stack
+        for i in range(self.n_layers):
+            # layer normalization
+            e = getattr(self, '{}{:d}'.format('lnorm1_', i))(e)
+            # self-attention
+            s = getattr(self, '{}{:d}'.format('self_att_', i))(e, x.shape[0])
+            # residual
+            e = e + self.dropout(s)
+            # layer normalization
+            e = getattr(self, '{}{:d}'.format('lnorm2_', i))(e)
+            # positionwise feed-forward
+            s = getattr(self, '{}{:d}'.format('ff_', i))(e)
+            # residual
+            e = e + self.dropout(s)
+        # final layer normalization
+        # output: (BT, F)
+        # output: (B, T, F)
+        output = self.lnorm_out(e).reshape(B_size, T_size, -1)
+        return output
+
+
+class OutLayer(nn.Module):
+
+    def __init__(self, n_units=256, num_anchors=2):
+        super(OutLayer, self).__init__()
+        self.rnn_combine = TransformerEncoder_out(num_anchors * n_units,
+                                                  n_units)
+        self.out_linear = nn.Linear(n_units // num_anchors, 1)
+
+    def forward(self, input):
+        # input: [B, num_anchors, T, dim]
+        bs, num, tframe, dim = input.size()
+        output = input.permute(0, 2, 1,
+                               3).reshape(bs, tframe,
+                                          -1)  # [Bs, t, num_anchors*dim]
+        output = self.rnn_combine(output)  # [Bs, t, n_units]
+        output = output.reshape(
+            bs, tframe, num, -1)  # [Bs, t, num_anchors, n_units//num_anchors]
+        output = self.out_linear(output).squeeze(-1)  # [Bs, t, num_anchors]
+
+        return output
+
+
+class TransformerDetector(nn.Module):
+
+    def __init__(self,
+                 frame_dim=512,
+                 anchor_dim=192,
+                 hidden_dim=256,
+                 max_seq_len=500):
+        super(TransformerDetector, self).__init__()
+        self.detection = TransformerEncoder(
+            idim=frame_dim + anchor_dim, n_units=hidden_dim)
+        self.output = OutLayer(n_units=hidden_dim)
+        self.pos_enc = PosEncoding(max_seq_len, hidden_dim)
+
+    def forward(self, feats, anchors):
+        # feats: [1, t, fdim]
+        num_frames = feats.shape[1]
+        num_anchors = anchors.shape[1]
+        bs = feats.shape[0]
+        feats = feats.unsqueeze(1).repeat(
+            1, num_anchors, 1, 1)  # shape: [Bs, num_anchors, t, fdim]
+        anchors = anchors.unsqueeze(2).repeat(
+            1, 1, num_frames, 1)  # shape: [Bs, num_anchors, t, xdim]
+        sd_in = torch.cat((feats, anchors),
+                          dim=-1)  # shape: [Bs, num_anchors, t, fdim+xdim]
+        sd_out = self.detection(sd_in)  # shape: [Bs, num_anchors, t, sd_dim]
+
+        # pos
+        pos_emb = self.pos_enc(torch.tensor([num_frames] * (bs * num_anchors)))
+        pos_emb = pos_emb.reshape(bs, num_anchors, num_frames, -1)
+        sd_out += pos_emb
+
+        # output
+        output = self.output(sd_out)  # shape: [Bs, t, num_anchors]
+
+        return output
+
+
+@MODELS.register_module(
+    Tasks.speaker_diarization, module_name=Models.scl_sd_xvector)
+class SpeakerChangeLocatorTransformer(TorchModel):
+    r"""A speaekr change locator using the transformer architecture as the backbone.
+    Args:
+        model_dir: A model dir.
+        model_config: The model config.
+    """
+
+    def __init__(self, model_dir, model_config: Dict[str, Any], *args,
+                 **kwargs):
+        super().__init__(model_dir, model_config, *args, **kwargs)
+        self.model_config = model_config
+
+        self.feature_dim = self.model_config['fbank_dim']
+        frame_size = self.model_config['frame_size']
+        anchor_size = self.model_config['anchor_size']
+        self.device = create_device(kwargs['device'])
+
+        self.encoder = Xvector(in_channels=self.feature_dim)
+        self.backend = TransformerDetector(
+            frame_dim=frame_size, anchor_dim=anchor_size)
+
+        pretrained_encoder = kwargs['pretrained_encoder']
+        pretrained_backend = kwargs['pretrained_backend']
+
+        self.__load_check_point(pretrained_encoder, pretrained_backend)
+
+        self.encoder.to(self.device)
+        self.backend.to(self.device)
+        self.encoder.eval()
+        self.backend.eval()
+
+    def forward(self, audio, anchors):
+        if isinstance(audio, np.ndarray):
+            audio = torch.from_numpy(audio)
+        if isinstance(anchors, np.ndarray):
+            anchors = torch.from_numpy(anchors)
+        assert len(audio.shape) == 2 and audio.shape[
+            0] == 1, 'modelscope error: the shape of input audio to model needs to be [1, T]'
+        assert len(
+            anchors.shape
+        ) == 3 and anchors.shape[0] == 1 and anchors.shape[
+            1] == 2, 'modelscope error: the shape of input anchors to model needs to be [1, 2, D]'
+        # audio shape: [1, T]
+        feature = self.__extract_feature(audio)
+        frame_state = self.encoder(feature.to(self.device))
+        output = self.backend(frame_state, anchors.to(self.device))
+        output = output.squeeze(0).detach().cpu().sigmoid()
+
+        time_scale_factor = int(np.ceil(feature.shape[1] / output.shape[0]))
+        output = output.unsqueeze(1).expand(-1, time_scale_factor,
+                                            -1).reshape(-1, output.shape[-1])
+        return output
+
+    def __extract_feature(self, audio):
+        feature = Kaldi.fbank(audio, num_mel_bins=self.feature_dim)
+        feature = feature - feature.mean(dim=0, keepdim=True)
+        feature = feature.unsqueeze(0)
+        return feature
+
+    def __load_check_point(
+        self,
+        pretrained_encoder,
+        pretrained_backend,
+    ):
+        self.encoder.load_state_dict(
+            torch.load(
+                os.path.join(self.model_dir, pretrained_encoder),
+                map_location=torch.device('cpu')))
+
+        self.backend.load_state_dict(
+            torch.load(
+                os.path.join(self.model_dir, pretrained_backend),
+                map_location=torch.device('cpu')))

modelscope/pipelines/audio/segmentation_clustering_pipeline.py

@@ -92,8 +92,9 @@ class SegmentationClusteringPipeline(Pipeline):
     def forward(self, input: list) -> np.ndarray:
         embeddings = []
         for s in input:
-            _, embs = self.sv_pipeline([s[2]], output_emb=True)
-            embeddings.append(embs)
+            save_dict = self.sv_pipeline([s[2]], output_emb=True)
+            if save_dict['embs'].shape == (1, 192):
+                embeddings.append(save_dict['embs'])
         embeddings = np.concatenate(embeddings)
         return embeddings
 

modelscope/pipelines/audio/speaker_verification_eres2net_pipeline.py

@@ -3,8 +3,10 @@
 import io
 from typing import Any, Dict, List, Union
 
+import numpy as np
 import soundfile as sf
 import torch
+import torchaudio
 
 from modelscope.fileio import File
 from modelscope.metainfo import Pipelines
@@ -46,64 +48,111 @@ class ERes2Net_Pipeline(Pipeline):
         self.model_config = self.model.model_config
         self.config = self.model.other_config
         self.thr = self.config['yesOrno_thr']
+        self.save_dict = {}
 
     def __call__(self,
-                 in_audios: List[str],
-                 thr: float = None) -> Dict[str, Any]:
+                 in_audios: Union[np.ndarray, list],
+                 save_dir: str = None,
+                 output_emb: bool = False,
+                 thr: float = None):
         if thr is not None:
             self.thr = thr
         if self.thr < -1 or self.thr > 1:
             raise ValueError(
                 'modelscope error: the thr value should be in [-1, 1], but found to be %f.'
                 % self.thr)
-        outputs = self.preprocess(in_audios)
-        outputs = self.forward(outputs)
-        outputs = self.postprocess(outputs)
-
-        return outputs
-
-    def forward(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
-        emb1 = self.model(inputs['data1'])
-        emb2 = self.model(inputs['data2'])
-
-        return {'emb1': emb1, 'emb2': emb2}
-
-    def postprocess(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
-        score = self.compute_cos_similarity(inputs['emb1'], inputs['emb2'])
-        score = round(score, 5)
-        if score >= self.thr:
-            ans = 'yes'
+        wavs = self.preprocess(in_audios)
+        embs = self.forward(wavs)
+        outputs = self.postprocess(embs, in_audios, save_dir)
+        if output_emb:
+            self.save_dict['outputs'] = outputs
+            self.save_dict['embs'] = embs.numpy()
+            return self.save_dict
         else:
-            ans = 'no'
+            return outputs
 
-        return {OutputKeys.SCORE: score, OutputKeys.TEXT: ans}
+    def forward(self, inputs: list):
+        embs = []
+        for x in inputs:
+            embs.append(self.model(x))
+        embs = torch.cat(embs)
+        return embs
 
-    def preprocess(self, inputs: List[str],
-                   **preprocess_params) -> Dict[str, Any]:
-        if len(inputs) != 2:
-            raise ValueError(
-                'modelscope error: Two input audio files are required.')
-        output = {}
+    def postprocess(self,
+                    inputs: torch.Tensor,
+                    in_audios: Union[np.ndarray, list],
+                    save_dir=None):
+        if isinstance(in_audios[0], str) and save_dir is not None:
+            # save the embeddings
+            os.makedirs(save_dir, exist_ok=True)
+            for i, p in enumerate(in_audios):
+                save_path = os.path.join(
+                    save_dir, '%s.npy' %
+                    (os.path.basename(p).rsplit('.', 1)[0]))
+                np.save(save_path, inputs[i].numpy())
+
+        if len(inputs) == 2:
+            # compute the score
+            score = self.compute_cos_similarity(inputs[0], inputs[1])
+            score = round(score, 5)
+            if score >= self.thr:
+                ans = 'yes'
+            else:
+                ans = 'no'
+            output = {OutputKeys.SCORE: score, OutputKeys.TEXT: ans}
+        else:
+            output = {OutputKeys.TEXT: 'No similarity score output'}
+
+        return output
+
+    def preprocess(self, inputs: Union[np.ndarray, list]):
+        output = []
         for i in range(len(inputs)):
             if isinstance(inputs[i], str):
                 file_bytes = File.read(inputs[i])
                 data, fs = sf.read(io.BytesIO(file_bytes), dtype='float32')
                 if len(data.shape) == 2:
                     data = data[:, 0]
+                data = torch.from_numpy(data).unsqueeze(0)
                 if fs != self.model_config['sample_rate']:
-                    raise ValueError(
-                        'modelscope error: Only support %d sample rate files'
-                        % self.model_cfg['sample_rate'])
-                output['data%d' %
-                       (i + 1)] = torch.from_numpy(data).unsqueeze(0)
+                    logger.warning(
+                        'The sample rate of audio is not %d, resample it.'
+                        % self.model_config['sample_rate'])
+                    data, fs = torchaudio.sox_effects.apply_effects_tensor(
+                        data,
+                        fs,
+                        effects=[[
+                            'rate',
+                            str(self.model_config['sample_rate'])
+                        ]])
+                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 temporarily restricted to audio file address'
-                    % i)
+                    'modelscope error: The input type is restricted to audio address and nump array.'
+                )
+            output.append(data)
         return output
 
-    def compute_cos_similarity(self, emb1: torch.Tensor,
-                               emb2: torch.Tensor) -> float:
+    def compute_cos_similarity(self, emb1: Union[np.ndarray, torch.Tensor],
+                               emb2: Union[np.ndarray, torch.Tensor]) -> float:
+        if isinstance(emb1, np.ndarray):
+            emb1 = torch.from_numpy(emb1)
+        if isinstance(emb2, np.ndarray):
+            emb2 = torch.from_numpy(emb2)
+        if len(emb1.shape):
+            emb1 = emb1.unsqueeze(0)
+        if len(emb2.shape):
+            emb2 = emb2.unsqueeze(0)
         assert len(emb1.shape) == 2 and len(emb2.shape) == 2
         cos = torch.nn.CosineSimilarity(dim=1, eps=1e-6)
         cosine = cos(emb1, emb2)

modelscope/pipelines/audio/speaker_verification_light_pipeline.py

@@ -50,6 +50,7 @@ class SpeakerVerificationPipeline(Pipeline):
         self.model_config = self.model.model_config
         self.config = self.model.other_config
         self.thr = self.config['yesOrno_thr']
+        self.save_dict = {}
 
     def __call__(self,
                  in_audios: Union[np.ndarray, list],
@@ -66,7 +67,9 @@ class SpeakerVerificationPipeline(Pipeline):
         embs = self.forward(wavs)
         outputs = self.postprocess(embs, in_audios, save_dir)
         if output_emb:
-            return outputs, embs.numpy()
+            self.save_dict['outputs'] = outputs
+            self.save_dict['embs'] = embs.numpy()
+            return self.save_dict
         else:
             return outputs
 

tests/pipelines/test_speaker_verification.py

@@ -23,8 +23,10 @@ class SpeakerVerificationTest(unittest.TestCase):
     campplus_voxceleb_16k_model_id = 'damo/speech_campplus_sv_en_voxceleb_16k'
     rdino_voxceleb_16k_model_id = 'damo/speech_rdino_ecapa_tdnn_sv_en_voxceleb_16k'
     speaker_change_locating_cn_model_id = 'damo/speech_campplus-transformer_scl_zh-cn_16k-common'
+    speaker_change_lcoating_xvector_cn_model_id = 'damo/speech_xvector_transformer_scl_zh-cn_16k-common'
     eres2net_voxceleb_16k_model_id = 'damo/speech_eres2net_sv_en_voxceleb_16k'
     speaker_diarization_model_id = 'damo/speech_campplus_speaker-diarization_common'
+    speaker_diarization_eres2net_model_id = 'damo/speech_eres2net-large_speaker-diarization_common'
     lre_campplus_en_cn_16k_model_id = 'damo/speech_campplus_lre_en-cn_16k'
     lre_eres2net_base_en_cn_16k_model_id = 'damo/speech_eres2net_base_lre_en-cn_16k'
     lre_eres2net_large_en_cn_16k_model_id = 'damo/speech_eres2net_large_lre_en-cn_16k'
@@ -123,6 +125,17 @@ class SpeakerVerificationTest(unittest.TestCase):
         print(result)
         self.assertTrue(OutputKeys.TEXT in result)
 
+    @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
+    def test_run_with_speaker_change_locating_xvector_cn_16k(self):
+        logger.info(
+            'Run speaker change locating for xvector-transformer model')
+        result = self.run_pipeline(
+            model_id=self.speaker_change_lcoating_xvector_cn_model_id,
+            task=Tasks.speaker_diarization,
+            audios=SCL_EXAMPLE_WAV)
+        print(result)
+        self.assertTrue(OutputKeys.TEXT in result)
+
     @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
     def test_run_with_speaker_verification_eres2net_voxceleb_16k(self):
         logger.info('Run speaker verification for eres2net_voxceleb_16k model')
@@ -140,7 +153,7 @@ class SpeakerVerificationTest(unittest.TestCase):
         result = self.run_pipeline(
             model_id=self.eres2net_aug_zh_cn_16k_common_model_id,
             audios=[SPEAKER1_A_EN_16K_WAV, SPEAKER1_B_EN_16K_WAV],
-            model_revision='v1.0.4')
+            model_revision='v1.0.5')
         print(result)
         self.assertTrue(OutputKeys.SCORE in result)
 
@@ -154,6 +167,16 @@ class SpeakerVerificationTest(unittest.TestCase):
         print(result)
         self.assertTrue(OutputKeys.TEXT in result)
 
+    @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
+    def test_run_with_eres2net_speaker_diarization_common(self):
+        logger.info('Run eres2net speaker diarization task')
+        result = self.run_pipeline(
+            model_id=self.speaker_diarization_eres2net_model_id,
+            task=Tasks.speaker_diarization,
+            audios=SD_EXAMPLE_WAV)
+        print(result)
+        self.assertTrue(OutputKeys.TEXT in result)
+
     @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
     def test_run_with_language_recognition_campplus_en_cn_16k(self):
         logger.info('Run language recognition for campplus_en_cn_16k')