update

main.py

@@ -23,6 +23,7 @@ from toolbox.os.command import Command
 from toolbox.torchaudio.models.vad.fsmn_vad.inference_fsmn_vad_onnx import InferenceFSMNVadOnnx
 from toolbox.torchaudio.models.vad.silero_vad.inference_silero_vad import InferenceSileroVad
 from toolbox.torchaudio.utils.visualization import process_speech_probs
+from toolbox.vad.utils import PostProcess
 
 log.setup_size_rotating(log_directory=log_directory, tz_info=time_zone_info)
 
@@ -98,7 +99,12 @@ def generate_image(signal: np.ndarray, speech_probs: np.ndarray, sample_rate: in
     return temp_file.name
 
 
-def when_click_vad_button(audio_file_t = None, audio_microphone_t = None, engine: str = None):
+def when_click_vad_button(audio_file_t = None, audio_microphone_t = None,
+                          start_ring_rate: float = 0.5, end_ring_rate: float = 0.3,
+                          min_silence_length: int = 2,
+                          max_speech_length: int = 10000, min_speech_length: int = 10,
+                          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:
@@ -136,15 +142,28 @@ def when_click_vad_button(audio_file_t = None, audio_microphone_t = None, engine
         lsnr = lsnr / 30
 
         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)
+        probs_ = process_speech_probs(audio, probs, frame_step)
+        probs_image = generate_image(audio, probs_)
+
+        lsnr_ = process_speech_probs(audio, lsnr, frame_step)
+        lsnr_image = generate_image(audio, lsnr_)
+
+        # post process
+        vad_post_process = PostProcess(
+            start_ring_rate=start_ring_rate,
+            end_ring_rate=end_ring_rate,
+            min_silence_length=min_silence_length,
+            max_speech_length=max_speech_length,
+            min_speech_length=min_speech_length
+        )
+        vad = vad_post_process.post_process(probs)
+        vad_ = process_speech_probs(audio, vad, frame_step)
+        vad_image = generate_image(audio, vad_)
 
     except Exception as e:
         raise gr.Error(f"vad failed, error type: {type(e)}, error text: {str(e)}.")
 
-    return probs_image, lsnr_image, message
+    return vad_image, probs_image, lsnr_image, message
 
 
 def main():
@@ -218,22 +237,36 @@ def main():
                             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")
+                        with gr.Row():
+                            vad_start_ring_rate = gr.Slider(minimum=0, maximum=1, value=0.5, step=0.1, label="start_ring_rate")
+                            vad_end_ring_rate = gr.Slider(minimum=0, maximum=1, value=0.3, step=0.1, label="end_ring_rate")
+                            vad_min_silence_length = gr.Number(value=2, label="min_silence_length")
+                        with gr.Row():
+                            vad_max_speech_length = gr.Number(value=100000, label="max_speech_length")
+                            vad_min_speech_length = gr.Number(value=10, label="min_speech_length")
+                            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_prob_image = gr.Image(label="prob")
                         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],
+                    inputs=[
+                        vad_audio_file, vad_audio_microphone,
+                        vad_start_ring_rate, vad_end_ring_rate,
+                        vad_min_silence_length,
+                        vad_max_speech_length, vad_min_speech_length,
+                        vad_engine,
+                    ],
+                    outputs=[vad_vad_image, vad_prob_image, vad_lsnr_image, vad_message],
                 )
                 gr.Examples(
                     examples=examples,
                     inputs=[vad_audio_file, vad_audio_microphone, vad_engine],
-                    outputs=[vad_vad_image, vad_lsnr_image, vad_message],
+                    outputs=[vad_vad_image, vad_prob_image, vad_lsnr_image, vad_message],
                     fn=when_click_vad_button,
                     # cache_examples=True,
                     # cache_mode="lazy",

toolbox/vad/utils.py

@@ -0,0 +1,135 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import collections
+
+from typing import List, Tuple
+
+
+class PostProcess(object):
+    def __init__(self,
+                 start_ring_rate: float = 0.5,
+                 end_ring_rate: float = 0.5,
+                 min_silence_length: int = 1,
+                 max_speech_length: float = 10,
+                 min_speech_length: float = 2,
+                 ):
+        self.start_ring_rate = start_ring_rate
+        self.end_ring_rate = end_ring_rate
+        self.max_speech_length = max_speech_length
+        self.min_speech_length = min_speech_length
+        self.min_silence_length = min_silence_length
+
+        # segments
+        self.ring_buffer = collections.deque(maxlen=10)
+        self.triggered = False
+
+        # vad segments
+        self.is_first_segment = True
+        self.start_idx: int = -1
+        self.end_idx: int = -1
+
+        # speech probs
+        self.voiced_frames: List[Tuple[int, float]] = list()
+
+    def segments_generator(self, probs: List[float]):
+        for idx, prob in enumerate(probs):
+            if not self.triggered:
+                self.ring_buffer.append((idx, prob))
+                num_voiced = sum([p for _, p in self.ring_buffer])
+
+                if num_voiced > self.start_ring_rate * self.ring_buffer.maxlen:
+                    self.triggered = True
+                    for idx_prob_t in self.ring_buffer:
+                        self.voiced_frames.append(idx_prob_t)
+                continue
+
+            idx_prob_t = (idx, prob)
+            self.voiced_frames.append(idx_prob_t)
+            self.ring_buffer.append(idx_prob_t)
+            num_voiced = sum([p for _, p in self.ring_buffer])
+
+            if num_voiced < self.end_ring_rate * self.ring_buffer.maxlen:
+                segment = [
+                    self.voiced_frames[0][0],
+                    self.voiced_frames[-1][0],
+                ]
+                yield segment
+                self.triggered = False
+                self.ring_buffer.clear()
+                self.voiced_frames: List[Tuple[int, float]] = list()
+                continue
+
+    def vad_segments_generator(self, segments_generator):
+        segments = list(segments_generator)
+
+        for i, segment in enumerate(segments):
+            start = segment[0]
+            end = segment[1]
+
+            if self.start_idx == -1 and self.end_idx == -1:
+                self.start_idx = start
+                self.end_idx = end
+                continue
+
+            if self.end_idx - self.start_idx > self.max_speech_length:
+                end_ = self.start_idx + self.max_speech_length
+                vad_segment = [self.start_idx, end_]
+                yield vad_segment
+                self.start_idx = end_
+
+            silence_length = start - self.end_idx
+            if silence_length < self.min_silence_length:
+                self.end_idx = end
+                continue
+
+            if self.end_idx - self.start_idx < self.min_speech_length:
+                self.start_idx = start
+                self.end_idx = end
+                continue
+
+            vad_segment = [self.start_idx, self.end_idx]
+            yield vad_segment
+            self.start_idx = start
+            self.end_idx = end
+
+    def vad(self, probs: List[float]) -> List[list]:
+        segments = self.segments_generator(probs)
+        vad_segments = self.vad_segments_generator(segments)
+        vad_segments = list(vad_segments)
+        return vad_segments
+
+    def last_vad_segments(self) -> List[list]:
+        # last segments
+        if len(self.voiced_frames) == 0:
+            segments = []
+        else:
+            segment = [
+                self.voiced_frames[0][0],
+                self.voiced_frames[-1][0]
+            ]
+            segments = [segment]
+
+        # last vad segments
+        vad_segments = self.vad_segments_generator(segments)
+        vad_segments = list(vad_segments)
+
+        if self.start_idx > 1e-5 and self.end_idx > 1e-5:
+            vad_segments = vad_segments + [[self.start_idx, self.end_idx]]
+        return vad_segments
+
+    def post_process(self, probs: List[float]):
+        vad_segments = list()
+        segments = self.vad(probs)
+        vad_segments += segments
+        segments = self.last_vad_segments()
+        vad_segments += segments
+
+        result = [0] * len(probs)
+        for begin, end in vad_segments:
+            result[begin: end] = [1] * (end - begin)
+
+        return result
+
+
+if __name__ == "__main__":
+    pass