update

examples/fsmn_vad_by_webrtcvad/step_4_train_model.py

@@ -270,7 +270,7 @@ def main():
             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 + 0.03 * lsnr_loss
+            loss = 1.0 * bce_loss + 1.0 * dice_loss + 0.3 * 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
@@ -351,7 +351,7 @@ def main():
                         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 + 0.03 * lsnr_loss
+                        loss = 1.0 * bce_loss + 1.0 * dice_loss + 0.3 * 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

examples/silero_vad_by_webrtcvad/step_4_train_model.py

@@ -270,7 +270,7 @@ def main():
             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 + 0.03 * lsnr_loss
+            loss = 1.0 * bce_loss + 1.0 * dice_loss + 0.3 * 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
@@ -351,7 +351,7 @@ def main():
                         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 + 0.03 * lsnr_loss
+                        loss = 1.0 * bce_loss + 1.0 * dice_loss + 0.3 * 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

main.py

@@ -1,14 +1,25 @@
 #!/usr/bin/python3
 # -*- coding: utf-8 -*-
 import argparse
+from functools import lru_cache
+import json
 import logging
 import platform
+import tempfile
+import time
+from typing import Dict, Tuple
 
 import gradio as gr
+import librosa
+import librosa.display
+import matplotlib.pyplot as plt
+import numpy as np
 
 import log
-from project_settings import environment, log_directory, time_zone_info
+from project_settings import environment, project_path, log_directory, time_zone_info
 from toolbox.os.command import Command
+from toolbox.torchaudio.models.vad.fsmn_vad.inference_fsmn_vad import InferenceFSMNVad
+from toolbox.torchaudio.utils.visualization import process_speech_probs
 
 log.setup_size_rotating(log_directory=log_directory, tz_info=time_zone_info)
 
@@ -36,13 +47,136 @@ def shell(cmd: str):
     return Command.popen(cmd)
 
 
+def get_infer_cls_by_model_name(model_name: str):
+    if model_name.__contains__("fsmn"):
+        infer_cls = InferenceFSMNVad
+    else:
+        raise AssertionError
+    return infer_cls
+
+
+vad_engines: Dict[str, dict] = None
+
+
+@lru_cache(maxsize=1)
+def load_vad_model(infer_cls, **kwargs):
+    infer_engine = infer_cls(**kwargs)
+    return infer_engine
+
+
+def generate_image(signal: np.ndarray, speech_probs: np.ndarray, sample_rate: int = 8000, title: str = ""):
+    duration = np.arange(0, len(signal)) / sample_rate
+    plt.figure(figsize=(12, 5))
+    plt.plot(duration, signal, color='b')
+    plt.plot(duration, speech_probs, color='gray')
+    plt.title(title)
+
+    temp_file = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
+    plt.savefig(temp_file.name, bbox_inches="tight")
+    plt.close()
+    return temp_file.name
+
+
+def when_click_vad_button(audio_file_t = None, audio_microphone_t = None, engine: str = None):
+    if audio_file_t is None and audio_microphone_t is None:
+        raise gr.Error(f"audio file and microphone is null.")
+    if audio_file_t is not None and audio_microphone_t is not None:
+        gr.Warning(f"both audio file and microphone file is provided, audio file taking priority.")
+    audio_t: Tuple = audio_file_t or audio_microphone_t
+
+    sample_rate, signal = audio_t
+    audio_duration = signal.shape[-1] // 8000
+    audio = np.array(signal / (1 << 15), dtype=np.float32)
+
+    infer_engine_param = vad_engines.get(engine)
+    if infer_engine_param is None:
+        raise gr.Error(f"invalid denoise engine: {engine}.")
+
+    try:
+        infer_cls = infer_engine_param["infer_cls"]
+        kwargs = infer_engine_param["kwargs"]
+        infer_engine = load_vad_model(infer_cls=infer_cls, **kwargs)
+
+        begin = time.time()
+        vad_info = infer_engine.infer(audio)
+        time_cost = time.time() - begin
+
+        fpr = time_cost / audio_duration
+        info = {
+            "time_cost": round(time_cost, 4),
+            "audio_duration": round(audio_duration, 4),
+            "fpr": round(fpr, 4)
+        }
+        message = json.dumps(info, ensure_ascii=False, indent=4)
+
+        probs = vad_info["probs"]
+        lsnr = vad_info["lsnr"]
+        lsnr = lsnr / np.max(np.abs(lsnr))
+
+        frame_step = infer_engine.config.hop_size
+        probs = process_speech_probs(audio, probs, frame_step)
+        lsnr = process_speech_probs(audio, lsnr, frame_step)
+        probs_image = generate_image(audio, probs)
+        lsnr_image = generate_image(audio, lsnr)
+
+    except Exception as e:
+        raise gr.Error(f"vad failed, error type: {type(e)}, error text: {str(e)}.")
+
+    return probs_image, lsnr_image, message
+
+
 def main():
     args = get_args()
 
+    # engines
+    global vad_engines
+    vad_engines = {
+        filename.stem: {
+            "infer_cls": get_infer_cls_by_model_name(filename.stem),
+            "kwargs": {
+                "pretrained_model_path_or_zip_file": filename.as_posix()
+            }
+        }
+        for filename in (project_path / "trained_models").glob("*.zip")
+        if filename.name not in (
+            "cnn-vad-by-webrtcvad-nx-dns3.zip",
+            "fsmn-vad-by-webrtcvad-nx-dns3.zip",
+            "examples.zip",
+            "sound-2-ch32.zip",
+            "sound-3-ch32.zip",
+            "sound-4-ch32.zip",
+            "sound-8-ch32.zip",
+        )
+    }
+
+    # choices
+    vad_engine_choices = list(vad_engines.keys())
+
     # ui
     with gr.Blocks() as blocks:
         gr.Markdown(value="vad.")
         with gr.Tabs():
+            with gr.TabItem("vad"):
+                with gr.Row():
+                    with gr.Column(variant="panel", scale=5):
+                        with gr.Tabs():
+                            with gr.TabItem("file"):
+                                vad_audio_file = gr.Audio(label="audio")
+                            with gr.TabItem("microphone"):
+                                vad_audio_microphone = gr.Audio(sources="microphone", label="audio")
+
+                        vad_engine = gr.Dropdown(choices=vad_engine_choices, value=vad_engine_choices[0], label="engine")
+                        vad_button = gr.Button(variant="primary")
+                    with gr.Column(variant="panel", scale=5):
+                        vad_vad_image = gr.Image(label="vad")
+                        vad_lsnr_image = gr.Image(label="lsnr")
+                        vad_message = gr.Textbox(lines=1, max_lines=20, label="message")
+
+                vad_button.click(
+                    when_click_vad_button,
+                    inputs=[vad_audio_file, vad_audio_microphone, vad_engine],
+                    outputs=[vad_vad_image, vad_lsnr_image, vad_message]
+                )
             with gr.TabItem("shell"):
                 shell_text = gr.Textbox(label="cmd")
                 shell_button = gr.Button("run")

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

@@ -60,7 +60,7 @@ class InferenceFSMNVad(object):
         shutil.rmtree(model_path)
         return config, model
 
-    def infer(self, signal: torch.Tensor) -> float:
+    def infer(self, signal: torch.Tensor) -> dict:
         # signal shape: [num_samples,], value between -1 and 1.
 
         inputs = torch.tensor(signal, dtype=torch.float32)
@@ -73,11 +73,20 @@ class InferenceFSMNVad(object):
         # 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
+
+        # lsnr shape: [b, t, 1]
+        lsnr = torch.squeeze(lsnr, dim=-1)
+        # lsnr shape: [b, t]
+        lsnr = lsnr.numpy()
+        lsnr = lsnr[0]
+
+        result = {
+            "probs": probs,
+            "lsnr": lsnr,
+        }
+        return result
 
     def post_process(self, probs: List[float]):
         return
@@ -88,11 +97,11 @@ def get_args():
     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/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=(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\speech\en-SG\2025-05-19\active_media_w_8b6e28e2-a238-4c8c-b2e3-426b1fca149b_6.wav",
         # default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\speech\en-SG\2025-05-19\active_media_r_0a56f035-40f6-4530-b852-613f057d718d_6.wav",
         # default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\speech\en-SG\2025-05-19\active_media_r_0ae70b76-3651-4a71-bc0c-9e1429e4c854_5.wav",
@@ -133,18 +142,29 @@ def main():
     )
     frame_step = infer.config.hop_size
 
-    speech_probs = infer.infer(signal)
+    vad_info = infer.infer(signal)
+    speech_probs = vad_info["probs"]
+    lsnr = vad_info["lsnr"]
+
+    lsnr = lsnr / np.max(np.abs(lsnr))
 
-    # print(speech_probs)
+    speech_probs = speech_probs.tolist()
+    lsnr = lsnr.tolist()
 
     speech_probs = process_speech_probs(
         signal=signal,
         speech_probs=speech_probs,
         frame_step=frame_step,
     )
+    lsnr = process_speech_probs(
+        signal=signal,
+        speech_probs=lsnr,
+        frame_step=frame_step,
+    )
 
     # plot
     make_visualization(signal, speech_probs, SAMPLE_RATE)
+    make_visualization(signal, lsnr, SAMPLE_RATE)
     return
 
 

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

@@ -60,7 +60,7 @@ class InferenceSileroVad(object):
         shutil.rmtree(model_path)
         return config, model
 
-    def infer(self, signal: torch.Tensor) -> float:
+    def infer(self, signal: torch.Tensor) -> dict:
         # signal shape: [num_samples,], value between -1 and 1.
 
         inputs = torch.tensor(signal, dtype=torch.float32)
@@ -68,16 +68,25 @@ class InferenceSileroVad(object):
         # inputs shape: [1, num_samples,]
 
         with torch.no_grad():
-            logits, probs = self.model.forward(inputs)
+            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
+
+        # lsnr shape: [b, t, 1]
+        lsnr = torch.squeeze(lsnr, dim=-1)
+        # lsnr shape: [b, t]
+        lsnr = lsnr.numpy()
+        lsnr = lsnr[0]
+
+        result = {
+            "probs": probs,
+            "lsnr": lsnr,
+        }
+        return result
 
     def post_process(self, probs: List[float]):
         return
@@ -87,11 +96,29 @@ 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-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",
+        # 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\speech\en-SG\2025-05-19\active_media_w_8b6e28e2-a238-4c8c-b2e3-426b1fca149b_6.wav",
+        # default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\speech\en-SG\2025-05-19\active_media_r_0a56f035-40f6-4530-b852-613f057d718d_6.wav",
+        # default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\speech\en-SG\2025-05-19\active_media_r_0ae70b76-3651-4a71-bc0c-9e1429e4c854_5.wav",
+        # default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\speech\en-SG\2025-05-19\active_media_r_0d483249-57f8-4d45-b4c6-bda82d6816ae_2.wav",
+        # default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\speech\en-SG\2025-05-19\active_media_r_0d952885-5bc2-4633-81b6-e0e809e113f1_2.wav",
+        # default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\speech\en-SG\2025-05-19\active_media_r_0ddac777-d986-4a5c-9c7c-ff64be0a463d_11.wav",
+
+        # 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",
+        # default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\noise\en-SG\2025-05-19\active_media_r_1aac396f-1661-4f26-ab49-1a4879684567_0.wav",
+        default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\noise\en-SG\2025-05-19\active_media_r_1aac396f-1661-4f26-ab49-1a4879684567_1.wav",
+        # default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\noise\en-SG\2025-05-19\active_media_r_1aff518b-4749-42fc-adfe-64046f9baeb6_0.wav",
+        # default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\noise\en-SG\2025-05-19\active_media_r_1b16f2a3-a8c9-4739-9a76-59faf1c64d79_0.wav",
+        # default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\noise\en-SG\2025-05-19\active_media_r_1b16f2a3-a8c9-4739-9a76-59faf1c64d79_1.wav",
+        # default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\noise\en-SG\2025-05-19\active_media_r_1bb1f22e-9c3a-4aea-b53f-71cc6547a6ee_0.wav",
+        # default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\noise\en-SG\2025-05-19\active_media_r_1dab161b-2a76-4491-abd1-60dba6172f8d_2.wav",
         type=str,
     )
     args = parser.parse_args()
@@ -115,18 +142,29 @@ def main():
     )
     frame_step = infer.model.hop_size
 
-    speech_probs = infer.infer(signal)
+    vad_info = infer.infer(signal)
+    speech_probs = vad_info["probs"]
+    lsnr = vad_info["lsnr"]
+
+    lsnr = lsnr / np.max(np.abs(lsnr))
 
-    # print(speech_probs)
+    speech_probs = speech_probs.tolist()
+    lsnr = lsnr.tolist()
 
     speech_probs = process_speech_probs(
         signal=signal,
         speech_probs=speech_probs,
         frame_step=frame_step,
     )
+    lsnr = process_speech_probs(
+        signal=signal,
+        speech_probs=lsnr,
+        frame_step=frame_step,
+    )
 
     # plot
     make_visualization(signal, speech_probs, SAMPLE_RATE)
+    make_visualization(signal, lsnr, SAMPLE_RATE)
     return