localsub/speech_segmenter.py at master · SmartestWashingMachine/localsub · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
from dataclasses import dataclass
from typing import List
from cacher import Cache
from tqdm import tqdm
import sys, os

binary_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "../include"))
print(f"Added TEN VAD binary: {binary_path}")

from ten_vad import TenVad
import librosa
import numpy as np

@dataclass
class SpeechSegment():
    start: float
    end: float

class SpeechSegmenter():
    def __init__(self):
        self.cache = Cache("speechsegmenter")

    def detect(self, audio_file_path: str):
        cached = None #self.cache.load(audio_file_path)
        if cached:
            return cached

        results = self.process(audio_file_path)

        self.cache.save(results, audio_file_path)
        return results

    def process(self, audio_file_path: str) -> List[SpeechSegment]:
        return []

# TODO: Better filter. Not sure about the "buffer = []" on a single silence frame.
class TenVadSpeechSegmenter(SpeechSegmenter):
    def process(self, audio_file_path):
        sampling_rate = 16000
        audio_data, _ = librosa.load(audio_file_path, sr=sampling_rate)

        audio_data = np.clip(audio_data * 32768, -32768, 32767).astype(np.int16)

        # ref: https://github.com/TEN-framework/ten-vad/blob/main/examples/test.py
        hop_size = 256  # 16 ms per frame
        threshold = 0.5
        num_frames = audio_data.shape[0] // hop_size

        vad = TenVad(hop_size, threshold)

        is_speaking = False
        start_i = 0
        buffer = []
        valid_frames = 3 # Wait for this many speech frames to start a speech segment.
        patience_frames = 3 # Wait for this many silence frames to consider a speech segment done.

        segments: List[SpeechSegment] = []

        for i in tqdm(range(num_frames), desc="VAD"):
            frame = audio_data[i * hop_size:(i + 1) * hop_size]
            out_probability, out_flag = vad.process(frame)

            has_speech = out_flag == 1

            #from time import sleep
            #sleep(1)
            #print("[%d] %0.6f, %d" % (i, out_probability, out_flag))

            if is_speaking:
                if not has_speech:
                    buffer.append(i)

                    # Waited long enough - there's no more speech. End this segment.
                    if len(buffer) >= patience_frames:
                        segments.append(SpeechSegment(start=start_i, end=i))

                        is_speaking = False
                        buffer = []
                else:
                    # Speech was still detected - someone is still speaking. Reset the patience threshold.
                    buffer = []
            else:
                if has_speech:
                    buffer.append(i)

                    if len(buffer) >= valid_frames:
                        # Enough speech frames were detected - starting a new segment and listening for its end...
                        is_speaking = True
                        buffer = []
                        start_i = i
                else:
                    # Not currently in a speech segment and silence was detected - reset the window.
                    buffer = []

        # Detect speech at end of audio; or alternatively: Detect speech that spans the entire audio.
        if is_speaking:
            segments.append(SpeechSegment(start=start_i, end=i + 1)) # i + 1 as last frame probably had speech too.

        # From frame indices to seconds.
        segments = [SpeechSegment(start=s.start * hop_size / sampling_rate, end=s.end * hop_size / sampling_rate) for s in segments]

        joined: List[SpeechSegment] = []
        join_second_threshold = 0.0 # Segments closer than this many seconds are joined together.
        for seg in segments:
            if len(joined) == 0:
                joined.append(seg)
                continue

            if seg.start - joined[-1].end <= join_second_threshold:
                joined[-1].end = seg.end
            else:
                joined.append(seg)

        return joined