From 33d3d2084403fd34b79c835d2f2fe04f6cd8f738 Mon Sep 17 00:00:00 2001
From: 游雁 <zhifu.gzf@alibaba-inc.com>
Date: 星期三, 13 九月 2023 09:33:54 +0800
Subject: [PATCH] Merge branch 'main' of github.com:alibaba-damo-academy/FunASR add
---
funasr/bin/diar_inference_launch.py | 136 ++++++++++++++++++---------------------------
1 files changed, 55 insertions(+), 81 deletions(-)
diff --git a/funasr/bin/diar_inference_launch.py b/funasr/bin/diar_inference_launch.py
index 08004e8..b655df5 100755
--- a/funasr/bin/diar_inference_launch.py
+++ b/funasr/bin/diar_inference_launch.py
@@ -1,4 +1,5 @@
-#!/usr/bin/env python3
+# !/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)
@@ -7,48 +8,27 @@
import logging
import os
import sys
-from typing import Union, Dict, Any
-
-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
-
-import argparse
-import logging
-import os
-import sys
-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 collections import OrderedDict
import numpy as np
import soundfile
import torch
-from torch.nn import functional as F
-from typeguard import check_argument_types
-from typeguard import check_return_type
from scipy.signal import medfilt
-from funasr.utils.cli_utils import get_commandline_args
-from funasr.tasks.diar import DiarTask
-from funasr.tasks.asr import ASRTask
-from funasr.tasks.diar import EENDOLADiarTask
-from funasr.torch_utils.device_funcs import to_device
+
+from funasr.bin.diar_infer import Speech2DiarizationSOND, Speech2DiarizationEEND
+from funasr.datasets.iterable_dataset import load_bytes
+from funasr.build_utils.build_streaming_iterator import build_streaming_iterator
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 scipy.ndimage import median_filter
-from funasr.utils.misc import statistic_model_parameters
-from funasr.datasets.iterable_dataset import load_bytes
-from funasr.bin.diar_infer import Speech2DiarizationSOND, Speech2DiarizationEEND
+
def inference_sond(
diar_train_config: str,
@@ -71,7 +51,6 @@
mode: str = "sond",
**kwargs,
):
- assert check_argument_types()
ncpu = kwargs.get("ncpu", 1)
torch.set_num_threads(ncpu)
if batch_size > 1:
@@ -94,7 +73,8 @@
set_all_random_seed(seed)
# 2a. Build speech2xvec [Optional]
- if mode == "sond_demo" and param_dict is not None and "extract_profile" in param_dict and param_dict["extract_profile"]:
+ if mode == "sond_demo" and param_dict is not None and "extract_profile" in param_dict and param_dict[
+ "extract_profile"]:
assert "sv_train_config" in param_dict, "sv_train_config must be provided param_dict."
assert "sv_model_file" in param_dict, "sv_model_file must be provided in param_dict."
sv_train_config = param_dict["sv_train_config"]
@@ -112,10 +92,7 @@
embedding_node="resnet1_dense"
)
logging.info("speech2xvector_kwargs: {}".format(speech2xvector_kwargs))
- speech2xvector = Speech2Xvector.from_pretrained(
- model_tag=model_tag,
- **speech2xvector_kwargs,
- )
+ speech2xvector = Speech2Xvector(**speech2xvector_kwargs)
speech2xvector.sv_model.eval()
# 2b. Build speech2diar
@@ -129,17 +106,14 @@
dur_threshold=dur_threshold,
)
logging.info("speech2diarization_kwargs: {}".format(speech2diar_kwargs))
- speech2diar = Speech2DiarizationSOND.from_pretrained(
- model_tag=model_tag,
- **speech2diar_kwargs,
- )
+ speech2diar = Speech2DiarizationSOND(**speech2diar_kwargs)
speech2diar.diar_model.eval()
def output_results_str(results: dict, uttid: str):
rst = []
mid = uttid.rsplit("-", 1)[0]
for key in results:
- results[key] = [(x[0]/100, x[1]/100) for x in results[key]]
+ results[key] = [(x[0] / 100, x[1] / 100) for x in results[key]]
if out_format == "vad":
for spk, segs in results.items():
rst.append("{} {}".format(spk, segs))
@@ -176,7 +150,7 @@
example = [x.numpy() if isinstance(example[0], torch.Tensor) else x
for x in example]
speech = example[0]
- logging.info("Extracting profiles for {} waveforms".format(len(example)-1))
+ logging.info("Extracting profiles for {} waveforms".format(len(example) - 1))
profile = [speech2xvector.calculate_embedding(x) for x in example[1:]]
profile = torch.cat(profile, dim=0)
yield ["test{}".format(idx)], {"speech": [speech], "profile": [profile]}
@@ -186,16 +160,15 @@
raise TypeError("raw_inputs must be a list or tuple in [speech, profile1, profile2, ...] ")
else:
# 3. Build data-iterator
- loader = ASRTask.build_streaming_iterator(
- data_path_and_name_and_type,
+ loader = build_streaming_iterator(
+ task_name="diar",
+ preprocess_args=None,
+ data_path_and_name_and_type=data_path_and_name_and_type,
dtype=dtype,
batch_size=batch_size,
key_file=key_file,
num_workers=num_workers,
- preprocess_fn=None,
- collate_fn=None,
- allow_variable_data_keys=allow_variable_data_keys,
- inference=True,
+ use_collate_fn=False,
)
# 7. Start for-loop
@@ -235,6 +208,7 @@
return _forward
+
def inference_eend(
diar_train_config: str,
diar_model_file: str,
@@ -251,7 +225,6 @@
param_dict: Optional[dict] = None,
**kwargs,
):
- assert check_argument_types()
ncpu = kwargs.get("ncpu", 1)
torch.set_num_threads(ncpu)
if batch_size > 1:
@@ -278,10 +251,7 @@
dtype=dtype,
)
logging.info("speech2diarization_kwargs: {}".format(speech2diar_kwargs))
- speech2diar = Speech2DiarizationEEND.from_pretrained(
- model_tag=model_tag,
- **speech2diar_kwargs,
- )
+ speech2diar = Speech2DiarizationEEND(**speech2diar_kwargs)
speech2diar.diar_model.eval()
def output_results_str(results: dict, uttid: str):
@@ -306,16 +276,14 @@
if isinstance(raw_inputs, torch.Tensor):
raw_inputs = raw_inputs.numpy()
data_path_and_name_and_type = [raw_inputs[0], "speech", "sound"]
- loader = EENDOLADiarTask.build_streaming_iterator(
- data_path_and_name_and_type,
+ loader = build_streaming_iterator(
+ task_name="diar",
+ preprocess_args=None,
+ data_path_and_name_and_type=data_path_and_name_and_type,
dtype=dtype,
batch_size=batch_size,
key_file=key_file,
num_workers=num_workers,
- preprocess_fn=EENDOLADiarTask.build_preprocess_fn(speech2diar.diar_train_args, False),
- collate_fn=EENDOLADiarTask.build_collate_fn(speech2diar.diar_train_args, False),
- allow_variable_data_keys=allow_variable_data_keys,
- inference=True,
)
# 3. Start for-loop
@@ -360,6 +328,29 @@
return result_list
return _forward
+
+
+def inference_launch(mode, **kwargs):
+ if mode == "sond":
+ return inference_sond(mode=mode, **kwargs)
+ elif mode == "sond_demo":
+ param_dict = {
+ "extract_profile": True,
+ "sv_train_config": "sv.yaml",
+ "sv_model_file": "sv.pb",
+ }
+ if "param_dict" in kwargs and kwargs["param_dict"] is not None:
+ for key in param_dict:
+ if key not in kwargs["param_dict"]:
+ kwargs["param_dict"][key] = param_dict[key]
+ else:
+ kwargs["param_dict"] = param_dict
+ return inference_sond(mode=mode, **kwargs)
+ elif mode == "eend-ola":
+ return inference_eend(mode=mode, **kwargs)
+ else:
+ logging.info("Unknown decoding mode: {}".format(mode))
+ return None
def get_parser():
@@ -462,36 +453,19 @@
help="The batch size for inference",
)
group.add_argument(
- "--diar_smooth_size",
+ "--smooth_size",
type=int,
default=121,
help="The smoothing size for post-processing"
)
+ group.add_argument(
+ "--dur_threshold",
+ type=int,
+ default=10,
+ help="The threshold of minimum duration"
+ )
return parser
-
-
-def inference_launch(mode, **kwargs):
- if mode == "sond":
- return inference_sond(mode=mode, **kwargs)
- elif mode == "sond_demo":
- param_dict = {
- "extract_profile": True,
- "sv_train_config": "sv.yaml",
- "sv_model_file": "sv.pb",
- }
- if "param_dict" in kwargs and kwargs["param_dict"] is not None:
- for key in param_dict:
- if key not in kwargs["param_dict"]:
- kwargs["param_dict"][key] = param_dict[key]
- else:
- kwargs["param_dict"] = param_dict
- return inference_sond(mode=mode, **kwargs)
- elif mode == "eend-ola":
- return inference_eend(mode=mode, **kwargs)
- else:
- logging.info("Unknown decoding mode: {}".format(mode))
- return None
def main(cmd=None):
--
Gitblit v1.9.1