Initial commit

Check in a shit ton of code.

Most of the audio processing logic in `app.py` is lifted/ported from
github.com/yum_food/TaSTT. I made some adjustments to make it work
better (removing normalization, adding volume filters) and also increase
fidelity.

1 files changed, 417 insertions, 0 deletions

diff --git a/app.py b/app.py
new file mode 100644
index 0000000..3cb3816
--- /dev/null
+++ b/app.py
@@ -0,0 +1,417 @@
+from datetime import datetime
+from pydub import AudioSegment
+
+import math
+import numpy as np
+import os
+import pyaudio
+import sys
+import time
+import typing
+import vad
+import wave
+
+class AudioStream():
+    FORMAT = pyaudio.paInt16
+    # Size of each frame (audio sample), in bytes. If you change FORMAT, make
+    # sure this stays up to date!
+    FRAME_SZ = 2
+    # Frames per second.
+    FPS = 16000
+    CHANNELS = 1
+    def __init__(self):
+        pass
+
+    def getSamples(self) -> bytes:
+        raise NotImplementedError("getSamples is not implemented!")
+
+class MicStream(AudioStream):
+    CHUNK_SZ = 1024
+
+    def __init__(self, which_mic: str, fps: int = AudioStream.FPS):
+        self.p = pyaudio.PyAudio()
+        self.stream = None
+        self.sample_rate = None
+        # Each time pyaudio gives us audio data, it's in the form of a chunk of
+        # samples. We keep these in a list to keep the audio callback as light
+        # as possible. Whenever downstream layers want data, we collapse the
+        # list into a single array of data (a bytes object).
+        self.chunks = []
+        # If set, incoming frames are simply discarded.
+        self.paused = False
+        self.fps = fps
+
+        print(f"Finding mic {which_mic}", file=sys.stderr)
+        self.dumpMicDevices()
+
+        got_match = False
+        device_index = -1
+        focusrite_str = "Focusrite"
+        index_str = "Digital Audio Interface"
+        if which_mic == "index":
+            target_str = index_str
+        elif which_mic == "focusrite":
+            target_str = focusrite_str
+        else:
+            print(f"Mic {which_mic} requested, treating it as a numerical " +
+                    "device ID", file=sys.stderr)
+            device_index = int(which_mic)
+            got_match = True
+        if not got_match:
+            info = self.p.get_host_api_info_by_index(0)
+            numdevices = info.get('deviceCount')
+            for i in range(0, numdevices):
+                if (self.p.get_device_info_by_host_api_device_index(0, i).get('maxInputChannels')) > 0:
+                    device_name = self.p.get_device_info_by_host_api_device_index(0, i).get('name')
+                    if target_str in device_name:
+                        print(f"Got matching mic: {device_name}",
+                                file=sys.stderr)
+                        device_index = i
+                        got_match = True
+                        break
+        if not got_match:
+            raise KeyError(f"Mic {which_mic} not found")
+
+        info = self.p.get_device_info_by_host_api_device_index(0, device_index)
+        print(f"Found mic {which_mic}: {info['name']}", file=sys.stderr)
+        self.sample_rate = int(info['defaultSampleRate'])
+        print(f"Mic sample rate: {self.sample_rate}", file=sys.stderr)
+
+        self.stream = self.p.open(
+                rate=self.sample_rate,
+                channels=self.CHANNELS,
+                format=self.FORMAT,
+                input=True,
+                frames_per_buffer=MicStream.CHUNK_SZ,
+                input_device_index=device_index,
+                stream_callback=self.onAudioFramesAvailable)
+
+        self.stream.start_stream()
+
+        AudioStream.__init__(self)
+
+    def pause(self, state: bool = True):
+        self.paused = state
+
+    def dumpMicDevices(self):
+        info = self.p.get_host_api_info_by_index(0)
+        numdevices = info.get('deviceCount')
+
+        for i in range(0, numdevices):
+            if (self.p.get_device_info_by_host_api_device_index(0, i).get('maxInputChannels')) > 0:
+                device_name = self.p.get_device_info_by_host_api_device_index(0, i).get('name')
+                print("Input Device id ", i, " - ", device_name)
+
+    def onAudioFramesAvailable(self,
+            frames,
+            frame_count,
+            time_info,
+            status_flags):
+        if self.paused:
+            # Don't literally pause, just start returning silence. This allows
+            # the `min_segment_age_s` check to work while paused.
+            n_frames = int(frame_count * self.fps /
+                    float(self.sample_rate))
+            self.chunks.append(np.zeros(n_frames,
+                dtype=np.int16).tobytes())
+            return (frames, pyaudio.paContinue)
+
+        decimated = b''
+        # In pyaudio, a `frame` is a single sample of audio data.
+        frame_len = self.FRAME_SZ
+        next_frame = 0.0
+        # The mic probably has a higher sample rate than Whisper wants, so
+        # decrease the sample rate by dropping samples. Note that this
+        # algorithm only works if the mic's rate is higher than whisper's
+        # expected rate.
+        keep_every = float(self.sample_rate) / self.fps
+        for i in range(frame_count):
+            if i >= next_frame:
+                decimated += frames[i*frame_len:(i+1)*frame_len]
+                next_frame += keep_every
+        self.chunks.append(decimated)
+
+        return (frames, pyaudio.paContinue)
+
+    # Get audio data and the corresponding timestamp.
+    def getSamples(self) -> bytes:
+        chunks = self.chunks
+        self.chunks = []
+        result = b''.join(chunks)
+        return result
+
+class AudioCollector:
+    def __init__(self, stream: AudioStream):
+        self.stream = stream
+        self.frames = b''
+        # Note: by design, this is the only spot where we anchor our timestamps
+        # against the real world. This is done to make it possible to profile
+        # test cases which read from disk (at much faster than real speed) in
+        # the same way that we profile real-time data.
+        self.wall_ts = time.time()
+
+    def getAudio(self) -> bytes:
+        frames = self.stream.getSamples()
+        if frames:
+            self.frames += frames
+        return self.frames
+
+    def dropAudioPrefix(self, dur_s: float) -> bytes:
+        n_bytes = int(dur_s * self.stream.fps) * self.stream.FRAME_SZ
+        n_bytes = min(n_bytes, len(self.frames))
+        cut_portion = self.frames[:n_bytes]
+        self.frames = self.frames[n_bytes:]
+        self.wall_ts += float(n_bytes / self.stream.FRAME_SZ) / self.stream.fps
+        return cut_portion
+
+    def dropAudioPrefixByFrames(self, dur_frames: int) -> bytes:
+        n_bytes = dur_frames * self.stream.FRAME_SZ
+        n_bytes = min(n_bytes, len(self.frames))
+        cut_portion = self.frames[:n_bytes]
+        self.frames = self.frames[n_bytes:]
+        self.wall_ts += float(n_bytes / self.stream.FRAME_SZ) / self.stream.fps
+        return cut_portion
+
+    def keepLast(self, dur_s: float) -> bytes:
+        drop_len = max(0, self.duration() - dur_s)
+        return self.dropAudioPrefix(drop_len)
+
+    def dropAudio(self):
+        self.wall_ts += self.duration()
+        cut_portion = self.frames
+        self.frames = b''
+        return cut_portion
+
+    def duration(self):
+        return len(self.frames) / (self.stream.fps * self.stream.FRAME_SZ)
+
+    def begin(self):
+        return self.wall_ts
+
+    def now(self):
+        return self.begin() + self.duration()
+
+class AudioCollectorFilter:
+    def __init__(self, parent: AudioCollector):
+        self.parent = parent
+        self.stream = self.parent.stream
+
+    def getAudio(self) -> bytes:
+        return self.parent.getAudio()
+    def dropAudioPrefix(self, dur_s: float):
+        return self.parent.dropAudioPrefix(dur_s)
+    def dropAudioPrefixByFrames(self, dur_frames: int):
+        return self.parent.dropAudioPrefixByFrames(dur_frames)
+    def keepLast(self, dur_s):
+        return self.parent.keepLast(dur_s)
+    def dropAudio(self):
+        return self.parent.dropAudio()
+    def duration(self):
+        return self.parent.duration()
+    def begin(self):
+        return self.parent.begin()
+    def now(self):
+        return self.parent.now()
+
+class NormalizingAudioCollector(AudioCollectorFilter):
+    def __init__(self, parent: AudioCollector):
+        AudioCollectorFilter.__init__(self, parent)
+
+    def getAudio(self) -> bytes:
+        audio = self.parent.getAudio()
+
+        audio = AudioSegment(audio, sample_width=self.stream.FRAME_SZ,
+                frame_rate=self.stream.fps, channels=self.stream.CHANNELS)
+        audio = audio.normalize()
+
+        frames = np.array(audio.get_array_of_samples())
+        frames = np.int16(frames).tobytes()
+
+        return frames
+
+class CompressingAudioCollector(AudioCollectorFilter):
+    def __init__(self, parent: AudioCollector):
+        AudioCollectorFilter.__init__(self, parent)
+
+    def getAudio(self) -> bytes:
+        audio = self.parent.getAudio()
+
+        audio = AudioSegment(audio,
+                sample_width=self.stream.FRAME_SZ,
+                frame_rate=self.stream.fps,
+                channels=self.stream.CHANNELS)
+        # subtle compression has a slight positive effect on my benchmark
+        audio = audio.compress_dynamic_range(threshold=-10, ratio=2.0)
+
+        frames = np.array(audio.get_array_of_samples())
+        frames = np.int16(frames).tobytes()
+
+        return frames
+
+class AudioSegmenter:
+    def __init__(self,
+            min_silence_ms=250,
+            max_speech_s=5,
+            stream: AudioStream = None):
+        self.vad_options = vad.VadOptions(
+                min_silence_duration_ms=min_silence_ms,
+                max_speech_duration_s=max_speech_s)
+        self.stream = stream
+        pass
+
+    def segmentAudio(self, audio: bytes):
+        audio = np.frombuffer(audio,
+                dtype=np.int16).flatten().astype(np.float32) / 32768.0
+        return vad.get_speech_timestamps(audio, vad_options=self.vad_options)
+
+    # Returns the stable cutoff (if any) and whether there are any segments.
+    def getStableCutoff(self, audio: bytes) -> typing.Tuple[int, bool]:
+        min_delta_frames = int((self.vad_options.min_silence_duration_ms *
+                self.stream.fps) / 1000)
+        cutoff = None
+
+        last_end = None
+        segments = self.segmentAudio(audio)
+
+        for i in range(len(segments)):
+            s = segments[i]
+            #print(f"s: {s}")
+            #print(f"last_end: {last_end}")
+
+            if last_end:
+                delta_frames = s['start'] - last_end
+                #print(f"delta frames: {delta_frames}")
+                if delta_frames > min_delta_frames:
+                    cutoff = s['start']
+            else:
+                last_end = s['end']
+
+            if i == len(segments) - 1:
+                now = int(len(audio) / self.stream.FRAME_SZ)
+                delta_frames = now - s['end']
+                if delta_frames > min_delta_frames:
+                    cutoff = now - int(min_delta_frames / 2)
+
+        return (cutoff, len(segments) > 0)
+
+def install_in_venv(pkgs: typing.List[str]) -> bool:
+    pkgs_str = " ".join(pkgs)
+    print(f"Installing {pkgs_str}")
+    pip_proc = subprocess.Popen(
+            f"Resources/Python/python.exe -m pip install {pkgs_str} --no-warn-script-location".split(),
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE)
+    pip_stdout, pip_stderr = pip_proc.communicate()
+    pip_stdout = pip_stdout.decode("utf-8")
+    pip_stderr = pip_stderr.decode("utf-8")
+    print(pip_stdout, file=sys.stderr)
+    print(pip_stderr, file=sys.stderr)
+    if pip_proc.returncode != 0:
+        print(f"`pip install {pkgs_str}` exited with {pip_proc.returncode}",
+                file=sys.stderr)
+        return False
+    return True
+
+def saveAudio(audio: bytes, path: str, stream: AudioStream):
+    with wave.open(path, 'wb') as wf:
+        print(f"Saving audio to {path}", file=sys.stderr)
+        wf.setnchannels(stream.CHANNELS)
+        wf.setsampwidth(stream.FRAME_SZ)
+        wf.setframerate(stream.fps)
+        wf.writeframes(audio)
+
+def concatenate_wav_files(output_path):
+    # List all .wav files in the CWD
+    wav_files = [f for f in os.listdir('.') if f.endswith('.wav')]
+
+    # Initialize parameters for wave file
+    params = None
+
+    # Open the output file
+    with wave.open(output_path, 'wb') as output_wav:
+        for wav_file in wav_files:
+            print(f"Processing {wav_file}")
+            with wave.open(wav_file, 'rb') as input_wav:
+                # Check if parameters are the same for each file
+                if params is None:
+                    params = input_wav.getparams()
+                    output_wav.setparams(params)
+
+                # Read and write frames
+                frames = input_wav.readframes(input_wav.getnframes())
+                output_wav.writeframes(frames)
+
+if __name__ == "__main__":
+    abspath = os.path.abspath(__file__)
+    dname = os.path.dirname(abspath)
+    os.chdir(dname)
+    print(f"Set cwd to {os.getcwd()}", file=sys.stderr)
+
+    concatenate_wav_files("concatenated.wav")
+    sys.exit(0)
+
+    stream = MicStream("index")
+    stream_hd = MicStream("index", fps=44100)
+
+    collector = AudioCollector(stream)
+    #collector = NormalizingAudioCollector(collector)
+    collector = CompressingAudioCollector(collector)
+
+    collector_hd = AudioCollector(stream_hd)
+    #collector_hd = NormalizingAudioCollector(collector_hd)
+    collector_hd = CompressingAudioCollector(collector_hd)
+
+    min_silence_ms = 1000
+    max_speech_s = 30
+    segmenter = AudioSegmenter(
+            min_silence_ms=min_silence_ms,
+            max_speech_s=max_speech_s,
+            stream=stream)
+
+    while True:
+        audio = collector.getAudio()
+        collector_hd.getAudio()
+        stable_cutoff, has_audio = segmenter.getStableCutoff(audio)
+
+        #print(f"has audio: {has_audio}")
+        #print(f"stable cutoff: {stable_cutoff}")
+
+        if has_audio and stable_cutoff:
+            commit_audio = collector.dropAudioPrefixByFrames(stable_cutoff)
+            print(f"stable cutoff: {stable_cutoff}")
+            hd_cutoff = int(math.floor(stable_cutoff * stream_hd.fps /
+                stream.fps))
+            print(f"hd cutoff: {hd_cutoff}")
+            commit_audio_hd = collector_hd.dropAudioPrefixByFrames(hd_cutoff)
+            print(f"hd audio len: {len(commit_audio_hd)}")
+
+            # Calculate naive measure of volume
+            audio_v = AudioSegment(commit_audio_hd,
+                    sample_width=stream_hd.FRAME_SZ,
+                    frame_rate=stream_hd.fps,
+                    channels=stream_hd.CHANNELS)
+            audio_v = np.array(audio_v.get_array_of_samples())
+            audio_v = np.int16(audio_v)
+            audio_v = np.sqrt(np.mean(np.square(audio_v)))
+            audio_v /= np.sqrt(len(commit_audio_hd) / stream_hd.FRAME_SZ)
+            audio_v = math.log(audio_v, 10)
+            print(f"volume: {audio_v}")
+            # cutoff is a fine-tuned value based on volumes seen while in vr
+            # (index mic)
+            if audio_v < -1.3 or audio_v > -0.8:
+                # Discard sample
+                print("Discarding too-quiet/too-loud segment")
+                collector.keepLast(1.0)
+                collector_hd.keepLast(1.0)
+                continue
+
+
+            ts = datetime.fromtimestamp(time.time())
+            filename = str(ts.strftime('%Y_%m_%d__%H-%M-%S')) + ".wav"
+            saveAudio(commit_audio_hd, filename, stream_hd)
+
+        if not has_audio:
+            #print("VAD detects no audio, skip transcription", file=sys.stderr)
+            collector.keepLast(1.0)
+            collector_hd.keepLast(1.0)
+