update

Dockerfile

@@ -10,10 +10,6 @@ RUN apt-get install -y ffmpeg build-essential
 RUN pip install --upgrade pip
 RUN pip install --no-cache-dir --upgrade -r /code/requirements.txt
 
-RUN pip install --upgrade pip
-
-RUN bash install.sh --stage 1 --stop_stage 2 --system_version centos
-
 USER user
 
 ENV HOME=/home/user \
@@ -23,4 +19,6 @@ WORKDIR $HOME/app
 
 COPY --chown=user . $HOME/app
 
+RUN bash install.sh --stage 1 --stop_stage 2 --system_version centos
+
 CMD ["python3", "main.py"]

examples/data_annotation/annotation_by_google.py

@@ -0,0 +1,158 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import argparse
+import json
+import os
+import re
+from pathlib import Path
+import sys
+
+pwd = os.path.abspath(os.path.dirname(__file__))
+sys.path.append(os.path.join(pwd, "../../"))
+
+from google import genai
+from google.genai import types
+
+from project_settings import environment, project_path
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--google_application_credentials",
+        default=(project_path / "dotenv/potent-veld-462405-t3-8091a29b2894.json").as_posix(),
+        type=str
+    )
+    parser.add_argument(
+        "--model_name",
+        default="gemini-2.5-pro",
+        type=str
+    )
+    parser.add_argument(
+        "--speech_audio_dir",
+        default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\speech\nx-speech\en-SG\2025-06-17",
+        type=str
+    )
+    parser.add_argument(
+        "--output_file",
+        # default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\noise\nx-noise\en-SG\2025-06-17\vad.jsonl",
+        default=r"vad.jsonl",
+        type=str
+    )
+    parser.add_argument(
+        "--gemini_api_key",
+        default=environment.get("GEMINI_API_KEY", dtype=str),
+        type=str
+    )
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = get_args()
+
+    speech_audio_dir = Path(args.speech_audio_dir)
+    output_file = Path(args.output_file)
+
+    os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = args.google_application_credentials
+    os.environ["gemini_api_key"] = args.gemini_api_key
+
+
+    developer_client = genai.Client(
+        api_key=args.gemini_api_key,
+    )
+    client = genai.Client(
+        vertexai=True,
+        project="potent-veld-462405-t3",
+        location="global",
+    )
+    generate_content_config = types.GenerateContentConfig(
+        temperature=1,
+        top_p=0.95,
+        max_output_tokens=8192,
+        response_modalities=["TEXT"],
+    )
+
+    # finished
+    finished_set = set()
+    if output_file.exists():
+        with open(output_file.as_posix(), "r", encoding="utf-8") as f:
+            for row in f:
+                row = json.loads(row)
+                name = row["name"]
+                finished_set.add(name)
+    print(f"finished count: {len(finished_set)}")
+
+    with open(output_file.as_posix(), "a+", encoding="utf-8") as f:
+
+        for filename in speech_audio_dir.glob("**/*.wav"):
+            name = filename.name
+            if name in finished_set:
+                continue
+            finished_set.add(name)
+
+            # upload
+            audio_file = developer_client.files.upload(
+                file=filename.as_posix(),
+                config=None
+            )
+            print(f"upload file: {audio_file.name}")
+
+            prompt = f"""
+    给我这段音频中的语音分段的开始和结束时间，单位为秒，精确到毫秒，并输出JSON格式，
+    例如：
+    ```json
+    [[0.254, 1.214], [2.200, 3.100]]，
+    ```
+    如果没有语音段则输出：
+    ```json
+    []
+    ```
+    """.strip()
+
+            try:
+                contents = [
+                    types.Content(
+                        role="user",
+                        parts=[
+                            types.Part(text=prompt),
+                            types.Part.from_uri(
+                                file_uri=audio_file.uri,
+                                mime_type=audio_file.mime_type,
+                            )
+                        ]
+                    )
+                ]
+                response: types.GenerateContentResponse = developer_client.models.generate_content(
+                    model=args.model_name,
+                    contents=contents,
+                    config=generate_content_config,
+                )
+                answer = response.candidates[0].content.parts[0].text
+                print(answer)
+            finally:
+                # delete
+                print(f"delete file: {audio_file.name}")
+                developer_client.files.delete(name=audio_file.name)
+
+            pattern = "```json(.+?)```"
+            match = re.search(pattern=pattern, string=answer, flags=re.DOTALL | re.IGNORECASE)
+            if match is None:
+                raise AssertionError(f"answer: {answer}")
+            vad_segments = match.group(1)
+            vad_segments = json.loads(vad_segments)
+            row = {
+                "name": name,
+                "filename": filename.as_posix(),
+                "vad_segments": vad_segments
+            }
+            row = json.dumps(row, ensure_ascii=False)
+
+            f.write(f"{row}\n")
+            exit(0)
+
+    return
+
+
+if __name__ == "__main__":
+    main()

examples/fsmn_vad_by_webrtcvad/run.sh

@@ -0,0 +1,174 @@
+#!/usr/bin/env bash
+
+: <<'END'
+
+bash run.sh --stage 2 --stop_stage 2 --system_version centos \
+--file_folder_name silero-vad-by-webrtcvad-nx2-dns3 \
+--final_model_name silero-vad-by-webrtcvad-nx2-dns3 \
+--noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise/" \
+--speech_dir "/data/tianxing/HuggingDatasets/nx_noise/data/speech/"
+
+bash run.sh --stage 3 --stop_stage 3 --system_version centos \
+--file_folder_name silero-vad-by-webrtcvad-nx2-dns3 \
+--final_model_name silero-vad-by-webrtcvad-nx2-dns3 \
+--noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise/" \
+--speech_dir "/data/tianxing/HuggingDatasets/nx_noise/data/speech/"
+
+
+END
+
+
+# params
+system_version="windows";
+verbose=true;
+stage=0 # start from 0 if you need to start from data preparation
+stop_stage=9
+
+work_dir="$(pwd)"
+file_folder_name=file_folder_name
+final_model_name=final_model_name
+config_file="yaml/config.yaml"
+limit=10
+
+noise_dir=/data/tianxing/HuggingDatasets/nx_noise/data/noise
+speech_dir=/data/tianxing/HuggingDatasets/aishell/data_aishell/wav/train
+
+max_count=-1
+
+nohup_name=nohup.out
+
+# model params
+batch_size=64
+max_epochs=200
+save_top_k=10
+patience=5
+
+
+# parse options
+while true; do
+  [ -z "${1:-}" ] && break;  # break if there are no arguments
+  case "$1" in
+    --*) name=$(echo "$1" | sed s/^--// | sed s/-/_/g);
+      eval '[ -z "${'"$name"'+xxx}" ]' && echo "$0: invalid option $1" 1>&2 && exit 1;
+      old_value="(eval echo \\$$name)";
+      if [ "${old_value}" == "true" ] || [ "${old_value}" == "false" ]; then
+        was_bool=true;
+      else
+        was_bool=false;
+      fi
+
+      # Set the variable to the right value-- the escaped quotes make it work if
+      # the option had spaces, like --cmd "queue.pl -sync y"
+      eval "${name}=\"$2\"";
+
+      # Check that Boolean-valued arguments are really Boolean.
+      if $was_bool && [[ "$2" != "true" && "$2" != "false" ]]; then
+        echo "$0: expected \"true\" or \"false\": $1 $2" 1>&2
+        exit 1;
+      fi
+      shift 2;
+      ;;
+
+    *) break;
+  esac
+done
+
+file_dir="${work_dir}/${file_folder_name}"
+final_model_dir="${work_dir}/../../trained_models/${final_model_name}";
+evaluation_audio_dir="${file_dir}/evaluation_audio"
+
+train_dataset="${file_dir}/train.jsonl"
+valid_dataset="${file_dir}/valid.jsonl"
+
+train_vad_dataset="${file_dir}/train-vad.jsonl"
+valid_vad_dataset="${file_dir}/valid-vad.jsonl"
+
+$verbose && echo "system_version: ${system_version}"
+$verbose && echo "file_folder_name: ${file_folder_name}"
+
+if [ $system_version == "windows" ]; then
+  alias python3='D:/Users/tianx/PycharmProjects/virtualenv/nx_denoise/Scripts/python.exe'
+elif [ $system_version == "centos" ] || [ $system_version == "ubuntu" ]; then
+  #source /data/local/bin/nx_denoise/bin/activate
+  alias python3='/data/local/bin/nx_denoise/bin/python3'
+fi
+
+
+if [ ${stage} -le 1 ] && [ ${stop_stage} -ge 1 ]; then
+  $verbose && echo "stage 1: prepare data"
+  cd "${work_dir}" || exit 1
+  python3 step_1_prepare_data.py \
+  --noise_dir "${noise_dir}" \
+  --speech_dir "${speech_dir}" \
+  --train_dataset "${train_dataset}" \
+  --valid_dataset "${valid_dataset}" \
+  --max_count "${max_count}" \
+
+fi
+
+
+if [ ${stage} -le 2 ] && [ ${stop_stage} -ge 2 ]; then
+  $verbose && echo "stage 2: make vad segments"
+  cd "${work_dir}" || exit 1
+  python3 step_2_make_vad_segments.py \
+  --train_dataset "${train_dataset}" \
+  --valid_dataset "${valid_dataset}" \
+  --train_vad_dataset "${train_vad_dataset}" \
+  --valid_vad_dataset "${valid_vad_dataset}" \
+
+fi
+
+
+if [ ${stage} -le 3 ] && [ ${stop_stage} -ge 3 ]; then
+  $verbose && echo "stage 3: train model"
+  cd "${work_dir}" || exit 1
+  python3 step_4_train_model.py \
+  --train_dataset "${train_vad_dataset}" \
+  --valid_dataset "${valid_vad_dataset}" \
+  --serialization_dir "${file_dir}" \
+  --config_file "${config_file}" \
+
+fi
+
+
+if [ ${stage} -le 4 ] && [ ${stop_stage} -ge 4 ]; then
+  $verbose && echo "stage 4: test model"
+  cd "${work_dir}" || exit 1
+  python3 step_3_evaluation.py \
+  --valid_dataset "${valid_dataset}" \
+  --model_dir "${file_dir}/best" \
+  --evaluation_audio_dir "${evaluation_audio_dir}" \
+  --limit "${limit}" \
+
+fi
+
+
+if [ ${stage} -le 5 ] && [ ${stop_stage} -ge 5 ]; then
+  $verbose && echo "stage 5: collect files"
+  cd "${work_dir}" || exit 1
+
+  mkdir -p ${final_model_dir}
+
+  cp "${file_dir}/best"/* "${final_model_dir}"
+  cp -r "${file_dir}/evaluation_audio" "${final_model_dir}"
+
+  cd "${final_model_dir}/.." || exit 1;
+
+  if [ -e "${final_model_name}.zip" ]; then
+    rm -rf "${final_model_name}_backup.zip"
+    mv "${final_model_name}.zip" "${final_model_name}_backup.zip"
+  fi
+
+  zip -r "${final_model_name}.zip" "${final_model_name}"
+  rm -rf "${final_model_name}"
+
+fi
+
+
+if [ ${stage} -le 6 ] && [ ${stop_stage} -ge 6 ]; then
+  $verbose && echo "stage 6: clear file_dir"
+  cd "${work_dir}" || exit 1
+
+  rm -rf "${file_dir}";
+
+fi

examples/fsmn_vad_by_webrtcvad/step_1_prepare_data.py

@@ -0,0 +1,231 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import argparse
+import json
+import os
+from pathlib import Path
+import random
+import sys
+import time
+
+pwd = os.path.abspath(os.path.dirname(__file__))
+sys.path.append(os.path.join(pwd, "../../"))
+
+import librosa
+import numpy as np
+from tqdm import tqdm
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--noise_dir",
+        default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\noise",
+        type=str
+    )
+    parser.add_argument(
+        "--speech_dir",
+        default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\speech",
+        type=str
+    )
+
+    parser.add_argument("--train_dataset", default="train.jsonl", type=str)
+    parser.add_argument("--valid_dataset", default="valid.jsonl", type=str)
+
+    parser.add_argument("--duration", default=8.0, type=float)
+    parser.add_argument("--min_speech_duration", default=6.0, type=float)
+    parser.add_argument("--max_speech_duration", default=8.0, type=float)
+    parser.add_argument("--min_snr_db", default=-10, type=float)
+    parser.add_argument("--max_snr_db", default=20, type=float)
+
+    parser.add_argument("--target_sample_rate", default=8000, type=int)
+
+    parser.add_argument("--max_count", default=-1, type=int)
+
+    args = parser.parse_args()
+    return args
+
+
+def target_second_noise_signal_generator(data_dir: str,
+                                         duration: int = 4,
+                                         sample_rate: int = 8000, max_epoch: int = 20000):
+    noise_list = list()
+    wait_duration = duration
+
+    data_dir = Path(data_dir)
+    for epoch_idx in range(max_epoch):
+        for filename in data_dir.glob("**/*.wav"):
+            signal, _ = librosa.load(filename.as_posix(), sr=sample_rate)
+
+            if signal.ndim != 1:
+                raise AssertionError(f"expected ndim 1, instead of {signal.ndim}")
+
+            raw_duration = librosa.get_duration(y=signal, sr=sample_rate)
+
+            offset = 0.
+            rest_duration = raw_duration
+
+            for _ in range(1000):
+                if rest_duration <= 0:
+                    break
+                if rest_duration <= wait_duration:
+                    noise_list.append({
+                        "epoch_idx": epoch_idx,
+                        "filename": filename.as_posix(),
+                        "raw_duration": round(raw_duration, 4),
+                        "offset": round(offset, 4),
+                        "duration": None,
+                        "duration_": round(rest_duration, 4),
+                    })
+                    wait_duration -= rest_duration
+                    offset = 0
+                    rest_duration = 0
+                elif rest_duration > wait_duration:
+                    noise_list.append({
+                        "epoch_idx": epoch_idx,
+                        "filename": filename.as_posix(),
+                        "raw_duration": round(raw_duration, 4),
+                        "offset": round(offset, 4),
+                        "duration": round(wait_duration, 4),
+                        "duration_": round(wait_duration, 4),
+                    })
+                    offset += wait_duration
+                    rest_duration -= wait_duration
+                    wait_duration = 0
+                else:
+                    raise AssertionError
+
+                if wait_duration <= 0:
+                    yield noise_list
+                    noise_list = list()
+                    wait_duration = duration
+
+
+def target_second_speech_signal_generator(data_dir: str,
+                                          min_duration: int = 4,
+                                          max_duration: int = 6,
+                                          sample_rate: int = 8000, max_epoch: int = 1):
+    data_dir = Path(data_dir)
+    for epoch_idx in range(max_epoch):
+        for filename in data_dir.glob("**/*.wav"):
+            signal, _ = librosa.load(filename.as_posix(), sr=sample_rate)
+            raw_duration = librosa.get_duration(y=signal, sr=sample_rate)
+
+            if signal.ndim != 1:
+                raise AssertionError(f"expected ndim 1, instead of {signal.ndim}")
+
+            if raw_duration < min_duration:
+                # print(f"duration less than {duration} s. skip filename: {filename.as_posix()}")
+                continue
+
+            if raw_duration < max_duration:
+                row = {
+                    "epoch_idx": epoch_idx,
+                    "filename": filename.as_posix(),
+                    "raw_duration": round(raw_duration, 4),
+                    "offset": 0.,
+                    "duration": round(raw_duration, 4),
+                }
+                yield row
+
+            signal_length = len(signal)
+            win_size = int(max_duration * sample_rate)
+            for begin in range(0, signal_length - win_size, win_size):
+                if np.sum(signal[begin: begin+win_size]) == 0:
+                    continue
+                row = {
+                    "epoch_idx": epoch_idx,
+                    "filename": filename.as_posix(),
+                    "raw_duration": round(raw_duration, 4),
+                    "offset": round(begin / sample_rate, 4),
+                    "duration": round(max_duration, 4),
+                }
+                yield row
+
+
+def main():
+    args = get_args()
+
+    noise_dir = Path(args.noise_dir)
+    speech_dir = Path(args.speech_dir)
+
+    train_dataset = Path(args.train_dataset)
+    valid_dataset = Path(args.valid_dataset)
+    train_dataset.parent.mkdir(parents=True, exist_ok=True)
+    valid_dataset.parent.mkdir(parents=True, exist_ok=True)
+
+    noise_generator = target_second_noise_signal_generator(
+        noise_dir.as_posix(),
+        duration=args.duration,
+        sample_rate=args.target_sample_rate,
+        max_epoch=100000,
+    )
+    speech_generator = target_second_speech_signal_generator(
+        speech_dir.as_posix(),
+        min_duration=args.min_speech_duration,
+        max_duration=args.max_speech_duration,
+        sample_rate=args.target_sample_rate,
+        max_epoch=1,
+    )
+
+    count = 0
+    process_bar = tqdm(desc="build dataset jsonl")
+    with open(args.train_dataset, "w", encoding="utf-8") as ftrain, open(args.valid_dataset, "w", encoding="utf-8") as fvalid:
+        for speech, noise_list in zip(speech_generator, noise_generator):
+            if count >= args.max_count > 0:
+                break
+
+            # row
+            speech_filename = speech["filename"]
+            speech_raw_duration = speech["raw_duration"]
+            speech_offset = speech["offset"]
+            speech_duration = speech["duration"]
+
+            noise_list = [
+                {
+                    "filename": noise["filename"],
+                    "raw_duration": noise["raw_duration"],
+                    "offset": noise["offset"],
+                    "duration": noise["duration"],
+                }
+                for noise in noise_list
+            ]
+
+            # row
+            random1 = random.random()
+            random2 = random.random()
+
+            row = {
+                "count": count,
+
+                "speech_filename": speech_filename,
+                "speech_raw_duration": speech_raw_duration,
+                "speech_offset": speech_offset,
+                "speech_duration": speech_duration,
+
+                "noise_list": noise_list,
+
+                "snr_db": random.uniform(args.min_snr_db, args.max_snr_db),
+
+                "random1": random1,
+            }
+            row = json.dumps(row, ensure_ascii=False)
+            if random2 < (1 / 300 / 1):
+                fvalid.write(f"{row}\n")
+            else:
+                ftrain.write(f"{row}\n")
+
+            count += 1
+            duration_seconds = count * args.duration
+            duration_hours = duration_seconds / 3600
+
+            process_bar.update(n=1)
+            process_bar.set_postfix({
+                "duration_hours": round(duration_hours, 4),
+            })
+
+    return
+
+
+if __name__ == "__main__":
+    main()

examples/fsmn_vad_by_webrtcvad/step_2_make_vad_segments.py

@@ -0,0 +1,205 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import argparse
+import json
+import os
+import sys
+
+pwd = os.path.abspath(os.path.dirname(__file__))
+sys.path.append(os.path.join(pwd, "../../"))
+
+import librosa
+import numpy as np
+from tqdm import tqdm
+
+from project_settings import project_path
+from toolbox.vad.vad import WebRTCVoiceClassifier, SileroVoiceClassifier, CCSoundsClassifier, RingVad
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("--train_dataset", default="train.jsonl", type=str)
+    parser.add_argument("--valid_dataset", default="valid.jsonl", type=str)
+
+    parser.add_argument("--train_vad_dataset", default="train-vad.jsonl", type=str)
+    parser.add_argument("--valid_vad_dataset", default="valid-vad.jsonl", type=str)
+
+    parser.add_argument("--duration", default=8.0, type=float)
+    parser.add_argument("--expected_sample_rate", default=8000, type=int)
+
+    parser.add_argument(
+        "--silero_model_path",
+        default=(project_path / "trained_models/silero_vad.jit").as_posix(),
+        type=str,
+    )
+    parser.add_argument(
+        "--cc_sounds_model_path",
+        default=(project_path / "trained_models/sound-2-ch32.zip").as_posix(),
+        type=str,
+    )
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = get_args()
+
+    # webrtcvad
+    # model = SileroVoiceClassifier(model_path=args.silero_model_path, sample_rate=args.expected_sample_rate)
+    # w_vad = RingVad(
+    #     model=model,
+    #     start_ring_rate=0.2,
+    #     end_ring_rate=0.1,
+    #     frame_size_ms=32,
+    #     frame_step_ms=32,
+    #     padding_length_ms=320,
+    #     max_silence_length_ms=320,
+    #     max_speech_length_s=100,
+    #     min_speech_length_s=0.1,
+    #     sample_rate=args.expected_sample_rate,
+    # )
+
+    # webrtcvad
+    model = WebRTCVoiceClassifier(agg=3, sample_rate=args.expected_sample_rate)
+    w_vad = RingVad(
+        model=model,
+        start_ring_rate=0.9,
+        end_ring_rate=0.1,
+        frame_size_ms=30,
+        frame_step_ms=30,
+        padding_length_ms=90,
+        max_silence_length_ms=100,
+        max_speech_length_s=100,
+        min_speech_length_s=0.1,
+        sample_rate=args.expected_sample_rate,
+    )
+
+    # cc sounds
+    # model = CCSoundsClassifier(model_path=args.cc_sounds_model_path, sample_rate=args.expected_sample_rate)
+    # w_vad = RingVad(
+    #     model=model,
+    #     start_ring_rate=0.5,
+    #     end_ring_rate=0.3,
+    #     frame_size_ms=300,
+    #     frame_step_ms=300,
+    #     padding_length_ms=300,
+    #     max_silence_length_ms=100,
+    #     max_speech_length_s=100,
+    #     min_speech_length_s=0.1,
+    #     sample_rate=args.expected_sample_rate,
+    # )
+
+    # valid
+    va_duration = 0
+    raw_duration = 0
+    use_duration = 0
+
+    count = 0
+    process_bar_valid = tqdm(desc="process valid dataset jsonl")
+    with (open(args.valid_dataset, "r", encoding="utf-8") as fvalid,
+          open(args.valid_vad_dataset, "w", encoding="utf-8") as fvalid_vad):
+        for row in fvalid:
+            row = json.loads(row)
+
+            speech_filename = row["speech_filename"]
+            speech_offset = row["speech_offset"]
+            speech_duration = row["speech_duration"]
+
+            waveform, _ = librosa.load(
+                speech_filename,
+                sr=args.expected_sample_rate,
+                offset=speech_offset,
+                duration=speech_duration,
+            )
+            waveform = np.array(waveform * (1 << 15), dtype=np.int16)
+
+            # vad
+            vad_segments = list()
+            segments = w_vad.vad(waveform)
+            vad_segments += segments
+            segments = w_vad.last_vad_segments()
+            vad_segments += segments
+            w_vad.reset()
+
+            row["vad_segments"] = vad_segments
+
+            row = json.dumps(row, ensure_ascii=False)
+            fvalid_vad.write(f"{row}\n")
+
+            va_duration += sum([vad_segment[1] - vad_segment[0] for vad_segment in vad_segments])
+            raw_duration += speech_duration
+            use_duration += args.duration
+
+            count += 1
+
+            va_rate = va_duration / use_duration
+            va_raw_rate = va_duration / raw_duration
+            use_duration_hours = use_duration / 3600
+
+            process_bar_valid.update(n=1)
+            process_bar_valid.set_postfix({
+                "va_rate": round(va_rate, 4),
+                "va_raw_rate": round(va_raw_rate, 4),
+                "duration_hours": round(use_duration_hours, 4),
+            })
+
+    # train
+    va_duration = 0
+    raw_duration = 0
+    use_duration = 0
+
+    count = 0
+    process_bar_train = tqdm(desc="process train dataset jsonl")
+    with (open(args.train_dataset, "r", encoding="utf-8") as ftrain,
+          open(args.train_vad_dataset, "w", encoding="utf-8") as ftrain_vad):
+        for row in ftrain:
+            row = json.loads(row)
+
+            speech_filename = row["speech_filename"]
+            speech_offset = row["speech_offset"]
+            speech_duration = row["speech_duration"]
+
+            waveform, _ = librosa.load(
+                speech_filename,
+                sr=args.expected_sample_rate,
+                offset=speech_offset,
+                duration=speech_duration,
+            )
+            waveform = np.array(waveform * (1 << 15), dtype=np.int16)
+
+            # vad
+            vad_segments = list()
+            segments = w_vad.vad(waveform)
+            vad_segments += segments
+            segments = w_vad.last_vad_segments()
+            vad_segments += segments
+            w_vad.reset()
+
+            row["vad_segments"] = vad_segments
+
+            row = json.dumps(row, ensure_ascii=False)
+            ftrain_vad.write(f"{row}\n")
+
+            va_duration += sum([vad_segment[1] - vad_segment[0] for vad_segment in vad_segments])
+            raw_duration += speech_duration
+            use_duration += args.duration
+
+            count += 1
+
+            va_rate = va_duration / use_duration
+            va_raw_rate = va_duration / raw_duration
+            use_duration_hours = use_duration / 3600
+
+            process_bar_train.update(n=1)
+            process_bar_train.set_postfix({
+                "va_rate": round(va_rate, 4),
+                "va_raw_rate": round(va_raw_rate, 4),
+                "duration_hours": round(use_duration_hours, 4),
+            })
+
+    return
+
+
+if __name__ == "__main__":
+    main()

examples/fsmn_vad_by_webrtcvad/step_3_check_vad.py

@@ -0,0 +1,68 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import argparse
+import json
+import os
+import sys
+
+pwd = os.path.abspath(os.path.dirname(__file__))
+sys.path.append(os.path.join(pwd, "../../"))
+
+import librosa
+import matplotlib.pyplot as plt
+import numpy as np
+from scipy.io import wavfile
+from tqdm import tqdm
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("--train_vad_dataset", default="train-vad.jsonl", type=str)
+    parser.add_argument("--valid_vad_dataset", default="valid-vad.jsonl", type=str)
+
+    parser.add_argument("--duration", default=8.0, type=float)
+    parser.add_argument("--expected_sample_rate", default=8000, type=int)
+
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = get_args()
+
+    SAMPLE_RATE = 8000
+
+    with open(args.train_vad_dataset, "r", encoding="utf-8") as f:
+        for row in f:
+            row = json.loads(row)
+
+            speech_filename = row["speech_filename"]
+            speech_offset = row["speech_offset"]
+            speech_duration = row["speech_duration"]
+
+            vad_segments = row["vad_segments"]
+
+            print(f"speech_filename: {speech_filename}")
+            signal, sample_rate = librosa.load(
+                speech_filename,
+                sr=SAMPLE_RATE,
+                offset=speech_offset,
+                duration=speech_duration,
+            )
+
+            # plot
+            time = np.arange(0, len(signal)) / sample_rate
+            plt.figure(figsize=(12, 5))
+            plt.plot(time, signal, color='b')
+            for start, end in vad_segments:
+                plt.axvline(x=start, ymin=0.25, ymax=0.75, color='g', linestyle='--', label='开始端点')  # 标记开始端点
+                plt.axvline(x=end, ymin=0.25, ymax=0.75, color='r', linestyle='--', label='结束端点')  # 标记结束端点
+
+            plt.show()
+
+    return
+
+
+if __name__ == "__main__":
+    main()

examples/fsmn_vad_by_webrtcvad/step_4_train_model.py

@@ -0,0 +1,453 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import argparse
+import json
+import logging
+from logging.handlers import TimedRotatingFileHandler
+import os
+import platform
+from pathlib import Path
+import random
+import sys
+import shutil
+from typing import List, Tuple
+
+pwd = os.path.abspath(os.path.dirname(__file__))
+sys.path.append(os.path.join(pwd, "../../"))
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+from torch.utils.data.dataloader import DataLoader
+from tqdm import tqdm
+
+from toolbox.torch.utils.data.dataset.vad_padding_jsonl_dataset import VadPaddingJsonlDataset
+from toolbox.torchaudio.models.vad.fsmn_vad.configuration_fsmn_vad import FSMNVadConfig
+from toolbox.torchaudio.models.vad.fsmn_vad.modeling_fsmn_vad import FSMNVadModel, FSMNVadPretrainedModel
+from toolbox.torchaudio.losses.vad_loss.base_vad_loss import BaseVadLoss
+from toolbox.torchaudio.losses.bce_loss import BCELoss
+from toolbox.torchaudio.losses.dice_loss import DiceLoss
+from toolbox.torchaudio.metrics.vad_metrics.vad_accuracy import VadAccuracy
+from toolbox.torchaudio.metrics.vad_metrics.vad_f1_score import VadF1Score
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--train_dataset", default="train-vad.jsonl", type=str)
+    parser.add_argument("--valid_dataset", default="valid-vad.jsonl", type=str)
+
+    parser.add_argument("--num_serialized_models_to_keep", default=15, type=int)
+    parser.add_argument("--patience", default=30, type=int)
+    parser.add_argument("--serialization_dir", default="serialization_dir", type=str)
+
+    parser.add_argument("--config_file", default="yaml/config.yaml", type=str)
+
+    args = parser.parse_args()
+    return args
+
+
+def logging_config(file_dir: str):
+    fmt = "%(asctime)s - %(name)s - %(levelname)s  %(filename)s:%(lineno)d >  %(message)s"
+
+    logging.basicConfig(format=fmt,
+                        datefmt="%m/%d/%Y %H:%M:%S",
+                        level=logging.INFO)
+    file_handler = TimedRotatingFileHandler(
+        filename=os.path.join(file_dir, "main.log"),
+        encoding="utf-8",
+        when="D",
+        interval=1,
+        backupCount=7
+    )
+    file_handler.setLevel(logging.INFO)
+    file_handler.setFormatter(logging.Formatter(fmt))
+    logger = logging.getLogger(__name__)
+    logger.addHandler(file_handler)
+
+    return logger
+
+
+class CollateFunction(object):
+    def __init__(self):
+        pass
+
+    def __call__(self, batch: List[dict]):
+        noisy_audios = list()
+        batch_vad_segments = list()
+
+        for sample in batch:
+            noisy_wave: torch.Tensor = sample["noisy_wave"]
+            vad_segments: List[Tuple[float, float]] = sample["vad_segments"]
+
+            noisy_audios.append(noisy_wave)
+            batch_vad_segments.append(vad_segments)
+
+        noisy_audios = torch.stack(noisy_audios)
+
+        # assert
+        if torch.any(torch.isnan(noisy_audios)) or torch.any(torch.isinf(noisy_audios)):
+            raise AssertionError("nan or inf in noisy_audios")
+
+        return noisy_audios, batch_vad_segments
+
+
+collate_fn = CollateFunction()
+
+
+def main():
+    args = get_args()
+
+    config = FSMNVadConfig.from_pretrained(
+        pretrained_model_name_or_path=args.config_file,
+    )
+
+    serialization_dir = Path(args.serialization_dir)
+    serialization_dir.mkdir(parents=True, exist_ok=True)
+
+    logger = logging_config(serialization_dir)
+
+    random.seed(config.seed)
+    np.random.seed(config.seed)
+    torch.manual_seed(config.seed)
+    logger.info(f"set seed: {config.seed}")
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    n_gpu = torch.cuda.device_count()
+    logger.info(f"GPU available count: {n_gpu}; device: {device}")
+
+    # datasets
+    train_dataset = VadPaddingJsonlDataset(
+        jsonl_file=args.train_dataset,
+        expected_sample_rate=config.sample_rate,
+        max_wave_value=32768.0,
+        min_snr_db=config.min_snr_db,
+        max_snr_db=config.max_snr_db,
+        # skip=225000,
+    )
+    valid_dataset = VadPaddingJsonlDataset(
+        jsonl_file=args.valid_dataset,
+        expected_sample_rate=config.sample_rate,
+        max_wave_value=32768.0,
+        min_snr_db=config.min_snr_db,
+        max_snr_db=config.max_snr_db,
+    )
+    train_data_loader = DataLoader(
+        dataset=train_dataset,
+        batch_size=config.batch_size,
+        # shuffle=True,
+        sampler=None,
+        # Linux 系统中可以使用多个子进程加载数据, 而在 Windows 系统中不能.
+        num_workers=0 if platform.system() == "Windows" else os.cpu_count() // 2,
+        collate_fn=collate_fn,
+        pin_memory=False,
+        prefetch_factor=None if platform.system() == "Windows" else 2,
+    )
+    valid_data_loader = DataLoader(
+        dataset=valid_dataset,
+        batch_size=config.batch_size,
+        # shuffle=True,
+        sampler=None,
+        # Linux 系统中可以使用多个子进程加载数据, 而在 Windows 系统中不能.
+        num_workers=0 if platform.system() == "Windows" else os.cpu_count() // 2,
+        collate_fn=collate_fn,
+        pin_memory=False,
+        prefetch_factor=None if platform.system() == "Windows" else 2,
+    )
+
+    # models
+    logger.info(f"prepare models. config_file: {args.config_file}")
+    model = FSMNVadPretrainedModel(config).to(device)
+    model.to(device)
+    model.train()
+
+    # optimizer
+    logger.info("prepare optimizer, lr_scheduler, loss_fn, evaluation_metric")
+    optimizer = torch.optim.AdamW(model.parameters(), config.lr)
+
+    # resume training
+    last_step_idx = -1
+    last_epoch = -1
+    for step_idx_str in serialization_dir.glob("steps-*"):
+        step_idx_str = Path(step_idx_str)
+        step_idx = step_idx_str.stem.split("-")[1]
+        step_idx = int(step_idx)
+        if step_idx > last_step_idx:
+            last_step_idx = step_idx
+    # last_epoch = 1
+
+    if last_step_idx != -1:
+        logger.info(f"resume from steps-{last_step_idx}.")
+        model_pt = serialization_dir / f"steps-{last_step_idx}/model.pt"
+
+        logger.info(f"load state dict for model.")
+        with open(model_pt.as_posix(), "rb") as f:
+            state_dict = torch.load(f, map_location="cpu", weights_only=True)
+        model.load_state_dict(state_dict, strict=True)
+
+    if config.lr_scheduler == "CosineAnnealingLR":
+        lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
+            optimizer,
+            last_epoch=last_epoch,
+            # T_max=10 * config.eval_steps,
+            # eta_min=0.01 * config.lr,
+            **config.lr_scheduler_kwargs,
+        )
+    elif config.lr_scheduler == "MultiStepLR":
+        lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
+            optimizer,
+            last_epoch=last_epoch,
+            milestones=[10000, 20000, 30000, 40000, 50000], gamma=0.5
+        )
+    else:
+        raise AssertionError(f"invalid lr_scheduler: {config.lr_scheduler}")
+
+    bce_loss_fn = BCELoss(reduction="mean").to(device)
+    dice_loss_fn = DiceLoss(reduction="mean").to(device)
+
+    vad_accuracy_metrics_fn = VadAccuracy(threshold=0.5)
+    vad_f1_score_metrics_fn = VadF1Score(threshold=0.5)
+
+    # training loop
+
+    # state
+    average_loss = 1000000000
+    average_bce_loss = 1000000000
+    average_dice_loss = 1000000000
+
+    accuracy = -1
+    f1 = -1
+    precision = -1
+    recall = -1
+
+    model_list = list()
+    best_epoch_idx = None
+    best_step_idx = None
+    best_metric = None
+    patience_count = 0
+
+    step_idx = 0 if last_step_idx == -1 else last_step_idx
+
+    logger.info("training")
+    early_stop_flag = False
+    for epoch_idx in range(max(0, last_epoch+1), config.max_epochs):
+        if early_stop_flag:
+            break
+
+        # train
+        model.train()
+        vad_accuracy_metrics_fn.reset()
+        vad_f1_score_metrics_fn.reset()
+
+        total_loss = 0.
+        total_bce_loss = 0.
+        total_dice_loss = 0.
+        total_batches = 0.
+
+        progress_bar_train = tqdm(
+            initial=step_idx,
+            desc="Training; epoch-{}".format(epoch_idx),
+        )
+        for train_batch in train_data_loader:
+            noisy_audios, batch_vad_segments = train_batch
+            noisy_audios: torch.Tensor = noisy_audios.to(device)
+            # noisy_audios shape: [b, num_samples]
+            num_samples = noisy_audios.shape[-1]
+
+            logits, probs = model.forward(noisy_audios)
+
+            targets = BaseVadLoss.get_targets(probs, batch_vad_segments, duration=num_samples / config.sample_rate)
+
+            bce_loss = bce_loss_fn.forward(probs, targets)
+            dice_loss = dice_loss_fn.forward(probs, targets)
+
+            loss = 1.0 * bce_loss + 1.0 * dice_loss
+            if torch.any(torch.isnan(loss)) or torch.any(torch.isinf(loss)):
+                logger.info(f"find nan or inf in loss. continue.")
+                continue
+
+            vad_accuracy_metrics_fn.__call__(probs, targets)
+            vad_f1_score_metrics_fn.__call__(probs, targets)
+
+            optimizer.zero_grad()
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=config.clip_grad_norm)
+            optimizer.step()
+            lr_scheduler.step()
+
+            total_loss += loss.item()
+            total_bce_loss += bce_loss.item()
+            total_dice_loss += dice_loss.item()
+            total_batches += 1
+
+            average_loss = round(total_loss / total_batches, 4)
+            average_bce_loss = round(total_bce_loss / total_batches, 4)
+            average_dice_loss = round(total_dice_loss / total_batches, 4)
+
+            metrics = vad_accuracy_metrics_fn.get_metric()
+            accuracy = metrics["accuracy"]
+            metrics = vad_f1_score_metrics_fn.get_metric()
+            f1 = metrics["f1"]
+            precision = metrics["precision"]
+            recall = metrics["recall"]
+
+            progress_bar_train.update(1)
+            progress_bar_train.set_postfix({
+                "lr": lr_scheduler.get_last_lr()[0],
+                "loss": average_loss,
+                "bce_loss": average_bce_loss,
+                "dice_loss": average_dice_loss,
+                "accuracy": accuracy,
+                "f1": f1,
+                "precision": precision,
+                "recall": recall,
+            })
+
+            # evaluation
+            step_idx += 1
+            if step_idx % config.eval_steps == 0:
+                with torch.no_grad():
+                    torch.cuda.empty_cache()
+
+                    model.eval()
+                    vad_accuracy_metrics_fn.reset()
+                    vad_f1_score_metrics_fn.reset()
+
+                    total_loss = 0.
+                    total_bce_loss = 0.
+                    total_dice_loss = 0.
+                    total_batches = 0.
+
+                    progress_bar_train.close()
+                    progress_bar_eval = tqdm(
+                        desc="Evaluation; steps-{}k".format(int(step_idx/1000)),
+                    )
+                    for eval_batch in valid_data_loader:
+                        noisy_audios, batch_vad_segments = train_batch
+                        noisy_audios: torch.Tensor = noisy_audios.to(device)
+                        # noisy_audios shape: [b, num_samples]
+                        num_samples = noisy_audios.shape[-1]
+
+                        logits, probs = model.forward(noisy_audios)
+
+                        targets = BaseVadLoss.get_targets(probs, batch_vad_segments, duration=num_samples / config.sample_rate)
+
+                        bce_loss = bce_loss_fn.forward(probs, targets)
+                        dice_loss = dice_loss_fn.forward(probs, targets)
+
+                        loss = 1.0 * bce_loss + 1.0 * dice_loss
+                        if torch.any(torch.isnan(loss)) or torch.any(torch.isinf(loss)):
+                            logger.info(f"find nan or inf in loss. continue.")
+                            continue
+
+                        vad_accuracy_metrics_fn.__call__(probs, targets)
+                        vad_f1_score_metrics_fn.__call__(probs, targets)
+
+                        total_loss += loss.item()
+                        total_bce_loss += bce_loss.item()
+                        total_dice_loss += dice_loss.item()
+                        total_batches += 1
+
+                        average_loss = round(total_loss / total_batches, 4)
+                        average_bce_loss = round(total_bce_loss / total_batches, 4)
+                        average_dice_loss = round(total_dice_loss / total_batches, 4)
+
+                        metrics = vad_accuracy_metrics_fn.get_metric()
+                        accuracy = metrics["accuracy"]
+                        metrics = vad_f1_score_metrics_fn.get_metric()
+                        f1 = metrics["f1"]
+                        precision = metrics["precision"]
+                        recall = metrics["recall"]
+
+                        progress_bar_eval.update(1)
+                        progress_bar_eval.set_postfix({
+                            "lr": lr_scheduler.get_last_lr()[0],
+                            "loss": average_loss,
+                            "bce_loss": average_bce_loss,
+                            "dice_loss": average_dice_loss,
+                            "accuracy": accuracy,
+                            "f1": f1,
+                            "precision": precision,
+                            "recall": recall,
+                        })
+
+                model.train()
+                vad_accuracy_metrics_fn.reset()
+                vad_f1_score_metrics_fn.reset()
+
+                total_loss = 0.
+                total_bce_loss = 0.
+                total_dice_loss = 0.
+                total_batches = 0.
+
+                progress_bar_eval.close()
+                progress_bar_train = tqdm(
+                    initial=progress_bar_train.n,
+                    postfix=progress_bar_train.postfix,
+                    desc=progress_bar_train.desc,
+                )
+
+                # save path
+                save_dir = serialization_dir / "steps-{}".format(step_idx)
+                save_dir.mkdir(parents=True, exist_ok=False)
+
+                # save models
+                model.save_pretrained(save_dir.as_posix())
+
+                model_list.append(save_dir)
+                if len(model_list) >= args.num_serialized_models_to_keep:
+                    model_to_delete: Path = model_list.pop(0)
+                    shutil.rmtree(model_to_delete.as_posix())
+
+                # save metric
+                if best_metric is None:
+                    best_epoch_idx = epoch_idx
+                    best_step_idx = step_idx
+                    best_metric = f1
+                elif f1 >= best_metric:
+                    # great is better.
+                    best_epoch_idx = epoch_idx
+                    best_step_idx = step_idx
+                    best_metric = f1
+                else:
+                    pass
+
+                metrics = {
+                    "epoch_idx": epoch_idx,
+                    "best_epoch_idx": best_epoch_idx,
+                    "best_step_idx": best_step_idx,
+                    "loss": average_loss,
+                    "bce_loss": average_bce_loss,
+                    "dice_loss": average_dice_loss,
+
+                    "accuracy": accuracy,
+                }
+                metrics_filename = save_dir / "metrics_epoch.json"
+                with open(metrics_filename, "w", encoding="utf-8") as f:
+                    json.dump(metrics, f, indent=4, ensure_ascii=False)
+
+                # save best
+                best_dir = serialization_dir / "best"
+                if best_epoch_idx == epoch_idx and best_step_idx == step_idx:
+                    if best_dir.exists():
+                        shutil.rmtree(best_dir)
+                    shutil.copytree(save_dir, best_dir)
+
+                # early stop
+                early_stop_flag = False
+                if best_epoch_idx == epoch_idx and best_step_idx == step_idx:
+                    patience_count = 0
+                else:
+                    patience_count += 1
+                if patience_count >= args.patience:
+                    early_stop_flag = True
+
+                # early stop
+                if early_stop_flag:
+                    break
+
+    return
+
+
+if __name__ == "__main__":
+    main()

examples/fsmn_vad_by_webrtcvad/yaml/config.yaml

@@ -0,0 +1,42 @@
+model_name: "fsmn_vad"
+
+# spec
+sample_rate: 8000
+nfft: 512
+win_size: 240
+hop_size: 80
+win_type: hann
+
+# model
+fsmn_input_size: 257
+fsmn_input_affine_size: 140
+fsmn_hidden_size: 250
+fsmn_basic_block_layers: 4
+fsmn_basic_block_hidden_size: 128
+fsmn_basic_block_lorder: 20
+fsmn_basic_block_rorder: 0
+fsmn_basic_block_lstride: 1
+fsmn_basic_block_rstride: 0
+fsmn_output_affine_size: 140
+fsmn_output_size: 1
+
+use_softmax: false
+
+# data
+min_snr_db: -10
+max_snr_db: 20
+
+# train
+lr: 0.001
+lr_scheduler: "CosineAnnealingLR"
+lr_scheduler_kwargs:
+  T_max: 250000
+  eta_min: 0.0001
+
+max_epochs: 100
+clip_grad_norm: 10.0
+seed: 1234
+
+num_workers: 4
+batch_size: 128
+eval_steps: 25000

examples/silero_vad_by_webrtcvad/run.sh

@@ -2,7 +2,7 @@
 
 : <<'END'
 
-bash run.sh --stage 1 --stop_stage 1 --system_version centos \
+bash run.sh --stage 2 --stop_stage 2 --system_version centos \
 --file_folder_name silero-vad-by-webrtcvad-nx2-dns3 \
 --final_model_name silero-vad-by-webrtcvad-nx2-dns3 \
 --noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise/" \

requirements.txt

@@ -11,3 +11,4 @@ torchaudio==2.5.1
 overrides==7.7.0
 webrtcvad==2.0.10
 matplotlib==3.10.3
+google-genai

toolbox/torchaudio/models/vad/cnn_vad/__init__.py

@@ -0,0 +1,6 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+
+
+if __name__ == "__main__":
+    pass

toolbox/torchaudio/models/vad/cnn_vad/modeling_cnn_vad.py

@@ -0,0 +1,6 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+
+
+if __name__ == "__main__":
+    pass

toolbox/torchaudio/models/vad/fsmn_vad/configuration_fsmn_vad.py

@@ -13,8 +13,19 @@ class FSMNVadConfig(PretrainedConfig):
                  hop_size: int = 80,
                  win_type: str = "hann",
 
-                 in_channels: int = 64,
-                 hidden_size: int = 128,
+                 fsmn_input_size: int = 257,
+                 fsmn_input_affine_size: int = 140,
+                 fsmn_hidden_size: int = 250,
+                 fsmn_basic_block_layers: int = 4,
+                 fsmn_basic_block_hidden_size: int = 128,
+                 fsmn_basic_block_lorder: int = 20,
+                 fsmn_basic_block_rorder: int = 0,
+                 fsmn_basic_block_lstride: int = 1,
+                 fsmn_basic_block_rstride: int = 0,
+                 fsmn_output_affine_size: int = 140,
+                 fsmn_output_size: int = 1,
+
+                 use_softmax: bool = False,
 
                  lr: float = 0.001,
                  lr_scheduler: str = "CosineAnnealingLR",
@@ -39,8 +50,19 @@ class FSMNVadConfig(PretrainedConfig):
         self.win_type = win_type
 
         # encoder
-        self.in_channels = in_channels
-        self.hidden_size = hidden_size
+        self.fsmn_input_size = fsmn_input_size
+        self.fsmn_input_affine_size = fsmn_input_affine_size
+        self.fsmn_hidden_size = fsmn_hidden_size
+        self.fsmn_basic_block_layers = fsmn_basic_block_layers
+        self.fsmn_basic_block_hidden_size = fsmn_basic_block_hidden_size
+        self.fsmn_basic_block_lorder = fsmn_basic_block_lorder
+        self.fsmn_basic_block_rorder = fsmn_basic_block_rorder
+        self.fsmn_basic_block_lstride = fsmn_basic_block_lstride
+        self.fsmn_basic_block_rstride = fsmn_basic_block_rstride
+        self.fsmn_output_affine_size = fsmn_output_affine_size
+        self.fsmn_output_size = fsmn_output_size
+
+        self.use_softmax = use_softmax
 
         # train
         self.lr = lr

toolbox/torchaudio/models/vad/fsmn_vad/fsmn_encoder.py

@@ -226,10 +226,6 @@ class FSMN(nn.Module):
         self.out_linear1 = AffineTransform(hidden_size, output_affine_size)
         self.out_linear2 = AffineTransform(output_affine_size, output_size)
 
-        self.use_softmax = use_softmax
-        if self.use_softmax:
-            self.softmax = nn.Softmax(dim=-1)
-
     def forward(self,
                 inputs: torch.Tensor,
                 cache_list: List[torch.Tensor] = None,
@@ -253,8 +249,6 @@ class FSMN(nn.Module):
         outputs = self.out_linear2.forward(x)
         # outputs shape: [b, t, f]
 
-        if self.use_softmax:
-            outputs = self.softmax(outputs)
         return outputs, new_cache_list
 
 
@@ -271,7 +265,6 @@ def main():
         basic_block_rstride=1,
         output_affine_size=16,
         output_size=32,
-        use_softmax=True,
     )
 
     inputs = torch.randn(size=(1, 198, 32), dtype=torch.float32)

toolbox/torchaudio/models/vad/fsmn_vad/modeling_fsmn_vad.py

@@ -41,20 +41,104 @@ class FSMNVadModel(nn.Module):
         )
 
         self.fsmn_encoder = FSMN(
-            input_size=400,
-            input_affine_size=140,
-            hidden_size=250,
-            basic_block_layers=4,
-            basic_block_hidden_size=128,
-            basic_block_lorder=20,
-            basic_block_rorder=0,
-            basic_block_lstride=1,
-            basic_block_rstride=0,
-            output_affine_size=140,
-            output_size=248,
-            use_softmax=True,
+            input_size=config.fsmn_input_size,
+            input_affine_size=config.fsmn_input_affine_size,
+            hidden_size=config.fsmn_hidden_size,
+            basic_block_layers=config.fsmn_basic_block_layers,
+            basic_block_hidden_size=config.fsmn_basic_block_hidden_size,
+            basic_block_lorder=config.fsmn_basic_block_lorder,
+            basic_block_rorder=config.fsmn_basic_block_rorder,
+            basic_block_lstride=config.fsmn_basic_block_lstride,
+            basic_block_rstride=config.fsmn_basic_block_rstride,
+            output_affine_size=config.fsmn_output_affine_size,
+            output_size=config.fsmn_output_size,
         )
 
+        self.use_softmax = config.use_softmax
+        self.sigmoid = nn.Sigmoid()
+        self.softmax = nn.Softmax()
+
+    def forward(self, signal: torch.Tensor):
+        if signal.dim() == 2:
+            signal = torch.unsqueeze(signal, dim=1)
+        _, _, num_samples = signal.shape
+        # signal shape [b, 1, num_samples]
+
+        mags = self.stft.forward(signal)
+        # mags shape: [b, f, t]
+
+        x = torch.transpose(mags, dim0=1, dim1=2)
+        # x shape: [b, t, f]
+
+        logits, _ = self.fsmn_encoder.forward(x)
+
+        if self.use_softmax:
+            probs = self.softmax.forward(logits)
+            # probs shape: [b, t, n]
+        else:
+            probs = self.sigmoid.forward(logits)
+            # probs shape: [b, t, 1]
+        return logits, probs
+
+
+class FSMNVadPretrainedModel(FSMNVadModel):
+    def __init__(self,
+                 config: FSMNVadConfig,
+                 ):
+        super(FSMNVadPretrainedModel, self).__init__(
+            config=config,
+        )
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+        config = FSMNVadConfig.from_pretrained(pretrained_model_name_or_path, **kwargs)
+
+        model = cls(config)
+
+        if os.path.isdir(pretrained_model_name_or_path):
+            ckpt_file = os.path.join(pretrained_model_name_or_path, MODEL_FILE)
+        else:
+            ckpt_file = pretrained_model_name_or_path
+
+        with open(ckpt_file, "rb") as f:
+            state_dict = torch.load(f, map_location="cpu", weights_only=True)
+        model.load_state_dict(state_dict, strict=True)
+        return model
+
+    def save_pretrained(self,
+                        save_directory: Union[str, os.PathLike],
+                        state_dict: Optional[dict] = None,
+                        ):
+
+        model = self
+
+        if state_dict is None:
+            state_dict = model.state_dict()
+
+        os.makedirs(save_directory, exist_ok=True)
+
+        # save state dict
+        model_file = os.path.join(save_directory, MODEL_FILE)
+        torch.save(state_dict, model_file)
+
+        # save config
+        config_file = os.path.join(save_directory, CONFIG_FILE)
+        self.config.to_yaml_file(config_file)
+        return save_directory
+
+
+def main():
+    config = FSMNVadConfig()
+    model = FSMNVadPretrainedModel(config=config)
+
+    noisy = torch.randn(size=(1, 16000), dtype=torch.float32)
+
+    logits, probs = model.forward(noisy)
+    print(f"probs: {probs}")
+    print(f"probs.shape: {logits.shape}")
+    print(f"use_softmax: {config.use_softmax}")
+    return
+
 
 if __name__ == "__main__":
-    pass
+    main()

toolbox/torchaudio/models/vad/fsmn_vad/yaml/config-sigmoid.yaml

@@ -0,0 +1,42 @@
+model_name: "fsmn_vad"
+
+# spec
+sample_rate: 8000
+nfft: 512
+win_size: 240
+hop_size: 80
+win_type: hann
+
+# model
+fsmn_input_size: 257
+fsmn_input_affine_size: 140
+fsmn_hidden_size: 250
+fsmn_basic_block_layers: 4
+fsmn_basic_block_hidden_size: 128
+fsmn_basic_block_lorder: 20
+fsmn_basic_block_rorder: 0
+fsmn_basic_block_lstride: 1
+fsmn_basic_block_rstride: 0
+fsmn_output_affine_size: 140
+fsmn_output_size: 1
+
+use_softmax: false
+
+# data
+min_snr_db: -10
+max_snr_db: 20
+
+# train
+lr: 0.001
+lr_scheduler: "CosineAnnealingLR"
+lr_scheduler_kwargs:
+  T_max: 250000
+  eta_min: 0.0001
+
+max_epochs: 100
+clip_grad_norm: 10.0
+seed: 1234
+
+num_workers: 4
+batch_size: 128
+eval_steps: 25000

toolbox/torchaudio/models/vad/silero_vad/inference_silero_vad.py

@@ -5,6 +5,7 @@ import logging
 from pathlib import Path
 import shutil
 import tempfile, time
+from typing import List
 import zipfile
 
 from scipy.io import wavfile
@@ -18,17 +19,14 @@ torch.set_num_threads(1)
 from project_settings import project_path
 from toolbox.torchaudio.models.vad.silero_vad.configuration_silero_vad import SileroVadConfig
 from toolbox.torchaudio.models.vad.silero_vad.modeling_silero_vad import SileroVadPretrainedModel, MODEL_FILE
-from toolbox.vad.vad import FrameVoiceClassifier, RingVad, process_speech_probs, make_visualization
+from toolbox.torchaudio.utils.visualization import process_speech_probs, make_visualization
 
 
 logger = logging.getLogger("toolbox")
 
 
-class SileroVadVoiceClassifier(FrameVoiceClassifier):
-    def __init__(self,
-                 pretrained_model_path_or_zip_file: str,
-                 device: str = "cpu",
-                 ):
+class InferenceSileroVad(object):
+    def __init__(self, pretrained_model_path_or_zip_file: str, device: str = "cpu"):
         self.pretrained_model_path_or_zip_file = pretrained_model_path_or_zip_file
         self.device = torch.device(device)
 
@@ -62,72 +60,38 @@ class SileroVadVoiceClassifier(FrameVoiceClassifier):
         shutil.rmtree(model_path)
         return config, model
 
-    def predict(self, chunk: np.ndarray) -> float:
-        if chunk.dtype != np.int16:
-            raise AssertionError("signal dtype should be np.int16, instead of {}".format(chunk.dtype))
+    def infer(self, signal: torch.Tensor) -> float:
+        # signal shape: [num_samples,], value between -1 and 1.
 
-        chunk = chunk / 32768
-
-        inputs = torch.tensor(chunk, dtype=torch.float32)
+        inputs = torch.tensor(signal, dtype=torch.float32)
         inputs = torch.unsqueeze(inputs, dim=0)
+        # inputs shape: [1, num_samples,]
 
-        try:
-            logits, _ = self.model.forward(inputs)
-        except RuntimeError as e:
-            print(inputs.shape)
-            raise e
-        # logits shape: [b, t, 1]
-        logits_ = torch.mean(logits, dim=1)
-        # logits_ shape: [b, 1]
-        probs = torch.sigmoid(logits_)
-
-        voice_prob = probs[0][0]
-        return float(voice_prob)
-
-
-class InferenceSileroVad(object):
-    def __init__(self, pretrained_model_path_or_zip_file: str, device: str = "cpu"):
-        self.pretrained_model_path_or_zip_file = pretrained_model_path_or_zip_file
-        self.device = torch.device(device)
-
-        self.voice_classifier = SileroVadVoiceClassifier(pretrained_model_path_or_zip_file, device=device)
-
-        self.ring_vad = RingVad(model=self.voice_classifier,
-                                start_ring_rate=0.2,
-                                end_ring_rate=0.1,
-                                frame_size_ms=30,
-                                frame_step_ms=30,
-                                padding_length_ms=300,
-                                max_silence_length_ms=300,
-                                sample_rate=SAMPLE_RATE,
-                                )
-
-    def vad(self, signal: np.ndarray) -> np.ndarray:
-        self.ring_vad.reset()
-
-        vad_segments = list()
+        with torch.no_grad():
+            logits, probs = self.model.forward(inputs)
 
-        segments = self.ring_vad.vad(signal)
-        vad_segments += segments
-        # last vad segment
-        segments = self.ring_vad.last_vad_segments()
-        vad_segments += segments
-        return vad_segments
+        # probs shape: [b, t, 1]
+        probs = torch.squeeze(probs, dim=-1)
+        # probs shape: [b, t]
 
-    def get_vad_speech_probs(self):
-        result = self.ring_vad.speech_probs
-        return result
+        probs = probs.numpy()
+        probs = probs[0]
+        probs = probs.tolist()
+        return probs
 
-    def get_vad_frame_step(self):
-        result = self.ring_vad.frame_step
-        return result
+    def post_process(self, probs: List[float]):
+        return
 
 
 def get_args():
     parser = argparse.ArgumentParser()
     parser.add_argument(
         "--wav_file",
-        default=(project_path / "data/examples/hado/2f16ca0b-baec-4601-8a1e-7893eb875623.wav").as_posix(),
+        default=(project_path / "data/examples/ai_agent/chinese-4.wav").as_posix(),
+        # default=(project_path / "data/examples/ai_agent/chinese-5.wav").as_posix(),
+        # default=(project_path / "data/examples/hado/b556437e-c68b-4f6d-9eed-2977c29db887.wav").as_posix(),
+        # default=(project_path / "data/examples/hado/eae93a33-8ee0-4d86-8f85-cac5116ae6ef.wav").as_posix(),
+        # default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\speech\nx-speech\en-SG\2025-06-17\active_media_r_0af6bd3a-9aef-4bef-935b-63abfb4d46d8_5.wav",
         type=str,
     )
     args = parser.parse_args()
@@ -143,17 +107,18 @@ def main():
     sample_rate, signal = wavfile.read(args.wav_file)
     if SAMPLE_RATE != sample_rate:
         raise AssertionError
+    signal = signal / (1 << 15)
 
     infer = InferenceSileroVad(
-        pretrained_model_path_or_zip_file=(project_path / "trained_models/silero-vad-webrtcvad-nx2-dns3.zip").as_posix()
+        pretrained_model_path_or_zip_file=(project_path / "trained_models/silero-vad-by-webrtcvad-nx2-dns3.zip").as_posix()
+        # pretrained_model_path_or_zip_file=(project_path / "trained_models/silero-vad-webrtcvad-nx2-dns3.zip").as_posix()
     )
+    frame_step = infer.model.hop_size
 
-    vad_segments = infer.vad(signal)
+    speech_probs = infer.infer(signal)
 
-    speech_probs = infer.get_vad_speech_probs()
-    frame_step = infer.get_vad_frame_step()
+    # print(speech_probs)
 
-    # speech_probs
     speech_probs = process_speech_probs(
         signal=signal,
         speech_probs=speech_probs,
@@ -161,7 +126,7 @@ def main():
     )
 
     # plot
-    make_visualization(signal, speech_probs, SAMPLE_RATE, vad_segments)
+    make_visualization(signal, speech_probs, SAMPLE_RATE)
     return
 
 

toolbox/torchaudio/models/vad/silero_vad/modeling_silero_vad.py

@@ -82,6 +82,11 @@ class Encoder(nn.Module):
 class SileroVadModel(nn.Module):
     def __init__(self, config: SileroVadConfig):
         super(SileroVadModel, self).__init__()
+        self.nfft = config.nfft
+        self.win_size = config.win_size
+        self.hop_size = config.hop_size
+        self.win_type = config.win_type
+
         self.config = config
         self.eps = 1e-12
 
@@ -120,6 +125,11 @@ class SileroVadModel(nn.Module):
         self.sigmoid = nn.Sigmoid()
 
     def forward(self, signal: torch.Tensor):
+        if signal.dim() == 2:
+            signal = torch.unsqueeze(signal, dim=1)
+        _, _, num_samples = signal.shape
+        # signal shape [b, 1, num_samples]
+
         mags = self.stft.forward(signal)
         # mags shape: [b, f, t]
 
@@ -139,6 +149,35 @@ class SileroVadModel(nn.Module):
         # probs shape: [b, t, 1]
         return logits, probs
 
+    def forward_chunk(self, chunk: torch.Tensor):
+        # chunk shape [b, 1, num_samples]
+
+        mags = self.stft.forward(chunk)
+        # mags shape: [b, f, t]
+
+        x = torch.transpose(mags, dim0=1, dim1=2)
+        # x shape: [b, t, f]
+
+        x = self.linear.forward(x)
+        # x shape: [b, t, f']
+
+        return
+
+    def forward_chunk_by_chunk(self, signal: torch.Tensor):
+        if signal.dim() == 2:
+            signal = torch.unsqueeze(signal, dim=1)
+        _, _, num_samples = signal.shape
+        # signal shape [b, 1, num_samples]
+
+        t = (num_samples - self.win_size) // self.hop_size + 1
+        waveform_list = list()
+        for i in range(int(t)):
+            begin = i * self.hop_size
+            end = begin + self.win_size
+            sub_signal = signal[:, :, begin: end]
+
+        return
+
 
 class SileroVadPretrainedModel(SileroVadModel):
     def __init__(self,

toolbox/torchaudio/models/vad/ten_vad/__init__.py

@@ -0,0 +1,12 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+"""
+https://huggingface.co/TEN-framework/ten-vad
+https://zhuanlan.zhihu.com/p/1906832842756976909
+https://github.com/TEN-framework/ten-vad
+
+"""
+
+
+if __name__ == "__main__":
+    pass

toolbox/torchaudio/models/vad/ten_vad/modeling_ten_vad.py

@@ -0,0 +1,6 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+
+
+if __name__ == "__main__":
+    pass

toolbox/torchaudio/models/vad/wav2vec2_vad/__init__.py

@@ -0,0 +1,6 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+
+
+if __name__ == "__main__":
+    pass

toolbox/torchaudio/models/vad/wav2vec2_vad/modeling_wav2vec2.py

@@ -0,0 +1,6 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+
+
+if __name__ == "__main__":
+    pass

toolbox/torchaudio/utils/__init__.py

@@ -0,0 +1,6 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+
+
+if __name__ == "__main__":
+    pass

toolbox/torchaudio/utils/visualization.py

@@ -0,0 +1,33 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+from typing import List
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+
+def process_speech_probs(signal: np.ndarray, speech_probs: List[float], frame_step: int) -> np.ndarray:
+    speech_probs_ = list()
+    for p in speech_probs[1:]:
+        speech_probs_.extend([p] * frame_step)
+
+    pad = (signal.shape[0] - len(speech_probs_))
+    speech_probs_ = speech_probs_ + [0.0] * pad
+    speech_probs_ = np.array(speech_probs_, dtype=np.float32)
+
+    if len(speech_probs_) != len(signal):
+        raise AssertionError
+    return speech_probs_
+
+
+def make_visualization(signal: np.ndarray, speech_probs, sample_rate: int):
+    time = np.arange(0, len(signal)) / sample_rate
+    plt.figure(figsize=(12, 5))
+    plt.plot(time, signal, color='b')
+    plt.plot(time, speech_probs, color='gray')
+    plt.show()
+    return
+
+
+if __name__ == "__main__":
+    pass