update

examples/silerovad/vad.py

@@ -6,6 +6,8 @@ https://github.com/snakers4/silero-vad
 """
 import argparse
 
+import matplotlib.pyplot as plt
+import numpy as np
 from scipy.io import wavfile
 import torch
 
@@ -35,6 +37,33 @@ def get_args():
     return args
 
 
+def make_visualization(probs, step):
+    import pandas as pd
+    pd.DataFrame({'probs': probs},
+                 index=[x * step for x in range(len(probs))]).plot(figsize=(16, 8),
+                                                                   kind='area', ylim=[0, 1.05], xlim=[0, len(probs) * step],
+                                                                   xlabel='seconds',
+                                                                   ylabel='speech probability',
+                                                                   colormap='tab20')
+
+
+def plot(signal, sample_rate, speeches):
+    time = np.arange(0, len(signal)) / sample_rate
+
+    plt.figure(figsize=(12, 5))
+
+    plt.plot(time, signal / 32768, color="b")
+
+    for speech in speeches:
+        start = speech["start"]
+        end = speech["end"]
+        plt.axvline(x=start, ymin=0.25, ymax=0.75, color="g", linestyle="--")
+        plt.axvline(x=end, ymin=0.25, ymax=0.75, color="r", linestyle="--")
+
+    plt.show()
+    return
+
+
 def main():
     args = get_args()
 
@@ -45,7 +74,6 @@ def main():
     sample_rate, signal = wavfile.read(args.wav_file)
     signal = signal / 32768
     signal = torch.tensor(signal, dtype=torch.float32)
-    print(signal)
 
     min_speech_samples = sample_rate * args.min_speech_duration_ms / 1000
     speech_pad_samples = sample_rate * args.speech_pad_ms / 1000
@@ -53,9 +81,11 @@ def main():
     min_silence_samples = sample_rate * args.min_silence_duration_ms / 1000
     min_silence_samples_at_max_speech = sample_rate * 98 / 1000
 
+    audio_length_samples = len(signal)
+
     # probs
     speech_probs = []
-    for start in range(0, len(signal), args.window_size_samples):
+    for start in range(0, audio_length_samples, args.window_size_samples):
         chunk = signal[start: start + args.window_size_samples]
         if len(chunk) < args.window_size_samples:
             chunk = torch.nn.functional.pad(chunk, (0, int(args.window_size_samples - len(chunk))))
@@ -63,8 +93,6 @@ def main():
         speech_prob = model(chunk, sample_rate).item()
         speech_probs.append(speech_prob)
 
-    print(speech_probs)
-
     # segments
     triggered = False
     speeches = list()
@@ -107,6 +135,7 @@ def main():
                 temp_end = args.window_size_samples * i
             if ((args.window_size_samples * i) - temp_end) > min_silence_samples_at_max_speech:
                 prev_end = temp_end
+
             if (args.window_size_samples * i) - temp_end < min_silence_samples:
                 continue
             else:
@@ -118,10 +147,32 @@ def main():
                 triggered = False
                 continue
 
-    if current_speech and (args.audio_length_samples - current_speech["start"]) > min_speech_samples:
-        current_speech["end"] = args.audio_length_samples
+    if current_speech and (audio_length_samples - current_speech["start"]) > min_speech_samples:
+        current_speech["end"] = audio_length_samples
         speeches.append(current_speech)
 
+    for i, speech in enumerate(speeches):
+        if i == 0:
+            speech["start"] = int(max(0, speech["start"] - speech_pad_samples))
+        if i != len(speeches) - 1:
+            silence_duration = speeches[i+1]["start"] - speech["end"]
+            if silence_duration < 2 * speech_pad_samples:
+                speech["end"] += int(silence_duration // 2)
+                speeches[i+1]["start"] = int(max(0, speeches[i+1]["start"] - silence_duration // 2))
+            else:
+                speech["end"] = int(min(audio_length_samples, speech["end"] + speech_pad_samples))
+                speeches[i+1]["start"] = int(max(0, speeches[i+1]["start"] - speech_pad_samples))
+        else:
+            speech["end"] = int(min(audio_length_samples, speech["end"] + speech_pad_samples))
+
+    # in seconds
+    for speech_dict in speeches:
+        speech_dict["start"] = round(speech_dict["start"] / sample_rate, 1)
+        speech_dict["end"] = round(speech_dict["end"] / sample_rate, 1)
+
+    print(speeches)
+    plot(signal, sample_rate, speeches)
+
     return
 
 

main.py

@@ -15,44 +15,65 @@ from PIL import Image
 
 from project_settings import project_path, temp_directory
 from toolbox.webrtcvad.vad import WebRTCVad
+from toolbox.vad.vad import Vad, WebRTCVoiceClassifier, SileroVoiceClassifier
 
 
 def get_args():
     parser = argparse.ArgumentParser()
     parser.add_argument(
-        "--webrtcvad_examples_file",
-        default=(project_path / "webrtcvad_examples.json").as_posix(),
+        "--ring_vad_examples_file",
+        default=(project_path / "ring_vad_examples.json").as_posix(),
         type=str
     )
     args = parser.parse_args()
     return args
 
 
-webrtcvad: WebRTCVad = None
+vad: Vad = None
 
 
-def click_webrtcvad_button(audio: Tuple[int, np.ndarray],
-                           agg: int = 3,
-                           frame_duration_ms: int = 30,
-                           padding_duration_ms: int = 300,
-                           silence_duration_threshold: float = 0.3,
-                           ):
-    global webrtcvad
+def click_ring_vad_button(audio: Tuple[int, np.ndarray],
+                          model_name: str,
+                          agg: int = 3,
+                          frame_duration_ms: int = 30,
+                          padding_duration_ms: int = 300,
+                          silence_duration_threshold: float = 0.3,
+                          start_ring_rate: float = 0.9,
+                          end_ring_rate: float = 0.1,
+                          ):
+    global vad
 
+    if audio is None:
+        return None, "please upload audio."
     sample_rate, signal = audio
 
-    webrtcvad = WebRTCVad(agg=int(agg),
-                          frame_duration_ms=frame_duration_ms,
-                          padding_duration_ms=padding_duration_ms,
-                          silence_duration_threshold=silence_duration_threshold,
-                          sample_rate=sample_rate,
-                          )
-
-    vad_segments = list()
-    segments = webrtcvad.vad(signal)
-    vad_segments += segments
-    segments = webrtcvad.last_vad_segments()
-    vad_segments += segments
+    if model_name == "webrtcvad" and frame_duration_ms not in (10, 20, 30):
+        return None, "only 10, 20, 30 available for `frame_duration_ms`."
+
+    if model_name == "webrtcvad":
+        model = WebRTCVoiceClassifier(agg=agg)
+    elif model_name == "silerovad":
+        model = SileroVoiceClassifier(model_name=(project_path / "pretrained_models/silero_vad/silero_vad.jit").as_posix())
+    else:
+        return None, "`model_name` not valid."
+
+    vad = Vad(model=model,
+              start_ring_rate=start_ring_rate,
+              end_ring_rate=end_ring_rate,
+              frame_duration_ms=frame_duration_ms,
+              padding_duration_ms=padding_duration_ms,
+              silence_duration_threshold=silence_duration_threshold,
+              sample_rate=sample_rate,
+              )
+
+    try:
+        vad_segments = list()
+        segments = vad.vad(signal)
+        vad_segments += segments
+        segments = vad.last_vad_segments()
+        vad_segments += segments
+    except Exception as e:
+        return None, str(e)
 
     time = np.arange(0, len(signal)) / sample_rate
     plt.figure(figsize=(12, 5))
@@ -77,8 +98,8 @@ def main():
     """
 
     # examples
-    with open(args.webrtcvad_examples_file, "r", encoding="utf-8") as f:
-        webrtcvad_examples = json.load(f)
+    with open(args.ring_vad_examples_file, "r", encoding="utf-8") as f:
+        ring_vad_examples = json.load(f)
 
     # ui
     with gr.Blocks() as blocks:
@@ -87,50 +108,62 @@ def main():
         with gr.Row():
             with gr.Column(scale=5):
                 with gr.Tabs():
-                    with gr.TabItem("webrtcvad"):
+                    with gr.TabItem("ring_vad"):
                         gr.Markdown(value="")
 
                         with gr.Row():
                             with gr.Column(scale=1):
-                                webrtcvad_wav = gr.Audio(label="wav")
+                                ring_wav = gr.Audio(label="wav")
+
+                                with gr.Row():
+                                    ring_model_name = gr.Dropdown(choices=["webrtcvad", "silerovad"], value="webrtcvad", label="model_name")
+
+                                with gr.Row():
+                                    ring_agg = gr.Dropdown(choices=[1, 2, 3], value=3, label="agg")
+                                    ring_frame_duration_ms = gr.Slider(minimum=0, maximum=100, value=30, label="frame_duration_ms")
 
                                 with gr.Row():
-                                    webrtcvad_agg = gr.Dropdown(choices=[1, 2, 3], value=3, label="agg")
-                                    webrtcvad_frame_duration_ms = gr.Slider(minimum=0, maximum=100, value=30, label="frame_duration_ms")
+                                    ring_padding_duration_ms = gr.Slider(minimum=0, maximum=1000, value=300, label="padding_duration_ms")
+                                    ring_silence_duration_threshold = gr.Slider(minimum=0, maximum=1.0, value=0.3, step=0.1, label="silence_duration_threshold")
 
                                 with gr.Row():
-                                    webrtcvad_padding_duration_ms = gr.Slider(minimum=0, maximum=1000, value=300, label="padding_duration_ms")
-                                    webrtcvad_silence_duration_threshold = gr.Slider(minimum=0, maximum=1.0, value=0.3, step=0.1, label="silence_duration_threshold")
+                                    ring_start_ring_rate = gr.Slider(minimum=0, maximum=1, value=0.9, step=0.1, label="start_ring_rate")
+                                    ring_end_ring_rate = gr.Slider(minimum=0, maximum=1, value=0.1, step=0.1, label="end_ring_rate")
 
-                                webrtcvad_button = gr.Button("retrieval", variant="primary")
+                                ring_button = gr.Button("retrieval", variant="primary")
 
                             with gr.Column(scale=1):
-                                webrtcvad_image = gr.Image(label="image", height=300, width=720, show_label=False)
-                                webrtcvad_end_points = gr.TextArea(label="end_points", max_lines=35)
+                                ring_image = gr.Image(label="image", height=300, width=720, show_label=False)
+                                ring_end_points = gr.TextArea(label="end_points", max_lines=35)
 
                         gr.Examples(
-                            examples=webrtcvad_examples,
+                            examples=ring_vad_examples,
                             inputs=[
-                                webrtcvad_wav, webrtcvad_agg, webrtcvad_frame_duration_ms,
-                                webrtcvad_padding_duration_ms, webrtcvad_silence_duration_threshold
+                                ring_wav,
+                                ring_model_name, ring_agg, ring_frame_duration_ms,
+                                ring_padding_duration_ms, ring_silence_duration_threshold,
+                                ring_start_ring_rate, ring_end_ring_rate
                             ],
-                            outputs=[webrtcvad_image, webrtcvad_end_points],
-                            fn=click_webrtcvad_button
+                            outputs=[ring_image, ring_end_points],
+                            fn=click_ring_vad_button
                         )
 
                         # click event
-                        webrtcvad_button.click(
-                            click_webrtcvad_button,
+                        ring_button.click(
+                            click_ring_vad_button,
                             inputs=[
-                                webrtcvad_wav, webrtcvad_agg, webrtcvad_frame_duration_ms,
-                                webrtcvad_padding_duration_ms, webrtcvad_silence_duration_threshold
+                                ring_wav,
+                                ring_model_name, ring_agg, ring_frame_duration_ms,
+                                ring_padding_duration_ms, ring_silence_duration_threshold,
+                                ring_start_ring_rate, ring_end_ring_rate
                             ],
-                            outputs=[webrtcvad_image, webrtcvad_end_points],
+                            outputs=[ring_image, ring_end_points],
                         )
 
     blocks.queue().launch(
         share=False if platform.system() == "Windows" else False,
-        server_name="0.0.0.0", server_port=7860
+        server_name="127.0.0.1" if platform.system() == "Windows" else "0.0.0.0",
+        server_port=7860
     )
     return
 

ring_vad_examples.json

@@ -0,0 +1,18 @@
+[
+    [
+        "data/early_media/3300999628164249998.wav",
+        "webrtcvad", 3, 30, 300, 0.3, 0.9, 0.1
+    ],
+    [
+        "data/early_media/3300999628164852605.wav",
+        "webrtcvad", 3, 30, 300, 0.3, 0.9, 0.1
+    ],
+    [
+        "data/early_media/3300999628164249998.wav",
+        "silerovad", 3, 35, 350, 0.35, 0.5, 0.5
+    ],
+    [
+        "data/early_media/3300999628164852605.wav",
+        "silerovad", 3, 35, 350, 0.35, 0.5, 0.5
+    ]
+]

toolbox/vad/__init__.py

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

toolbox/vad/vad.py

@@ -0,0 +1,299 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import argparse
+import collections
+from typing import List
+
+import matplotlib.pyplot as plt
+import numpy as np
+from scipy.io import wavfile
+import torch
+import webrtcvad
+
+from project_settings import project_path
+
+
+class FrameVoiceClassifier(object):
+    def predict(self, chunk: np.ndarray) -> float:
+        raise NotImplementedError
+
+
+class WebRTCVoiceClassifier(FrameVoiceClassifier):
+    def __init__(self,
+                 agg: int = 3,
+                 sample_rate: int = 8000
+                 ):
+        self.agg = agg
+        self.sample_rate = sample_rate
+
+        self.model = webrtcvad.Vad(mode=agg)
+
+    def predict(self, chunk: np.ndarray) -> float:
+        if chunk.dtype != np.int16:
+            raise AssertionError("signal dtype should be np.int16, instead of {}".format(chunk.dtype))
+
+        audio_bytes = bytes(chunk)
+        is_speech = self.model.is_speech(audio_bytes, self.sample_rate)
+        return 1.0 if is_speech else 0.0
+
+
+class SileroVoiceClassifier(FrameVoiceClassifier):
+    def __init__(self,
+                 model_name: str,
+                 sample_rate: int = 8000):
+        self.model_name = model_name
+        self.sample_rate = sample_rate
+
+        with open(self.model_name, "rb") as f:
+            model = torch.jit.load(f, map_location="cpu")
+        self.model = model
+        self.model.reset_states()
+
+    def predict(self, chunk: np.ndarray) -> float:
+        if self.sample_rate / len(chunk) > 31.25:
+            raise AssertionError("chunk samples number {} is less than {}".format(len(chunk), self.sample_rate / 31.25))
+        if chunk.dtype != np.int16:
+            raise AssertionError("signal dtype should be np.int16, instead of {}".format(chunk.dtype))
+
+        chunk = chunk / 32768
+        chunk = torch.tensor(chunk, dtype=torch.float32)
+        speech_prob = self.model(chunk, self.sample_rate).item()
+        return float(speech_prob)
+
+
+class Frame(object):
+    def __init__(self, signal: np.ndarray, timestamp, duration):
+        self.signal = signal
+        self.timestamp = timestamp
+        self.duration = duration
+
+
+class Vad(object):
+    def __init__(self,
+                 model: FrameVoiceClassifier,
+                 start_ring_rate: float = 0.5,
+                 end_ring_rate: float = 0.5,
+                 frame_duration_ms: int = 30,
+                 padding_duration_ms: int = 300,
+                 silence_duration_threshold: float = 0.3,
+                 sample_rate: int = 8000
+                 ):
+        self.model = model
+        self.start_ring_rate = start_ring_rate
+        self.end_ring_rate = end_ring_rate
+        self.frame_duration_ms = frame_duration_ms
+        self.padding_duration_ms = padding_duration_ms
+        self.silence_duration_threshold = silence_duration_threshold
+        self.sample_rate = sample_rate
+
+        # frames
+        self.frame_length = int(sample_rate * (frame_duration_ms / 1000.0))
+        self.frame_timestamp = 0.0
+        self.signal_cache = None
+
+        # segments
+        self.num_padding_frames = int(padding_duration_ms / frame_duration_ms)
+        self.ring_buffer = collections.deque(maxlen=self.num_padding_frames)
+        self.triggered = False
+        self.voiced_frames: List[Frame] = list()
+        self.segments = list()
+
+        # vad segments
+        self.is_first_segment = True
+        self.timestamp_start = 0.0
+        self.timestamp_end = 0.0
+
+    def signal_to_frames(self, signal: np.ndarray):
+        frames = list()
+
+        l = len(signal)
+
+        duration = float(self.frame_length) / self.sample_rate
+
+        for offset in range(0, l, self.frame_length):
+            sub_signal = signal[offset:offset+self.frame_length]
+
+            frame = Frame(sub_signal, self.frame_timestamp, duration)
+            self.frame_timestamp += duration
+
+            frames.append(frame)
+        return frames
+
+    def segments_generator(self, signal: np.ndarray):
+        # signal rounding
+        if self.signal_cache is not None:
+            signal = np.concatenate([self.signal_cache, signal])
+
+        rest = len(signal) % self.frame_length
+
+        if rest == 0:
+            self.signal_cache = None
+            signal_ = signal
+        else:
+            self.signal_cache = signal[-rest:]
+            signal_ = signal[:-rest]
+
+        # frames
+        frames = self.signal_to_frames(signal_)
+
+        for frame in frames:
+            speech_prob = self.model.predict(frame.signal)
+
+            if not self.triggered:
+                self.ring_buffer.append((frame, speech_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 f, _ in self.ring_buffer:
+                        self.voiced_frames.append(f)
+                    self.ring_buffer.clear()
+            else:
+                self.voiced_frames.append(frame)
+                self.ring_buffer.append((frame, speech_prob))
+                num_voiced = sum([p for _, p in self.ring_buffer])
+
+                if num_voiced < self.end_ring_rate * self.ring_buffer.maxlen:
+                    self.triggered = False
+                    segment = [
+                        np.concatenate([f.signal for f in self.voiced_frames]),
+                        self.voiced_frames[0].timestamp,
+                        self.voiced_frames[-1].timestamp,
+                    ]
+                    yield segment
+                    self.ring_buffer.clear()
+                    self.voiced_frames = []
+
+    def vad_segments_generator(self, segments_generator):
+        segments = list(segments_generator)
+
+        for i, segment in enumerate(segments):
+            start = round(segment[1], 4)
+            end = round(segment[2], 4)
+
+            if self.is_first_segment:
+                self.timestamp_start = start
+                self.timestamp_end = end
+                self.is_first_segment = False
+                continue
+
+            if self.timestamp_start:
+                sil_duration = start - self.timestamp_end
+                if sil_duration > self.silence_duration_threshold:
+                    vad_segment = [self.timestamp_start, self.timestamp_end]
+                    yield vad_segment
+
+                    self.timestamp_start = start
+                    self.timestamp_end = end
+                else:
+                    self.timestamp_end = end
+
+    def vad(self, signal: np.ndarray) -> List[list]:
+        segments = self.segments_generator(signal)
+        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 = [
+                np.concatenate([f.signal for f in self.voiced_frames]),
+                self.voiced_frames[0].timestamp,
+                self.voiced_frames[-1].timestamp
+            ]
+            segments = [segment]
+
+        # last vad segments
+        vad_segments = self.vad_segments_generator(segments)
+        vad_segments = list(vad_segments)
+
+        vad_segments = vad_segments + [[self.timestamp_start, self.timestamp_end]]
+        return vad_segments
+
+
+def make_visualization(signal: np.ndarray, sample_rate: int, vad_segments: list):
+    time = np.arange(0, len(signal)) / sample_rate
+    plt.figure(figsize=(12, 5))
+    plt.plot(time, signal / 32768, color='b')
+    for start, end in vad_segments:
+        plt.axvline(x=start, ymin=0.25, ymax=0.75, color='g', linestyle='--', label='开始端点')  # 标记开始端点
+        plt.axvline(x=end, ymin=0.25, ymax=0.75, color='r', linestyle='--', label='结束端点')  # 标记结束端点
+
+    plt.show()
+    return
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--wav_file",
+        default=(project_path / "data/early_media/3300999628164249998.wav").as_posix(),
+        type=str,
+    )
+    parser.add_argument(
+        "--model_name",
+        default=(project_path / "pretrained_models/silero_vad/silero_vad.jit").as_posix(),
+        type=str,
+    )
+    parser.add_argument(
+        "--frame_duration_ms",
+        default=30,
+        type=int,
+    )
+    parser.add_argument(
+        "--silence_duration_threshold",
+        default=0.3,
+        type=float,
+        help="minimum silence duration, in seconds."
+    )
+    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
+
+    # model = SileroVoiceClassifier(model_name=args.model_name, sample_rate=SAMPLE_RATE)
+    model = WebRTCVoiceClassifier(agg=1, sample_rate=SAMPLE_RATE)
+
+    vad = Vad(model=model,
+              start_ring_rate=0.9,
+              end_ring_rate=0.1,
+              frame_duration_ms=30,
+              padding_duration_ms=300,
+              silence_duration_threshold=0.30,
+              sample_rate=SAMPLE_RATE,
+              )
+    print(vad)
+
+    vad_segments = list()
+
+    segments = vad.vad(signal)
+    vad_segments += segments
+    for segment in segments:
+        print(segment)
+
+    # last vad segment
+    segments = vad.last_vad_segments()
+    vad_segments += segments
+    for segment in segments:
+        print(segment)
+
+    # plot
+    make_visualization(signal, SAMPLE_RATE, vad_segments)
+    return
+
+
+if __name__ == '__main__':
+    main()

toolbox/webrtcvad/vad.py

@@ -168,7 +168,7 @@ def get_args():
     parser = argparse.ArgumentParser()
     parser.add_argument(
         "--wav_file",
-        default=(project_path / "data/3300999628164249998.wav").as_posix(),
+        default=(project_path / "data/early_media/3300999628164249998.wav").as_posix(),
         type=str,
     )
     parser.add_argument(

webrtcvad_examples.json

@@ -1,8 +0,0 @@
-[
-    [
-        "data/early_media/3300999628164249998.wav"
-    ],
-    [
-        "data/early_media/3300999628164852605.wav"
-    ]
-]