update

examples/cnn_vad_by_webrtcvad/run.sh

@@ -5,14 +5,16 @@
 bash run.sh --stage 2 --stop_stage 2 --system_version centos \
 --file_folder_name cnn-vad-by-webrtcvad-nx-dns3 \
 --final_model_name cnn-vad-by-webrtcvad-nx-dns3 \
---noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise/" \
---speech_dir "/data/tianxing/HuggingDatasets/nx_noise/data/speech/"
+--noise_patterns "/data/tianxing/HuggingDatasets/nx_noise/data/noise/**/*.wav" \
+--speech_patterns "/data/tianxing/HuggingDatasets/nx_noise/data/speech/dns3-speech/**/*.wav \
+/data/tianxing/HuggingDatasets/nx_noise/data/speech/nx-speech2/**/*.wav"
 
 bash run.sh --stage 3 --stop_stage 3 --system_version centos \
 --file_folder_name cnn-vad-by-webrtcvad-nx-dns3 \
 --final_model_name cnn-vad-by-webrtcvad-nx-dns3 \
---noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise/" \
---speech_dir "/data/tianxing/HuggingDatasets/nx_noise/data/speech/"
+--noise_patterns "/data/tianxing/HuggingDatasets/nx_noise/data/noise/**/*.wav" \
+--speech_patterns "/data/tianxing/HuggingDatasets/nx_noise/data/speech/dns3-speech/**/*.wav \
+/data/tianxing/HuggingDatasets/nx_noise/data/speech/nx-speech2/**/*.wav"
 
 
 END
@@ -30,8 +32,8 @@ 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
+noise_patterns=/data/tianxing/HuggingDatasets/nx_noise/data/noise/**/*.wav
+speech_patterns=/data/tianxing/HuggingDatasets/nx_noise/data/speech/**/*.wav
 
 max_count=-1
 
@@ -98,8 +100,8 @@ 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}" \
+  --noise_patterns "${noise_patterns}" \
+  --speech_patterns "${speech_patterns}" \
   --train_dataset "${train_dataset}" \
   --valid_dataset "${valid_dataset}" \
   --max_count "${max_count}" \

examples/cnn_vad_by_webrtcvad/step_1_prepare_data.py

@@ -1,12 +1,14 @@
 #!/usr/bin/python3
 # -*- coding: utf-8 -*-
 import argparse
+from glob import glob
 import json
 import os
 from pathlib import Path
 import random
 import sys
 import time
+from typing import List
 
 pwd = os.path.abspath(os.path.dirname(__file__))
 sys.path.append(os.path.join(pwd, "../../"))
@@ -19,13 +21,13 @@ 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",
+        "--noise_patterns",
+        default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\noise\**\*.wav",
         type=str
     )
     parser.add_argument(
-        "--speech_dir",
-        default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\speech",
+        "--speech_patterns",
+        default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\speech\**\*.wav",
         type=str
     )
 
@@ -46,108 +48,112 @@ def get_args():
     return args
 
 
-def target_second_noise_signal_generator(data_dir: str,
+def target_second_noise_signal_generator(filename_patterns: List[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}")
+        for filename_pattern in filename_patterns:
+            for filename in glob(filename_pattern):
+                signal, _ = librosa.load(filename, 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,
+                            "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,
+                            "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(filename_patterns: List[str],
+                                          min_duration: int = 4,
+                                          max_duration: int = 6,
+                                          sample_rate: int = 8000, max_epoch: int = 1):
+    for epoch_idx in range(max_epoch):
+        for filename_pattern in filename_patterns:
+            for filename in glob(filename_pattern):
+                signal, _ = librosa.load(filename, sr=sample_rate)
+                raw_duration = librosa.get_duration(y=signal, 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}")
 
-            offset = 0.
-            rest_duration = raw_duration
+                if raw_duration < min_duration:
+                    # print(f"duration less than {duration} s. skip filename: {filename.as_posix()}")
+                    continue
 
-            for _ in range(1000):
-                if rest_duration <= 0:
-                    break
-                if rest_duration <= wait_duration:
-                    noise_list.append({
+                if raw_duration < max_duration:
+                    row = {
                         "epoch_idx": epoch_idx,
-                        "filename": filename.as_posix(),
+                        "filename": filename,
                         "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({
+                        "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(),
+                        "filename": filename,
                         "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
+                        "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)
+    noise_patterns = args.noise_patterns
+    noise_patterns = noise_patterns.split(" ")
+    print(f"noise_patterns: {noise_patterns}")
+    speech_patterns = args.speech_patterns
+    speech_patterns = speech_patterns.split(" ")
+    print(f"speech_patterns: {speech_patterns}")
 
     train_dataset = Path(args.train_dataset)
     valid_dataset = Path(args.valid_dataset)
@@ -155,13 +161,13 @@ def main():
     valid_dataset.parent.mkdir(parents=True, exist_ok=True)
 
     noise_generator = target_second_noise_signal_generator(
-        noise_dir.as_posix(),
+        noise_patterns,
         duration=args.duration,
         sample_rate=args.target_sample_rate,
         max_epoch=100000,
     )
     speech_generator = target_second_speech_signal_generator(
-        speech_dir.as_posix(),
+        speech_patterns,
         min_duration=args.min_speech_duration,
         max_duration=args.max_speech_duration,
         sample_rate=args.target_sample_rate,

examples/cnn_vad_by_webrtcvad/step_4_train_model.py

@@ -17,8 +17,6 @@ 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
 

examples/fsmn_vad_by_webrtcvad/run.sh

@@ -2,17 +2,19 @@
 
 : <<'END'
 
-bash run.sh --stage 2 --stop_stage 2 --system_version centos \
+bash run.sh --stage 1 --stop_stage 1 --system_version centos \
 --file_folder_name fsmn-vad-by-webrtcvad-nx2-dns3 \
 --final_model_name fsmn-vad-by-webrtcvad-nx2-dns3 \
---noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise/" \
---speech_dir "/data/tianxing/HuggingDatasets/nx_noise/data/speech/"
+--noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise/**/*.wav" \
+--speech_dir "/data/tianxing/HuggingDatasets/nx_noise/data/speech/dns3-speech/**/*.wav \
+/data/tianxing/HuggingDatasets/nx_noise/data/speech/nx-speech2/**/*.wav"
 
 bash run.sh --stage 3 --stop_stage 3 --system_version centos \
 --file_folder_name fsmn-vad-by-webrtcvad-nx2-dns3 \
 --final_model_name fsmn-vad-by-webrtcvad-nx2-dns3 \
---noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise/" \
---speech_dir "/data/tianxing/HuggingDatasets/nx_noise/data/speech/"
+--noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise/**/*.wav" \
+--speech_dir "/data/tianxing/HuggingDatasets/nx_noise/data/speech/dns3-speech/**/*.wav \
+/data/tianxing/HuggingDatasets/nx_noise/data/speech/nx-speech2/**/*.wav"
 
 
 END
@@ -30,8 +32,8 @@ 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
+noise_patterns=/data/tianxing/HuggingDatasets/nx_noise/data/noise/**/*.wav
+speech_patterns=/data/tianxing/HuggingDatasets/nx_noise/data/speech/**/*.wav
 
 max_count=-1
 
@@ -98,8 +100,8 @@ 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}" \
+  --noise_patterns "${noise_patterns}" \
+  --speech_patterns "${speech_patterns}" \
   --train_dataset "${train_dataset}" \
   --valid_dataset "${valid_dataset}" \
   --max_count "${max_count}" \

examples/fsmn_vad_by_webrtcvad/step_1_prepare_data.py

@@ -1,12 +1,14 @@
 #!/usr/bin/python3
 # -*- coding: utf-8 -*-
 import argparse
+from glob import glob
 import json
 import os
 from pathlib import Path
 import random
 import sys
 import time
+from typing import List
 
 pwd = os.path.abspath(os.path.dirname(__file__))
 sys.path.append(os.path.join(pwd, "../../"))
@@ -19,13 +21,13 @@ 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",
+        "--noise_patterns",
+        default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\noise\**\*.wav",
         type=str
     )
     parser.add_argument(
-        "--speech_dir",
-        default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\speech",
+        "--speech_patterns",
+        default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\speech\**\*.wav",
         type=str
     )
 
@@ -46,108 +48,112 @@ def get_args():
     return args
 
 
-def target_second_noise_signal_generator(data_dir: str,
+def target_second_noise_signal_generator(filename_patterns: List[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}")
+        for filename_pattern in filename_patterns:
+            for filename in glob(filename_pattern):
+                signal, _ = librosa.load(filename, 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,
+                            "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,
+                            "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(filename_patterns: List[str],
+                                          min_duration: int = 4,
+                                          max_duration: int = 6,
+                                          sample_rate: int = 8000, max_epoch: int = 1):
+    for epoch_idx in range(max_epoch):
+        for filename_pattern in filename_patterns:
+            for filename in glob(filename_pattern):
+                signal, _ = librosa.load(filename, sr=sample_rate)
+                raw_duration = librosa.get_duration(y=signal, 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}")
 
-            offset = 0.
-            rest_duration = raw_duration
+                if raw_duration < min_duration:
+                    # print(f"duration less than {duration} s. skip filename: {filename.as_posix()}")
+                    continue
 
-            for _ in range(1000):
-                if rest_duration <= 0:
-                    break
-                if rest_duration <= wait_duration:
-                    noise_list.append({
+                if raw_duration < max_duration:
+                    row = {
                         "epoch_idx": epoch_idx,
-                        "filename": filename.as_posix(),
+                        "filename": filename,
                         "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({
+                        "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(),
+                        "filename": filename,
                         "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
+                        "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)
+    noise_patterns = args.noise_patterns
+    noise_patterns = noise_patterns.split(" ")
+    print(f"noise_patterns: {noise_patterns}")
+    speech_patterns = args.speech_patterns
+    speech_patterns = speech_patterns.split(" ")
+    print(f"speech_patterns: {speech_patterns}")
 
     train_dataset = Path(args.train_dataset)
     valid_dataset = Path(args.valid_dataset)
@@ -155,13 +161,13 @@ def main():
     valid_dataset.parent.mkdir(parents=True, exist_ok=True)
 
     noise_generator = target_second_noise_signal_generator(
-        noise_dir.as_posix(),
+        noise_patterns,
         duration=args.duration,
         sample_rate=args.target_sample_rate,
         max_epoch=100000,
     )
     speech_generator = target_second_speech_signal_generator(
-        speech_dir.as_posix(),
+        speech_patterns,
         min_duration=args.min_speech_duration,
         max_duration=args.max_speech_duration,
         sample_rate=args.target_sample_rate,
@@ -210,7 +216,7 @@ def main():
                 "random1": random1,
             }
             row = json.dumps(row, ensure_ascii=False)
-            if random2 < (1 / 300 / 1):
+            if random2 < (2 / 300):
                 fvalid.write(f"{row}\n")
             else:
                 ftrain.write(f"{row}\n")

examples/fsmn_vad_by_webrtcvad/step_2_make_vad_segments.py

@@ -4,6 +4,7 @@ import argparse
 import json
 import os
 import sys
+from typing import List, Tuple
 
 pwd = os.path.abspath(os.path.dirname(__file__))
 sys.path.append(os.path.join(pwd, "../../"))
@@ -42,6 +43,54 @@ def get_args():
     return args
 
 
+def get_non_silence_segments(waveform: np.ndarray, sample_rate: int = 8000):
+    non_silent_intervals = librosa.effects.split(
+        waveform,
+        top_db=40,  # 静音阈值（单位：dB）
+        frame_length=512,  # 分析帧长
+        hop_length=128  # 帧移
+    )
+
+    # 输出非静音段的时间区间（单位：秒）
+    result = [(start / sample_rate, end / sample_rate) for (start, end) in non_silent_intervals]
+    return result
+
+
+def get_intersection(non_silence: list[tuple[float, float]],
+                     speech: list[tuple[float, float]]) -> list[tuple[float, float]]:
+    """
+    计算语音段与非静音段的交集
+    :param non_silence: 非静音段列表，格式 [(start1, end1), ...]
+    :param speech: 语音检测段列表，格式 [(start2, end2), ...]
+    :return: 交集段列表，格式 [(start, end), ...]
+    """
+    # 按起始时间排序（假设输入已排序可不排）
+    non_silence = sorted(non_silence, key=lambda x: x[0])
+    speech = sorted(speech, key=lambda x: x[0])
+
+    result = []
+    i = j = 0
+
+    while i < len(non_silence) and j < len(speech):
+        ns_start, ns_end = non_silence[i]
+        sp_start, sp_end = speech[j]
+
+        # 计算重叠区间
+        overlap_start = max(ns_start, sp_start)
+        overlap_end = min(ns_end, sp_end)
+
+        if overlap_start < overlap_end:
+            result.append((overlap_start, overlap_end))
+
+        # 移动指针策略：优先处理先结束的区间
+        if ns_end < sp_end:
+            i += 1  # 非静音段先结束
+        else:
+            j += 1  # 语音段先结束
+
+    return result
+
+
 def main():
     args = get_args()
 
@@ -68,8 +117,8 @@ def main():
         end_ring_rate=0.1,
         frame_size_ms=30,
         frame_step_ms=30,
-        padding_length_ms=90,
-        max_silence_length_ms=100,
+        padding_length_ms=30,
+        max_silence_length_ms=0,
         max_speech_length_s=100,
         min_speech_length_s=0.1,
         sample_rate=args.expected_sample_rate,
@@ -114,6 +163,9 @@ def main():
             )
             waveform = np.array(waveform * (1 << 15), dtype=np.int16)
 
+            # non_silence_segments
+            non_silence_segments = get_non_silence_segments(waveform, sample_rate=args.expected_sample_rate)
+
             # vad
             vad_segments = list()
             segments = w_vad.vad(waveform)
@@ -122,6 +174,7 @@ def main():
             vad_segments += segments
             w_vad.reset()
 
+            vad_segments = get_intersection(non_silence_segments, vad_segments)
             row["vad_segments"] = vad_segments
 
             row = json.dumps(row, ensure_ascii=False)
@@ -168,6 +221,9 @@ def main():
             )
             waveform = np.array(waveform * (1 << 15), dtype=np.int16)
 
+            # non_silence_segments
+            non_silence_segments = get_non_silence_segments(waveform, sample_rate=args.expected_sample_rate)
+
             # vad
             vad_segments = list()
             segments = w_vad.vad(waveform)
@@ -176,6 +232,7 @@ def main():
             vad_segments += segments
             w_vad.reset()
 
+            vad_segments = get_intersection(non_silence_segments, vad_segments)
             row["vad_segments"] = vad_segments
 
             row = json.dumps(row, ensure_ascii=False)

examples/fsmn_vad_by_webrtcvad/step_4_train_model.py

@@ -17,8 +17,6 @@ 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
 
@@ -38,7 +36,7 @@ def get_args():
     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("--patience", default=10, type=int)
     parser.add_argument("--serialization_dir", default="serialization_dir", type=str)
 
     parser.add_argument("--config_file", default="yaml/config.yaml", type=str)
@@ -74,22 +72,28 @@ class CollateFunction(object):
 
     def __call__(self, batch: List[dict]):
         noisy_audios = list()
+        clean_audios = list()
         batch_vad_segments = list()
 
         for sample in batch:
             noisy_wave: torch.Tensor = sample["noisy_wave"]
+            clean_wave: torch.Tensor = sample["clean_wave"]
             vad_segments: List[Tuple[float, float]] = sample["vad_segments"]
 
             noisy_audios.append(noisy_wave)
+            clean_audios.append(clean_wave)
             batch_vad_segments.append(vad_segments)
 
         noisy_audios = torch.stack(noisy_audios)
+        clean_audios = torch.stack(clean_audios)
 
         # assert
         if torch.any(torch.isnan(noisy_audios)) or torch.any(torch.isinf(noisy_audios)):
             raise AssertionError("nan or inf in noisy_audios")
+        if torch.any(torch.isnan(clean_audios)) or torch.any(torch.isinf(clean_audios)):
+            raise AssertionError("nan or inf in clean_audios")
 
-        return noisy_audios, batch_vad_segments
+        return noisy_audios, clean_audios, batch_vad_segments
 
 
 collate_fn = CollateFunction()
@@ -214,6 +218,7 @@ def main():
     average_loss = 1000000000
     average_bce_loss = 1000000000
     average_dice_loss = 1000000000
+    average_lsnr_loss = 1000000000
 
     accuracy = -1
     f1 = -1
@@ -242,6 +247,7 @@ def main():
         total_loss = 0.
         total_bce_loss = 0.
         total_dice_loss = 0.
+        total_lsnr_loss = 0.
         total_batches = 0.
 
         progress_bar_train = tqdm(
@@ -249,19 +255,22 @@ def main():
             desc="Training; epoch-{}".format(epoch_idx),
         )
         for train_batch in train_data_loader:
-            noisy_audios, batch_vad_segments = train_batch
+            noisy_audios, clean_audios, batch_vad_segments = train_batch
             noisy_audios: torch.Tensor = noisy_audios.to(device)
+            clean_audios: torch.Tensor = clean_audios.to(device)
             # noisy_audios shape: [b, num_samples]
             num_samples = noisy_audios.shape[-1]
 
-            logits, probs = model.forward(noisy_audios)
+            logits, probs, lsnr = model.forward(noisy_audios)
+            lsnr = torch.squeeze(lsnr, dim=-1)
 
             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)
+            lsnr_loss = model.lsnr_loss_fn(lsnr, clean_audios, noisy_audios)
 
-            loss = 1.0 * bce_loss + 1.0 * dice_loss
+            loss = 1.0 * bce_loss + 1.0 * dice_loss + 0.03 * lsnr_loss
             if torch.any(torch.isnan(loss)) or torch.any(torch.isinf(loss)):
                 logger.info(f"find nan or inf in loss. continue.")
                 continue
@@ -278,11 +287,13 @@ def main():
             total_loss += loss.item()
             total_bce_loss += bce_loss.item()
             total_dice_loss += dice_loss.item()
+            total_lsnr_loss += lsnr_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)
+            average_lsnr_loss = round(total_lsnr_loss / total_batches, 4)
 
             metrics = vad_accuracy_metrics_fn.get_metric()
             accuracy = metrics["accuracy"]
@@ -297,6 +308,7 @@ def main():
                 "loss": average_loss,
                 "bce_loss": average_bce_loss,
                 "dice_loss": average_dice_loss,
+                "lsnr_loss": average_lsnr_loss,
                 "accuracy": accuracy,
                 "f1": f1,
                 "precision": precision,
@@ -316,6 +328,7 @@ def main():
                     total_loss = 0.
                     total_bce_loss = 0.
                     total_dice_loss = 0.
+                    total_lsnr_loss = 0.
                     total_batches = 0.
 
                     progress_bar_train.close()
@@ -323,19 +336,22 @@ def main():
                         desc="Evaluation; steps-{}k".format(int(step_idx/1000)),
                     )
                     for eval_batch in valid_data_loader:
-                        noisy_audios, batch_vad_segments = eval_batch
+                        noisy_audios, clean_audios, batch_vad_segments = eval_batch
                         noisy_audios: torch.Tensor = noisy_audios.to(device)
+                        clean_audios: torch.Tensor = clean_audios.to(device)
                         # noisy_audios shape: [b, num_samples]
                         num_samples = noisy_audios.shape[-1]
 
-                        logits, probs = model.forward(noisy_audios)
+                        logits, probs, lsnr = model.forward(noisy_audios)
+                        lsnr = torch.squeeze(lsnr, dim=-1)
 
                         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)
+                        lsnr_loss = model.lsnr_loss_fn(lsnr, clean_audios, noisy_audios)
 
-                        loss = 1.0 * bce_loss + 1.0 * dice_loss
+                        loss = 1.0 * bce_loss + 1.0 * dice_loss + 0.03 * lsnr_loss
                         if torch.any(torch.isnan(loss)) or torch.any(torch.isinf(loss)):
                             logger.info(f"find nan or inf in loss. continue.")
                             continue
@@ -346,11 +362,13 @@ def main():
                         total_loss += loss.item()
                         total_bce_loss += bce_loss.item()
                         total_dice_loss += dice_loss.item()
+                        total_lsnr_loss += lsnr_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)
+                        average_lsnr_loss = round(total_lsnr_loss / total_batches, 4)
 
                         metrics = vad_accuracy_metrics_fn.get_metric()
                         accuracy = metrics["accuracy"]
@@ -365,6 +383,7 @@ def main():
                             "loss": average_loss,
                             "bce_loss": average_bce_loss,
                             "dice_loss": average_dice_loss,
+                            "lsnr_loss": average_lsnr_loss,
                             "accuracy": accuracy,
                             "f1": f1,
                             "precision": precision,
@@ -378,6 +397,7 @@ def main():
                 total_loss = 0.
                 total_bce_loss = 0.
                 total_dice_loss = 0.
+                total_lsnr_loss = 0.
                 total_batches = 0.
 
                 progress_bar_eval.close()
@@ -419,8 +439,12 @@ def main():
                     "loss": average_loss,
                     "bce_loss": average_bce_loss,
                     "dice_loss": average_dice_loss,
+                    "lsnr_loss": average_lsnr_loss,
 
                     "accuracy": accuracy,
+                    "f1": f1,
+                    "precision": precision,
+                    "recall": recall,
                 }
                 metrics_filename = save_dir / "metrics_epoch.json"
                 with open(metrics_filename, "w", encoding="utf-8") as f:

examples/fsmn_vad_by_webrtcvad/yaml/config.yaml

@@ -18,9 +18,13 @@ fsmn_basic_block_rorder: 0
 fsmn_basic_block_lstride: 1
 fsmn_basic_block_rstride: 0
 fsmn_output_affine_size: 140
-fsmn_output_size: 1
+fsmn_output_size: 2
 
-use_softmax: false
+# lsnr
+n_frame: 3
+min_local_snr_db: -15
+max_local_snr_db: 30
+norm_tau: 1.
 
 # data
 min_snr_db: -10

examples/silero_vad_by_webrtcvad/run.sh

@@ -5,14 +5,16 @@
 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/"
+--noise_patterns "/data/tianxing/HuggingDatasets/nx_noise/data/noise/**/*.wav" \
+--speech_patterns "/data/tianxing/HuggingDatasets/nx_noise/data/speech/dns3-speech/**/*.wav \
+/data/tianxing/HuggingDatasets/nx_noise/data/speech/nx-speech2/**/*.wav"
 
 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/"
+--noise_patterns "/data/tianxing/HuggingDatasets/nx_noise/data/noise/**/*.wav" \
+--speech_patterns "/data/tianxing/HuggingDatasets/nx_noise/data/speech/dns3-speech/**/*.wav \
+/data/tianxing/HuggingDatasets/nx_noise/data/speech/nx-speech2/**/*.wav"
 
 
 END
@@ -30,8 +32,8 @@ 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
+noise_patterns=/data/tianxing/HuggingDatasets/nx_noise/data/noise/**/*.wav
+speech_patterns=/data/tianxing/HuggingDatasets/nx_noise/data/speech/**/*.wav
 
 max_count=-1
 
@@ -98,8 +100,8 @@ 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}" \
+  --noise_patterns "${noise_patterns}" \
+  --speech_patterns "${speech_patterns}" \
   --train_dataset "${train_dataset}" \
   --valid_dataset "${valid_dataset}" \
   --max_count "${max_count}" \

examples/silero_vad_by_webrtcvad/step_1_prepare_data.py

@@ -1,12 +1,14 @@
 #!/usr/bin/python3
 # -*- coding: utf-8 -*-
 import argparse
+from glob import glob
 import json
 import os
 from pathlib import Path
 import random
 import sys
 import time
+from typing import List
 
 pwd = os.path.abspath(os.path.dirname(__file__))
 sys.path.append(os.path.join(pwd, "../../"))
@@ -19,13 +21,13 @@ 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",
+        "--noise_patterns",
+        default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\noise\**\*.wav",
         type=str
     )
     parser.add_argument(
-        "--speech_dir",
-        default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\speech",
+        "--speech_patterns",
+        default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\speech\**\*.wav",
         type=str
     )
 
@@ -46,108 +48,112 @@ def get_args():
     return args
 
 
-def target_second_noise_signal_generator(data_dir: str,
+def target_second_noise_signal_generator(filename_patterns: List[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}")
+        for filename_pattern in filename_patterns:
+            for filename in glob(filename_pattern):
+                signal, _ = librosa.load(filename, 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,
+                            "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,
+                            "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(filename_patterns: List[str],
+                                          min_duration: int = 4,
+                                          max_duration: int = 6,
+                                          sample_rate: int = 8000, max_epoch: int = 1):
+    for epoch_idx in range(max_epoch):
+        for filename_pattern in filename_patterns:
+            for filename in glob(filename_pattern):
+                signal, _ = librosa.load(filename, sr=sample_rate)
+                raw_duration = librosa.get_duration(y=signal, 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}")
 
-            offset = 0.
-            rest_duration = raw_duration
+                if raw_duration < min_duration:
+                    # print(f"duration less than {duration} s. skip filename: {filename.as_posix()}")
+                    continue
 
-            for _ in range(1000):
-                if rest_duration <= 0:
-                    break
-                if rest_duration <= wait_duration:
-                    noise_list.append({
+                if raw_duration < max_duration:
+                    row = {
                         "epoch_idx": epoch_idx,
-                        "filename": filename.as_posix(),
+                        "filename": filename,
                         "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({
+                        "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(),
+                        "filename": filename,
                         "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
+                        "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)
+    noise_patterns = args.noise_patterns
+    noise_patterns = noise_patterns.split(" ")
+    print(f"noise_patterns: {noise_patterns}")
+    speech_patterns = args.speech_patterns
+    speech_patterns = speech_patterns.split(" ")
+    print(f"speech_patterns: {speech_patterns}")
 
     train_dataset = Path(args.train_dataset)
     valid_dataset = Path(args.valid_dataset)
@@ -155,13 +161,13 @@ def main():
     valid_dataset.parent.mkdir(parents=True, exist_ok=True)
 
     noise_generator = target_second_noise_signal_generator(
-        noise_dir.as_posix(),
+        noise_patterns,
         duration=args.duration,
         sample_rate=args.target_sample_rate,
         max_epoch=100000,
     )
     speech_generator = target_second_speech_signal_generator(
-        speech_dir.as_posix(),
+        speech_patterns,
         min_duration=args.min_speech_duration,
         max_duration=args.max_speech_duration,
         sample_rate=args.target_sample_rate,
@@ -210,7 +216,7 @@ def main():
                 "random1": random1,
             }
             row = json.dumps(row, ensure_ascii=False)
-            if random2 < (1 / 300 / 1):
+            if random2 < (2 / 300):
                 fvalid.write(f"{row}\n")
             else:
                 ftrain.write(f"{row}\n")

examples/silero_vad_by_webrtcvad/step_4_train_model.py

@@ -17,8 +17,6 @@ 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
 
@@ -38,7 +36,7 @@ def get_args():
     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("--patience", default=10, type=int)
     parser.add_argument("--serialization_dir", default="serialization_dir", type=str)
 
     parser.add_argument("--config_file", default="yaml/config.yaml", type=str)
@@ -74,22 +72,28 @@ class CollateFunction(object):
 
     def __call__(self, batch: List[dict]):
         noisy_audios = list()
+        clean_audios = list()
         batch_vad_segments = list()
 
         for sample in batch:
             noisy_wave: torch.Tensor = sample["noisy_wave"]
+            clean_wave: torch.Tensor = sample["clean_wave"]
             vad_segments: List[Tuple[float, float]] = sample["vad_segments"]
 
             noisy_audios.append(noisy_wave)
+            clean_audios.append(clean_wave)
             batch_vad_segments.append(vad_segments)
 
         noisy_audios = torch.stack(noisy_audios)
+        clean_audios = torch.stack(clean_audios)
 
         # assert
         if torch.any(torch.isnan(noisy_audios)) or torch.any(torch.isinf(noisy_audios)):
             raise AssertionError("nan or inf in noisy_audios")
+        if torch.any(torch.isnan(clean_audios)) or torch.any(torch.isinf(clean_audios)):
+            raise AssertionError("nan or inf in clean_audios")
 
-        return noisy_audios, batch_vad_segments
+        return noisy_audios, clean_audios, batch_vad_segments
 
 
 collate_fn = CollateFunction()
@@ -214,6 +218,7 @@ def main():
     average_loss = 1000000000
     average_bce_loss = 1000000000
     average_dice_loss = 1000000000
+    average_lsnr_loss = 1000000000
 
     accuracy = -1
     f1 = -1
@@ -242,6 +247,7 @@ def main():
         total_loss = 0.
         total_bce_loss = 0.
         total_dice_loss = 0.
+        total_lsnr_loss = 0.
         total_batches = 0.
 
         progress_bar_train = tqdm(

examples/silero_vad_by_webrtcvad/yaml/config.yaml

@@ -11,6 +11,12 @@ win_type: hann
 in_channels: 64
 hidden_size: 128
 
+# lsnr
+n_frame: 3
+min_local_snr_db: -15
+max_local_snr_db: 30
+norm_tau: 1.
+
 # data
 min_snr_db: -10
 max_snr_db: 20

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

@@ -0,0 +1,138 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import argparse
+import logging
+from pathlib import Path
+import shutil
+import tempfile, time
+from typing import List
+import zipfile
+
+from scipy.io import wavfile
+import librosa
+import numpy as np
+import torch
+import torchaudio
+
+torch.set_num_threads(1)
+
+from project_settings import project_path
+from toolbox.torchaudio.models.vad.cnn_vad.configuration_cnn_vad import CNNVadConfig
+from toolbox.torchaudio.models.vad.cnn_vad.modeling_cnn_vad import CNNVadPretrainedModel, MODEL_FILE
+from toolbox.torchaudio.utils.visualization import process_speech_probs, make_visualization
+
+
+logger = logging.getLogger("toolbox")
+
+
+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)
+
+        logger.info(f"loading model; model_file: {self.pretrained_model_path_or_zip_file}")
+        config, model = self.load_models(self.pretrained_model_path_or_zip_file)
+        logger.info(f"model loading completed; model_file: {self.pretrained_model_path_or_zip_file}")
+
+        self.config = config
+        self.model = model
+        self.model.to(device)
+        self.model.eval()
+
+    def load_models(self, model_path: str):
+        model_path = Path(model_path)
+        if model_path.name.endswith(".zip"):
+            with zipfile.ZipFile(model_path.as_posix(), "r") as f_zip:
+                out_root = Path(tempfile.gettempdir()) / "cc_vad"
+                out_root.mkdir(parents=True, exist_ok=True)
+                f_zip.extractall(path=out_root)
+            model_path = out_root / model_path.stem
+
+        config = CNNVadConfig.from_pretrained(
+            pretrained_model_name_or_path=model_path.as_posix(),
+        )
+        model = CNNVadPretrainedModel.from_pretrained(
+            pretrained_model_name_or_path=model_path.as_posix(),
+        )
+        model.to(self.device)
+        model.eval()
+
+        shutil.rmtree(model_path)
+        return config, model
+
+    def infer(self, signal: torch.Tensor) -> float:
+        # signal shape: [num_samples,], value between -1 and 1.
+
+        inputs = torch.tensor(signal, dtype=torch.float32)
+        inputs = torch.unsqueeze(inputs, dim=0)
+        # inputs shape: [1, num_samples,]
+
+        with torch.no_grad():
+            logits, probs, lsnr = self.model.forward(inputs)
+
+        # probs shape: [b, t, 1]
+        probs = torch.squeeze(probs, dim=-1)
+        # probs shape: [b, t]
+
+        probs = probs.numpy()
+        probs = probs[0]
+        probs = probs.tolist()
+        return probs
+
+    def post_process(self, probs: List[float]):
+        return
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--wav_file",
+        # 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=(project_path / "data/examples/speech/active_media_r_0ba69730-66a4-4ecd-8929-ef58f18f4612_2.wav").as_posix(),
+        # default=(project_path / "data/examples/speech/active_media_r_2a2f472b-a0b8-4fd5-b1c4-1aedc5d2ce57_0.wav").as_posix(),
+        # default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\noise\en-SG\2025-05-19\active_media_r_0b8a8e80-52af-423b-8877-03a78b1e6e43_0.wav",
+        # default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\noise\en-SG\2025-05-19\active_media_r_0ebffb68-6490-4a8b-8eb6-eb82443d7d75_0.wav",
+        default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\noise\en-SG\2025-05-19\active_media_r_0f6ec933-90df-447b-aca4-6ddc149452ab_0.wav",
+        type=str,
+    )
+    args = parser.parse_args()
+    return args
+
+
+SAMPLE_RATE = 8000
+
+
+def main():
+    args = get_args()
+
+    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/cnn-vad-by-webrtcvad-nx-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
+
+    speech_probs = infer.infer(signal)
+
+    # print(speech_probs)
+
+    speech_probs = process_speech_probs(
+        signal=signal,
+        speech_probs=speech_probs,
+        frame_step=frame_step,
+    )
+
+    # plot
+    make_visualization(signal, speech_probs, SAMPLE_RATE)
+    return
+
+
+if __name__ == "__main__":
+    main()

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

@@ -23,9 +23,12 @@ class FSMNVadConfig(PretrainedConfig):
                  fsmn_basic_block_lstride: int = 1,
                  fsmn_basic_block_rstride: int = 0,
                  fsmn_output_affine_size: int = 140,
-                 fsmn_output_size: int = 1,
+                 fsmn_output_size: int = 2,
 
-                 use_softmax: bool = False,
+                 n_frame: int = 3,
+                 min_local_snr_db: float = -15,
+                 max_local_snr_db: float = 30,
+                 norm_tau: float = 1.,
 
                  min_snr_db: float = -10,
                  max_snr_db: float = 20,
@@ -65,7 +68,11 @@ class FSMNVadConfig(PretrainedConfig):
         self.fsmn_output_affine_size = fsmn_output_affine_size
         self.fsmn_output_size = fsmn_output_size
 
-        self.use_softmax = use_softmax
+        # lsnr
+        self.n_frame = n_frame
+        self.min_local_snr_db = min_local_snr_db
+        self.max_local_snr_db = max_local_snr_db
+        self.norm_tau = norm_tau
 
         # data snr
         self.min_snr_db = min_snr_db

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

@@ -15,48 +15,111 @@ from typing import Optional, Union
 
 import torch
 import torch.nn as nn
+from torch.nn import functional as F
 
 from toolbox.torchaudio.configuration_utils import CONFIG_FILE
 from toolbox.torchaudio.models.vad.fsmn_vad.configuration_fsmn_vad import FSMNVadConfig
 from toolbox.torchaudio.modules.conv_stft import ConvSTFT
 from toolbox.torchaudio.models.vad.fsmn_vad.fsmn_encoder import FSMN
+from toolbox.torchaudio.modules.local_snr_target import LocalSnrTarget
 
 
 MODEL_FILE = "model.pt"
 
 
 class FSMNVadModel(nn.Module):
-    def __init__(self, config: FSMNVadConfig):
+    def __init__(self,
+                 sample_rate: int,
+                 nfft: int,
+                 win_size: int,
+                 hop_size: int,
+                 win_type: int,
+
+                 fsmn_input_size: int,
+                 fsmn_input_affine_size: int,
+                 fsmn_hidden_size: int,
+                 fsmn_basic_block_layers: int,
+                 fsmn_basic_block_hidden_size: int,
+                 fsmn_basic_block_lorder: int,
+                 fsmn_basic_block_rorder: int,
+                 fsmn_basic_block_lstride: int,
+                 fsmn_basic_block_rstride: int,
+                 fsmn_output_affine_size: int,
+
+                 n_frame: int,
+                 min_local_snr_db: float,
+                 max_local_snr_db: float,
+                 ):
         super(FSMNVadModel, self).__init__()
-        self.config = config
+        self.sample_rate = sample_rate
+        self.nfft = nfft
+        self.win_size = win_size
+        self.hop_size = hop_size
+        self.win_type = win_type
+
+        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.n_frame = n_frame
+        self.min_local_snr_db = min_local_snr_db
+        self.max_local_snr_db = max_local_snr_db
+
         self.eps = 1e-12
 
         self.stft = ConvSTFT(
-            nfft=config.nfft,
-            win_size=config.win_size,
-            hop_size=config.hop_size,
-            win_type=config.win_type,
+            nfft=self.nfft,
+            win_size=self.win_size,
+            hop_size=self.hop_size,
+            win_type=self.win_type,
             power=1,
             requires_grad=False
         )
+        self.complex_stft = ConvSTFT(
+            nfft=self.nfft,
+            win_size=self.win_size,
+            hop_size=self.hop_size,
+            win_type=self.win_type,
+            power=None,
+            requires_grad=False
+        )
 
         self.fsmn_encoder = FSMN(
-            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,
+            input_size=self.fsmn_input_size,
+            input_affine_size=self.fsmn_input_affine_size,
+            hidden_size=self.fsmn_hidden_size,
+            basic_block_layers=self.fsmn_basic_block_layers,
+            basic_block_hidden_size=self.fsmn_basic_block_hidden_size,
+            basic_block_lorder=self.fsmn_basic_block_lorder,
+            basic_block_rorder=self.fsmn_basic_block_rorder,
+            basic_block_lstride=self.fsmn_basic_block_lstride,
+            basic_block_rstride=self.fsmn_basic_block_rstride,
+            output_affine_size=self.fsmn_output_affine_size,
+            output_size=2,
+            # output_size=self.fsmn_output_size,
         )
 
-        self.use_softmax = config.use_softmax
-        self.sigmoid = nn.Sigmoid()
-        self.softmax = nn.Softmax()
+        # lsnr
+        self.lsnr_scale = self.max_local_snr_db - self.min_local_snr_db
+        self.lsnr_offset = self.min_local_snr_db
+
+        self.lsnr_fn = LocalSnrTarget(
+            sample_rate=self.sample_rate,
+            nfft=self.nfft,
+            win_size=self.win_size,
+            hop_size=self.hop_size,
+            n_frame=self.n_frame,
+            min_local_snr=self.min_local_snr_db,
+            max_local_snr=self.max_local_snr_db,
+            db=True,
+        )
 
     def forward(self, signal: torch.Tensor):
         if signal.dim() == 2:
@@ -71,14 +134,49 @@ class FSMNVadModel(nn.Module):
         # x shape: [b, t, f]
 
         logits, _ = self.fsmn_encoder.forward(x)
+        # logits shape: [b, t, 2]
 
-        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
+        splits = torch.split(logits, split_size_or_sections=[1, 1], dim=-1)
+        vad_logits = splits[0]
+        snr_logits = splits[1]
+        # shape: [b, t, 1]
+        vad_probs = F.sigmoid(vad_logits)
+        # vad_probs shape: [b, t, 1]
+
+        lsnr = F.sigmoid(snr_logits) * self.lsnr_scale + self.lsnr_offset
+        # lsnr shape: [b, t, 1]
+
+        return vad_logits, vad_probs, lsnr
+
+    def lsnr_loss_fn(self, lsnr: torch.Tensor, clean: torch.Tensor, noisy: torch.Tensor):
+        if noisy.shape != clean.shape:
+            raise AssertionError("Input signals must have the same shape")
+        noise = noisy - clean
+
+        if clean.dim() == 2:
+            clean = torch.unsqueeze(clean, dim=1)
+        if noise.dim() == 2:
+            noise = torch.unsqueeze(noise, dim=1)
+
+        stft_clean = self.complex_stft.forward(clean)
+        stft_noise = self.complex_stft.forward(noise)
+        # shape: [b, f, t]
+        stft_clean = torch.transpose(stft_clean, dim0=1, dim1=2)
+        stft_noise = torch.transpose(stft_noise, dim0=1, dim1=2)
+        # shape: [b, t, f]
+        stft_clean = torch.unsqueeze(stft_clean, dim=1)
+        stft_noise = torch.unsqueeze(stft_noise, dim=1)
+        # shape: [b, 1, t, f]
+
+        # lsnr shape: [b, 1, t]
+        lsnr = lsnr.squeeze(1)
+        # lsnr shape: [b, t]
+
+        lsnr_gth = self.lsnr_fn.forward(stft_clean, stft_noise)
+        # lsnr_gth shape: [b, t]
+
+        loss = F.mse_loss(lsnr, lsnr_gth)
+        return loss
 
 
 class FSMNVadPretrainedModel(FSMNVadModel):
@@ -86,8 +184,26 @@ class FSMNVadPretrainedModel(FSMNVadModel):
                  config: FSMNVadConfig,
                  ):
         super(FSMNVadPretrainedModel, self).__init__(
-            config=config,
+            sample_rate=config.sample_rate,
+            nfft=config.nfft,
+            win_size=config.win_size,
+            hop_size=config.hop_size,
+            win_type=config.win_type,
+            fsmn_input_size=config.fsmn_input_size,
+            fsmn_input_affine_size=config.fsmn_input_affine_size,
+            fsmn_hidden_size=config.fsmn_hidden_size,
+            fsmn_basic_block_layers=config.fsmn_basic_block_layers,
+            fsmn_basic_block_hidden_size=config.fsmn_basic_block_hidden_size,
+            fsmn_basic_block_lorder=config.fsmn_basic_block_lorder,
+            fsmn_basic_block_rorder=config.fsmn_basic_block_rorder,
+            fsmn_basic_block_lstride=config.fsmn_basic_block_lstride,
+            fsmn_basic_block_rstride=config.fsmn_basic_block_rstride,
+            fsmn_output_affine_size=config.fsmn_output_affine_size,
+            n_frame=config.n_frame,
+            min_local_snr_db=config.min_local_snr_db,
+            max_local_snr_db=config.max_local_snr_db,
         )
+        self.config = config
 
     @classmethod
     def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
@@ -133,10 +249,11 @@ def main():
 
     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}")
+    logits, probs, lsnr = model.forward(noisy)
+    print(f"logits.shape: {logits.shape}")
+    print(f"probs.shape: {probs.shape}")
+    print(f"lsnr.shape: {lsnr.shape}")
+
     return
 
 

toolbox/torchaudio/models/vad/fsmn_vad/yaml/{config-sigmoid.yaml → config.yaml}

@@ -18,9 +18,13 @@ fsmn_basic_block_rorder: 0
 fsmn_basic_block_lstride: 1
 fsmn_basic_block_rstride: 0
 fsmn_output_affine_size: 140
-fsmn_output_size: 1
+fsmn_output_size: 2
 
-use_softmax: false
+# lsnr
+n_frame: 3
+min_local_snr_db: -15
+max_local_snr_db: 30
+norm_tau: 1.
 
 # data
 min_snr_db: -10

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

@@ -16,6 +16,11 @@ class SileroVadConfig(PretrainedConfig):
                  in_channels: int = 64,
                  hidden_size: int = 128,
 
+                 n_frame: int = 3,
+                 min_local_snr_db: float = -15,
+                 max_local_snr_db: float = 30,
+                 norm_tau: float = 1.,
+
                  min_snr_db: float = -10,
                  max_snr_db: float = 20,
 
@@ -45,6 +50,12 @@ class SileroVadConfig(PretrainedConfig):
         self.in_channels = in_channels
         self.hidden_size = hidden_size
 
+        # lsnr
+        self.n_frame = n_frame
+        self.min_local_snr_db = min_local_snr_db
+        self.max_local_snr_db = max_local_snr_db
+        self.norm_tau = norm_tau
+
         # data snr
         self.min_snr_db = min_snr_db
         self.max_snr_db = max_snr_db

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

@@ -13,10 +13,12 @@ from typing import Optional, Union
 
 import torch
 import torch.nn as nn
+from torch.nn import functional as F
 
 from toolbox.torchaudio.configuration_utils import CONFIG_FILE
 from toolbox.torchaudio.models.vad.silero_vad.configuration_silero_vad import SileroVadConfig
 from toolbox.torchaudio.modules.conv_stft import ConvSTFT
+from toolbox.torchaudio.modules.local_snr_target import LocalSnrTarget
 
 
 MODEL_FILE = "model.pt"
@@ -80,50 +82,99 @@ class Encoder(nn.Module):
 
 
 class SileroVadModel(nn.Module):
-    def __init__(self, config: SileroVadConfig):
+    def __init__(self,
+                 sample_rate: int,
+                 nfft: int,
+                 win_size: int,
+                 hop_size: int,
+                 win_type: int,
+
+                 in_channels: int,
+                 hidden_size: int,
+
+                 n_frame: int,
+                 min_local_snr_db: float,
+                 max_local_snr_db: float,
+
+                 ):
         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.sample_rate = sample_rate
+        self.nfft = nfft
+        self.win_size = win_size
+        self.hop_size = hop_size
+        self.win_type = win_type
+
+        self.in_channels = in_channels
+        self.hidden_size = hidden_size
+
+        self.n_frame = n_frame
+        self.min_local_snr_db = min_local_snr_db
+        self.max_local_snr_db = max_local_snr_db
 
-        self.config = config
         self.eps = 1e-12
 
         self.stft = ConvSTFT(
-            nfft=config.nfft,
-            win_size=config.win_size,
-            hop_size=config.hop_size,
-            win_type=config.win_type,
+            nfft=nfft,
+            win_size=win_size,
+            hop_size=hop_size,
+            win_type=win_type,
             power=1,
             requires_grad=False
         )
+        self.complex_stft = ConvSTFT(
+            nfft=nfft,
+            win_size=win_size,
+            hop_size=hop_size,
+            win_type=win_type,
+            power=None,
+            requires_grad=False
+        )
 
         self.linear = nn.Linear(
-            in_features=(config.nfft // 2 + 1),
-            out_features=config.in_channels,
+            in_features=(self.nfft // 2 + 1),
+            out_features=self.in_channels,
         )
 
         self.encoder = Encoder(
-            in_channels=config.in_channels,
-            out_channels=config.hidden_size,
+            in_channels=self.in_channels,
+            out_channels=self.hidden_size,
         )
 
         self.lstm = nn.LSTM(
-            input_size=config.hidden_size,
-            hidden_size=config.hidden_size,
+            input_size=self.hidden_size,
+            hidden_size=self.hidden_size,
             bidirectional=False,
             batch_first=True
         )
 
-        self.classifier = nn.Sequential(
-            nn.Linear(config.hidden_size, 32),
+        # vad
+        self.vad_fc = nn.Sequential(
+            nn.Linear(self.hidden_size, 32),
             nn.ReLU(),
             nn.Linear(32, 1),
         )
-
         self.sigmoid = nn.Sigmoid()
 
+        # lsnr
+        self.lsnr_fc = nn.Sequential(
+            nn.Linear(self.hidden_size, 1),
+            nn.Sigmoid()
+        )
+        self.lsnr_scale = self.max_local_snr_db - self.min_local_snr_db
+        self.lsnr_offset = self.min_local_snr_db
+
+        # lsnr
+        self.lsnr_fn = LocalSnrTarget(
+            sample_rate=self.sample_rate,
+            nfft=self.nfft,
+            win_size=self.win_size,
+            hop_size=self.hop_size,
+            n_frame=self.n_frame,
+            min_local_snr=self.min_local_snr_db,
+            max_local_snr=self.max_local_snr_db,
+            db=True,
+        )
+
     def forward(self, signal: torch.Tensor):
         if signal.dim() == 2:
             signal = torch.unsqueeze(signal, dim=1)
@@ -143,40 +194,46 @@ class SileroVadModel(nn.Module):
         # x shape: [b, t, f]
 
         x, _ = self.lstm.forward(x)
-        logits = self.classifier.forward(x)
+
+        logits = self.vad_fc.forward(x)
         # logits shape: [b, t, 1]
         probs = self.sigmoid.forward(logits)
         # probs shape: [b, t, 1]
-        return logits, probs
 
-    def forward_chunk(self, chunk: torch.Tensor):
-        # chunk shape [b, 1, num_samples]
+        lsnr = self.lsnr_fc.forward(x) * self.lsnr_scale + self.lsnr_offset
+        # lsnr shape: [b, t, 1]
 
-        mags = self.stft.forward(chunk)
-        # mags shape: [b, f, t]
+        return logits, probs, lsnr
 
-        x = torch.transpose(mags, dim0=1, dim1=2)
-        # x shape: [b, t, f]
+    def lsnr_loss_fn(self, lsnr: torch.Tensor, clean: torch.Tensor, noisy: torch.Tensor):
+        if noisy.shape != clean.shape:
+            raise AssertionError("Input signals must have the same shape")
+        noise = noisy - clean
 
-        x = self.linear.forward(x)
-        # x shape: [b, t, f']
+        if clean.dim() == 2:
+            clean = torch.unsqueeze(clean, dim=1)
+        if noise.dim() == 2:
+            noise = torch.unsqueeze(noise, dim=1)
 
-        return
+        stft_clean = self.complex_stft.forward(clean)
+        stft_noise = self.complex_stft.forward(noise)
+        # shape: [b, f, t]
+        stft_clean = torch.transpose(stft_clean, dim0=1, dim1=2)
+        stft_noise = torch.transpose(stft_noise, dim0=1, dim1=2)
+        # shape: [b, t, f]
+        stft_clean = torch.unsqueeze(stft_clean, dim=1)
+        stft_noise = torch.unsqueeze(stft_noise, dim=1)
+        # shape: [b, 1, t, f]
 
-    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]
+        # lsnr shape: [b, 1, t]
+        lsnr = lsnr.squeeze(1)
+        # lsnr shape: [b, t]
 
-        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]
+        lsnr_gth = self.lsnr_fn.forward(stft_clean, stft_noise)
+        # lsnr_gth shape: [b, t]
 
-        return
+        loss = F.mse_loss(lsnr, lsnr_gth)
+        return loss
 
 
 class SileroVadPretrainedModel(SileroVadModel):
@@ -184,8 +241,18 @@ class SileroVadPretrainedModel(SileroVadModel):
                  config: SileroVadConfig,
                  ):
         super(SileroVadPretrainedModel, self).__init__(
-            config=config,
+            sample_rate=config.sample_rate,
+            nfft=config.nfft,
+            win_size=config.win_size,
+            hop_size=config.hop_size,
+            win_type=config.win_type,
+            in_channels=config.in_channels,
+            hidden_size=config.hidden_size,
+            n_frame=config.n_frame,
+            min_local_snr_db=config.min_local_snr_db,
+            max_local_snr_db=config.max_local_snr_db,
         )
+        self.config = config
 
     @classmethod
     def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
@@ -227,13 +294,14 @@ class SileroVadPretrainedModel(SileroVadModel):
 
 def main():
     config = SileroVadConfig()
-    model = SileroVadModel(config=config)
+    model = SileroVadPretrainedModel(config=config)
 
     noisy = torch.randn(size=(1, 16000), dtype=torch.float32)
 
-    logits, probs = model.forward(noisy)
-    print(f"logits: {probs}")
+    logits, probs, lsnr = model.forward(noisy)
     print(f"logits.shape: {logits.shape}")
+    print(f"probs.shape: {probs.shape}")
+    print(f"lsnr.shape: {lsnr.shape}")
 
     return
 

toolbox/torchaudio/models/vad/silero_vad/yaml/config.yaml

@@ -11,6 +11,12 @@ win_type: hann
 in_channels: 64
 hidden_size: 128
 
+# lsnr
+n_frame: 3
+min_local_snr_db: -15
+max_local_snr_db: 30
+norm_tau: 1.
+
 # data
 min_snr_db: -10
 max_snr_db: 20