mirror of
https://github.com/modelscope/FunASR
synced 2025-09-15 14:48:36 +08:00
202 lines
7.1 KiB
Python
202 lines
7.1 KiB
Python
# -*- encoding: utf-8 -*-
|
|
#!/usr/bin/env python3
|
|
# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved.
|
|
# MIT License (https://opensource.org/licenses/MIT)
|
|
|
|
import argparse
|
|
import logging
|
|
import os
|
|
import sys
|
|
import json
|
|
from pathlib import Path
|
|
from typing import Any
|
|
from typing import List
|
|
from typing import Optional
|
|
from typing import Sequence
|
|
from typing import Tuple
|
|
from typing import Union
|
|
from typing import Dict
|
|
|
|
import math
|
|
import numpy as np
|
|
import torch
|
|
from typeguard import check_argument_types
|
|
from typeguard import check_return_type
|
|
|
|
from funasr.fileio.datadir_writer import DatadirWriter
|
|
from funasr.modules.scorers.scorer_interface import BatchScorerInterface
|
|
from funasr.modules.subsampling import TooShortUttError
|
|
from funasr.tasks.vad import VADTask
|
|
from funasr.torch_utils.device_funcs import to_device
|
|
from funasr.torch_utils.set_all_random_seed import set_all_random_seed
|
|
from funasr.utils import config_argparse
|
|
from funasr.utils.cli_utils import get_commandline_args
|
|
from funasr.utils.types import str2bool
|
|
from funasr.utils.types import str2triple_str
|
|
from funasr.utils.types import str_or_none
|
|
from funasr.utils import asr_utils, wav_utils, postprocess_utils
|
|
from funasr.models.frontend.wav_frontend import WavFrontend, WavFrontendOnline
|
|
|
|
|
|
|
|
class Speech2VadSegment:
|
|
"""Speech2VadSegment class
|
|
|
|
Examples:
|
|
>>> import soundfile
|
|
>>> speech2segment = Speech2VadSegment("vad_config.yml", "vad.pt")
|
|
>>> audio, rate = soundfile.read("speech.wav")
|
|
>>> speech2segment(audio)
|
|
[[10, 230], [245, 450], ...]
|
|
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
vad_infer_config: Union[Path, str] = None,
|
|
vad_model_file: Union[Path, str] = None,
|
|
vad_cmvn_file: Union[Path, str] = None,
|
|
device: str = "cpu",
|
|
batch_size: int = 1,
|
|
dtype: str = "float32",
|
|
**kwargs,
|
|
):
|
|
assert check_argument_types()
|
|
|
|
# 1. Build vad model
|
|
vad_model, vad_infer_args = VADTask.build_model_from_file(
|
|
vad_infer_config, vad_model_file, device
|
|
)
|
|
frontend = None
|
|
if vad_infer_args.frontend is not None:
|
|
frontend = WavFrontend(cmvn_file=vad_cmvn_file, **vad_infer_args.frontend_conf)
|
|
|
|
logging.info("vad_model: {}".format(vad_model))
|
|
logging.info("vad_infer_args: {}".format(vad_infer_args))
|
|
vad_model.to(dtype=getattr(torch, dtype)).eval()
|
|
|
|
self.vad_model = vad_model
|
|
self.vad_infer_args = vad_infer_args
|
|
self.device = device
|
|
self.dtype = dtype
|
|
self.frontend = frontend
|
|
self.batch_size = batch_size
|
|
|
|
@torch.no_grad()
|
|
def __call__(
|
|
self, speech: Union[torch.Tensor, np.ndarray], speech_lengths: Union[torch.Tensor, np.ndarray] = None,
|
|
in_cache: Dict[str, torch.Tensor] = dict()
|
|
) -> Tuple[List[List[int]], Dict[str, torch.Tensor]]:
|
|
"""Inference
|
|
|
|
Args:
|
|
speech: Input speech data
|
|
Returns:
|
|
text, token, token_int, hyp
|
|
|
|
"""
|
|
assert check_argument_types()
|
|
|
|
# Input as audio signal
|
|
if isinstance(speech, np.ndarray):
|
|
speech = torch.tensor(speech)
|
|
|
|
if self.frontend is not None:
|
|
self.frontend.filter_length_max = math.inf
|
|
fbanks, fbanks_len = self.frontend.forward_fbank(speech, speech_lengths)
|
|
feats, feats_len = self.frontend.forward_lfr_cmvn(fbanks, fbanks_len)
|
|
fbanks = to_device(fbanks, device=self.device)
|
|
feats = to_device(feats, device=self.device)
|
|
feats_len = feats_len.int()
|
|
else:
|
|
raise Exception("Need to extract feats first, please configure frontend configuration")
|
|
|
|
# b. Forward Encoder streaming
|
|
t_offset = 0
|
|
step = min(feats_len.max(), 6000)
|
|
segments = [[]] * self.batch_size
|
|
for t_offset in range(0, feats_len, min(step, feats_len - t_offset)):
|
|
if t_offset + step >= feats_len - 1:
|
|
step = feats_len - t_offset
|
|
is_final = True
|
|
else:
|
|
is_final = False
|
|
batch = {
|
|
"feats": feats[:, t_offset:t_offset + step, :],
|
|
"waveform": speech[:, t_offset * 160:min(speech.shape[-1], (t_offset + step - 1) * 160 + 400)],
|
|
"is_final": is_final,
|
|
"in_cache": in_cache
|
|
}
|
|
# a. To device
|
|
#batch = to_device(batch, device=self.device)
|
|
segments_part, in_cache = self.vad_model(**batch)
|
|
if segments_part:
|
|
for batch_num in range(0, self.batch_size):
|
|
segments[batch_num] += segments_part[batch_num]
|
|
return fbanks, segments
|
|
|
|
class Speech2VadSegmentOnline(Speech2VadSegment):
|
|
"""Speech2VadSegmentOnline class
|
|
|
|
Examples:
|
|
>>> import soundfile
|
|
>>> speech2segment = Speech2VadSegmentOnline("vad_config.yml", "vad.pt")
|
|
>>> audio, rate = soundfile.read("speech.wav")
|
|
>>> speech2segment(audio)
|
|
[[10, 230], [245, 450], ...]
|
|
|
|
"""
|
|
def __init__(self, **kwargs):
|
|
super(Speech2VadSegmentOnline, self).__init__(**kwargs)
|
|
vad_cmvn_file = kwargs.get('vad_cmvn_file', None)
|
|
self.frontend = None
|
|
if self.vad_infer_args.frontend is not None:
|
|
self.frontend = WavFrontendOnline(cmvn_file=vad_cmvn_file, **self.vad_infer_args.frontend_conf)
|
|
|
|
|
|
@torch.no_grad()
|
|
def __call__(
|
|
self, speech: Union[torch.Tensor, np.ndarray], speech_lengths: Union[torch.Tensor, np.ndarray] = None,
|
|
in_cache: Dict[str, torch.Tensor] = dict(), is_final: bool = False, max_end_sil: int = 800
|
|
) -> Tuple[torch.Tensor, List[List[int]], torch.Tensor]:
|
|
"""Inference
|
|
|
|
Args:
|
|
speech: Input speech data
|
|
Returns:
|
|
text, token, token_int, hyp
|
|
|
|
"""
|
|
assert check_argument_types()
|
|
|
|
# Input as audio signal
|
|
if isinstance(speech, np.ndarray):
|
|
speech = torch.tensor(speech)
|
|
batch_size = speech.shape[0]
|
|
segments = [[]] * batch_size
|
|
if self.frontend is not None:
|
|
feats, feats_len = self.frontend.forward(speech, speech_lengths, is_final)
|
|
fbanks, _ = self.frontend.get_fbank()
|
|
else:
|
|
raise Exception("Need to extract feats first, please configure frontend configuration")
|
|
if feats.shape[0]:
|
|
feats = to_device(feats, device=self.device)
|
|
feats_len = feats_len.int()
|
|
waveforms = self.frontend.get_waveforms()
|
|
|
|
batch = {
|
|
"feats": feats,
|
|
"waveform": waveforms,
|
|
"in_cache": in_cache,
|
|
"is_final": is_final,
|
|
"max_end_sil": max_end_sil
|
|
}
|
|
# a. To device
|
|
batch = to_device(batch, device=self.device)
|
|
segments, in_cache = self.vad_model.forward_online(**batch)
|
|
# in_cache.update(batch['in_cache'])
|
|
# in_cache = {key: value for key, value in batch['in_cache'].items()}
|
|
return fbanks, segments, in_cache
|
|
|
|
|