eres2net_lre_v2

Link: https://code.alibaba-inc.com/Ali-MaaS/MaaS-lib/codereview/13602081

* add eres2net_base_large_lre

* eres2net_language_identification

* eres2net_lre_v2

modelscope/metainfo.py

@@ -195,6 +195,7 @@ class Models(object):
     eres2net_aug_sv = 'eres2net-aug-sv'
     scl_sd = 'scl-sd'
     campplus_lre = 'cam++-lre'
+    eres2net_lre = 'eres2net-lre'
     cluster_backend = 'cluster-backend'
     rdino_tdnn_sv = 'rdino_ecapa-tdnn-sv'
     generic_lm = 'generic-lm'
@@ -506,6 +507,7 @@ class Pipelines(object):
     speaker_verification_rdino = 'speaker-verification-rdino'
     speaker_verification_eres2net = 'speaker-verification-eres2net'
     speech_language_recognition = 'speech-language-recognition'
+    speech_language_recognition_eres2net = 'speech-language-recognition-eres2net'
     speaker_change_locating = 'speaker-change-locating'
     speaker_diarization_dialogue_detection = 'speaker-diarization-dialogue-detection'
     speaker_diarization_semantic_speaker_turn_detection = 'speaker-diarization-semantic-speaker-turn-detection'

modelscope/models/audio/sv/lanuage_recognition_eres2net.py

@@ -0,0 +1,117 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+
+import os
+from typing import Any, Dict
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torchaudio.compliance.kaldi as Kaldi
+
+from modelscope.metainfo import Models
+from modelscope.models import MODELS, TorchModel
+from modelscope.models.audio.sv.DTDNN import CAMPPlus
+from modelscope.models.audio.sv.DTDNN_layers import DenseLayer
+from modelscope.models.audio.sv.ERes2Net import ERes2Net
+from modelscope.utils.constant import Tasks
+from modelscope.utils.device import create_device
+
+
+class LinearClassifier(nn.Module):
+
+    def __init__(
+        self,
+        input_dim,
+        num_blocks=0,
+        inter_dim=512,
+        out_neurons=1000,
+    ):
+
+        super().__init__()
+        self.blocks = nn.ModuleList()
+
+        self.nonlinear = nn.ReLU(inplace=True)
+        for _ in range(num_blocks):
+            self.blocks.append(DenseLayer(input_dim, inter_dim, bias=True))
+            input_dim = inter_dim
+
+        self.linear = nn.Linear(input_dim, out_neurons, bias=True)
+
+    def forward(self, x):
+        # x: [B, dim]
+        x = self.nonlinear(x)
+        for layer in self.blocks:
+            x = layer(x)
+        x = self.linear(x)
+        return x
+
+
+@MODELS.register_module(
+    Tasks.speech_language_recognition, module_name=Models.eres2net_lre)
+class LanguageRecognitionERes2Net(TorchModel):
+    r"""A speech language recognition model using the ERes2Net 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.embed_dim = self.model_config['embed_dim']
+        self.m_channels = self.model_config['channels']
+        self.feature_dim = self.model_config['fbank_dim']
+        self.device = create_device(kwargs['device'])
+
+        self.encoder = ERes2Net(
+            embed_dim=self.embed_dim, m_channels=self.m_channels)
+        self.backend = LinearClassifier(
+            input_dim=self.embed_dim,
+            out_neurons=len(self.model_config['languages']))
+
+        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):
+        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)
+        embs = self.encoder(feature.to(self.device))
+        output = self.backend(embs)
+        output = output.detach().cpu().argmax(-1)
+        return output
+
+    def _extract_feature(self, audio):
+        features = []
+        for au in audio:
+            feature = Kaldi.fbank(
+                au.unsqueeze(0), num_mel_bins=self.feature_dim)
+            feature = feature - feature.mean(dim=0, keepdim=True)
+            features.append(feature.unsqueeze(0))
+        features = torch.cat(features)
+        return features
+
+    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/language_recognition_eres2net_pipeline.py

@@ -0,0 +1,144 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+
+import io
+import os
+from typing import Union
+
+import numpy as np
+import soundfile as sf
+import torch
+import torchaudio
+
+from modelscope.fileio import File
+from modelscope.metainfo import Pipelines
+from modelscope.outputs import OutputKeys
+from modelscope.pipelines.base import InputModel, Pipeline
+from modelscope.pipelines.builder import PIPELINES
+from modelscope.utils.constant import Tasks
+from modelscope.utils.logger import get_logger
+
+logger = get_logger()
+
+__all__ = ['LanguageRecognitionPipeline']
+
+
+@PIPELINES.register_module(
+    Tasks.speech_language_recognition,
+    module_name=Pipelines.speech_language_recognition_eres2net)
+class LanguageRecognitionPipeline(Pipeline):
+    """Language Recognition Inference Pipeline
+    use `model` to create a Language Recognition pipeline.
+
+    Args:
+        model (LanguageRecognitionPipeline): A model instance, or a model local dir, or a model id in the model hub.
+        kwargs (dict, `optional`):
+            Extra kwargs passed into the pipeline's constructor.
+    Example:
+    >>> from modelscope.pipelines import pipeline
+    >>> from modelscope.utils.constant import Tasks
+    >>> p = pipeline(
+    >>>    task=Tasks.speech_language_recognition, model='damo/speech_eres2net_base_lre_en-cn_16k')
+    >>> print(p(audio_in))
+
+    """
+
+    def __init__(self, model: InputModel, **kwargs):
+        """use `model` to create a Language Recognition pipeline for prediction
+        Args:
+            model (str): a valid offical model id
+        """
+        super().__init__(model=model, **kwargs)
+        self.model_config = self.model.model_config
+        self.languages = self.model_config['languages']
+
+    def __call__(self,
+                 in_audios: Union[str, list, np.ndarray],
+                 out_file: str = None):
+        wavs = self.preprocess(in_audios)
+        results = self.forward(wavs)
+        outputs = self.postprocess(results, in_audios, out_file)
+        return outputs
+
+    def forward(self, inputs: list):
+        results = []
+        for x in inputs:
+            results.append(self.model(x).item())
+        return results
+
+    def postprocess(self,
+                    inputs: list,
+                    in_audios: Union[str, list, np.ndarray],
+                    out_file=None):
+        if isinstance(in_audios, str):
+            output = {OutputKeys.TEXT: self.languages[inputs[0]]}
+        else:
+            output = {OutputKeys.TEXT: [self.languages[i] for i in inputs]}
+            if out_file is not None:
+                out_lines = []
+                for i, audio in enumerate(in_audios):
+                    if isinstance(audio, str):
+                        audio_id = os.path.basename(audio).rsplit('.', 1)[0]
+                    else:
+                        audio_id = i
+                    out_lines.append('%s %s\n' %
+                                     (audio_id, self.languages[inputs[i]]))
+                with open(out_file, 'w') as f:
+                    for i in out_lines:
+                        f.write(i)
+        return output
+
+    def preprocess(self, inputs: Union[str, list, np.ndarray]):
+        output = []
+        if isinstance(inputs, str):
+            file_bytes = File.read(inputs)
+            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']:
+                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)
+            output.append(data)
+        else:
+            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']:
+                        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 restricted to audio address and nump array.'
+                    )
+                output.append(data)
+        return output

tests/pipelines/test_speaker_verification.py

@@ -26,6 +26,8 @@ class SpeakerVerificationTest(unittest.TestCase):
     eres2net_voxceleb_16k_model_id = 'damo/speech_eres2net_sv_en_voxceleb_16k'
     speaker_diarization_model_id = 'damo/speech_campplus_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'
     eres2net_aug_zh_cn_16k_common_model_id = 'damo/speech_eres2net_sv_zh-cn_16k-common'
     rdino_3dspeaker_16k_model_id = 'damo/speech_rdino_ecapa_tdnn_sv_zh-cn_3dspeaker_16k'
     eres2net_base_3dspeaker_16k_model_id = 'damo/speech_eres2net_base_sv_zh-cn_3dspeaker_16k'
@@ -161,6 +163,28 @@ 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_language_recognition_eres2net_base_en_cn_16k(self):
+        logger.info('Run language recognition for eres2net_base_en_cn_16k')
+        result = self.run_pipeline(
+            model_id=self.lre_eres2net_base_en_cn_16k_model_id,
+            task=Tasks.speech_language_recognition,
+            audios=SPEAKER1_A_EN_16K_WAV,
+            model_revision='v1.0.2')
+        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_eres2net_large_en_cn_16k(self):
+        logger.info('Run language recognition for eres2net_large_en_cn_16k')
+        result = self.run_pipeline(
+            model_id=self.lre_eres2net_large_en_cn_16k_model_id,
+            task=Tasks.speech_language_recognition,
+            audios=SPEAKER1_A_EN_16K_WAV,
+            model_revision='v1.0.0')
+        print(result)
+        self.assertTrue(OutputKeys.TEXT in result)
+
 
 if __name__ == '__main__':
     unittest.main()