From e8b34a3bcdc2d0d919b2156a81193f5c58bc086d Mon Sep 17 00:00:00 2001 From: "chenyafeng.cyf" Date: Wed, 25 Oct 2023 20:36:36 +0800 Subject: [PATCH] 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 | 1 + modelscope/models/audio/sv/ERes2Net.py | 14 +- modelscope/models/audio/sv/ERes2Net_aug.py | 14 +- modelscope/models/audio/sv/TDNN.py | 303 ++++++++++++++++ .../sv/speaker_change_locator_xvector.py | 329 ++++++++++++++++++ .../audio/segmentation_clustering_pipeline.py | 5 +- .../speaker_verification_eres2net_pipeline.py | 121 +++++-- .../speaker_verification_light_pipeline.py | 5 +- tests/pipelines/test_speaker_verification.py | 25 +- 9 files changed, 769 insertions(+), 48 deletions(-) create mode 100644 modelscope/models/audio/sv/TDNN.py create mode 100644 modelscope/models/audio/sv/speaker_change_locator_xvector.py diff --git a/modelscope/metainfo.py b/modelscope/metainfo.py index 377ade9b..87d5f312 100644 --- a/modelscope/metainfo.py +++ b/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' diff --git a/modelscope/models/audio/sv/ERes2Net.py b/modelscope/models/audio/sv/ERes2Net.py index 54f81ac9..0119783c 100644 --- a/modelscope/models/audio/sv/ERes2Net.py +++ b/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) diff --git a/modelscope/models/audio/sv/ERes2Net_aug.py b/modelscope/models/audio/sv/ERes2Net_aug.py index 08f1a8a1..d0739cad 100644 --- a/modelscope/models/audio/sv/ERes2Net_aug.py +++ b/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) diff --git a/modelscope/models/audio/sv/TDNN.py b/modelscope/models/audio/sv/TDNN.py new file mode 100644 index 00000000..9cc35c1f --- /dev/null +++ b/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 diff --git a/modelscope/models/audio/sv/speaker_change_locator_xvector.py b/modelscope/models/audio/sv/speaker_change_locator_xvector.py new file mode 100644 index 00000000..72e7b3d0 --- /dev/null +++ b/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'))) diff --git a/modelscope/pipelines/audio/segmentation_clustering_pipeline.py b/modelscope/pipelines/audio/segmentation_clustering_pipeline.py index 326d8787..e4810bcf 100644 --- a/modelscope/pipelines/audio/segmentation_clustering_pipeline.py +++ b/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 diff --git a/modelscope/pipelines/audio/speaker_verification_eres2net_pipeline.py b/modelscope/pipelines/audio/speaker_verification_eres2net_pipeline.py index ef91d83b..ba28ed6e 100644 --- a/modelscope/pipelines/audio/speaker_verification_eres2net_pipeline.py +++ b/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) diff --git a/modelscope/pipelines/audio/speaker_verification_light_pipeline.py b/modelscope/pipelines/audio/speaker_verification_light_pipeline.py index e3d1968a..7bfc7964 100644 --- a/modelscope/pipelines/audio/speaker_verification_light_pipeline.py +++ b/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 diff --git a/tests/pipelines/test_speaker_verification.py b/tests/pipelines/test_speaker_verification.py index 34e5a9d9..c5fe0004 100644 --- a/tests/pipelines/test_speaker_verification.py +++ b/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')