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)