path: root/app.py

from datetime import datetime
from pydub import AudioSegment

import gradio as gr
import math
import numpy as np
import os
import pyaudio
import subprocess
import sys
import time
import typing
import vad
import wave

class Logger:
    def __init__(self, filename):
        self.terminal = sys.stdout
        self.log = open(filename, "w")

    def write(self, message):
        self.terminal.write(message)
        self.log.write(message)

    def flush(self):
        self.terminal.flush()
        self.log.flush()

    def isatty(self):
        return False

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)

        got_match = False
        device_index = -1
        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 which_mic 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 getMicDevices() -> typing.List[str]:
        p = pyaudio.PyAudio()
        info = p.get_host_api_info_by_index(0)
        numdevices = info.get('deviceCount')

        result = []
        for i in range(0, numdevices):
            if (p.get_device_info_by_host_api_device_index(0, i).get('maxInputChannels')) > 0:
                device_name = p.get_device_info_by_host_api_device_index(0, i).get('name')
                result.append(device_name)
        return result

    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 DiskStream(AudioStream):
    def __init__(self, path: str):
        fmt = None
        if path.endswith(".mp3"):
            fmt = "mp3"
        elif path.endswith(".wav"):
            fmt = "wav"
        else:
            raise NotImplementedError(f"Requested file type {path} " + \
                    "is not supported")
        print(f"Loading audio data", file=sys.stderr)
        audio = AudioSegment.from_file(path, format=fmt)
        audio = audio.set_channels(1)
        audio = audio.set_frame_rate(16000)
        frames = np.array(audio.get_array_of_samples())
        frames = np.int16(frames).tobytes()

        self.frames = frames
        self.fps = 16000

    def getSamples(self) -> bytes:
        frames = self.frames
        self.frames = b''
        return frames

        if len(frames) < nframes:
            frames += np.zeros(nframes - len(frames), dtype=np.int16).tobytes()

        return frames

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 concatenateWavFiles(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:
            if os.path.abspath(wav_file) == os.path.abspath(output_path):
                print(f"Skip adding output file ({wav_file}) to itself")
                continue
            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)

class AppControl:
    run = True
app_ctrl = AppControl()

def recordAudio(
        mic_device: str,
        min_volume: float = -1.3,
        max_volume: float = -0.8
        ):
    app_ctrl.run = True

    stream = MicStream(mic_device)
    stream_hd = MicStream(mic_device, 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 = 250
    max_speech_s = 30
    segmenter = AudioSegmenter(
            min_silence_ms=min_silence_ms,
            max_speech_s=max_speech_s,
            stream=stream)

    while app_ctrl.run:
        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 < min_volume or audio_v > max_volume:
                # 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)
    print("Stopped recording")

class Segment:
    def __init__(self,
            transcript: str,
            start_ts: float,
            end_ts: float,
            wall_ts: float,
            avg_logprob: float,
            no_speech_prob: float,
            compression_ratio: float):
        self.transcript = transcript
        # start_ts, end_ts are timestamps in seconds relative to `wall_ts`.
        self.start_ts = start_ts
        self.end_ts = end_ts
        # wall_ts is the time.time() at which the oldest audio sample leading
        # to this transcript was collected.
        self.wall_ts = wall_ts
        self.avg_logprob = avg_logprob
        self.no_speech_prob = no_speech_prob
        self.compression_ratio = compression_ratio

    def __str__(self):
        ts = f"(ts: {self.start_ts}-{self.end_ts}) "

        wall_ts_start = datetime.utcfromtimestamp(self.start_ts + self.wall_ts).strftime('%H:%M:%S')
        wall_ts_end = datetime.utcfromtimestamp(self.end_ts + self.wall_ts).strftime('%H:%M:%S')
        wall_ts = f"(wall ts: {wall_ts_start}-{wall_ts_end}) "

        no_speech = f"(no_speech: {self.no_speech_prob}) "
        avg_logprob = f"(avg_logprob: {self.avg_logprob}) "
        return f"{self.transcript} " + ts + wall_ts + no_speech + avg_logprob

def pipInstall(pkgs: typing.List[str]) -> bool:
    pkgs_str = " ".join(pkgs)
    print(f"Installing {pkgs_str}")
    env = os.environ.copy()
    # cwd is set at top of __main__. We set PATH to ensure that installed
    # Python packages have access to any binaries that come with them.
    env["PATH"] = os.getcwd() + "/Python/Scripts;" + env['PATH']
    pip_proc = subprocess.Popen(
            f"./Python/python.exe -m pip install {pkgs_str} --no-warn-script-location".split(),
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            env=env)
    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

class Whisper:
    def __init__(self,
            collector: AudioCollector):
        self.collector = collector

        import torch
        from transformers import pipeline

        whisper_model = "openai/whisper-large-v2"
        print(f"Loading pipeline for {whisper_model}...")
        self.pipe = pipeline(
                "automatic-speech-recognition",
                model="distil-whisper/distil-large-v2",
                torch_dtype=torch.float16,
                device="cuda",
                )
        print(f"Done.")

    def transcribe(self, frames: bytes = None) -> typing.List[Segment]:
        if frames is None:
            frames = self.collector.getAudio()
        # Convert from signed 16-bit int [-32768, 32767] to signed 32-bit float on
        # [-1, 1].
        audio = np.frombuffer(frames,
                dtype=np.int16).flatten().astype(np.float32) / 32768.0

        t0 = time.time()
        res = self.pipe(
                audio,
                chunk_length_s=30,
                batch_size=1)

        result = [Segment(res["text"],
            0,
            0,
            self.collector.begin(),
            0,
            0,
            0)]

        t1 = time.time()
        print(f"Transcription latency (s): {t1 - t0}: {result[0].transcript}")
        return result

def getOutput() -> str:
    sys.stdout.flush()
    with open("output.log", "r") as f:
        return f.read()

def stopApp():
    print("Requesting app stop")
    app_ctrl.run = False

def transcribeAudio(concatenated_path: str):
    # Step 1: Install Whisper requirements
    print("Installing Whisper dependencies, this will take several minutes")
    with open("whisper_requirements.txt", "r") as file:
        requirements = file.read().splitlines()
    if not pipInstall(requirements):
        return

    # Step 2: Iterate over .wav files in the current working directory
    print("Loading Whisper model, this will take several minutes")
    whisper = Whisper(None)
    for wav_file in os.listdir('.'):
        if wav_file.endswith('.wav'):
            if wav_file.endswith(os.path.basename(concatenated_path)):
                print("Skipping concatenated file")
                continue

            # Step 3: Transcription pipeline
            # TODO parameterize high fidelity framerate
            print(f"Transcribing {wav_file}")
            disk_stream = DiskStream(wav_file)
            collector = CompressingAudioCollector(AudioCollector(disk_stream))
            whisper.collector = collector

            transcript_filename = wav_file.replace('.wav', '.txt')
            if os.path.exists(transcript_filename):
                print(f"Transcript already exists - skipping")
                continue

            # Transcribe the audio
            segments = whisper.transcribe()

            # Step 4: Save transcriptions
            with open(transcript_filename, 'w') as txt_file:
                for segment in segments:
                    txt_file.write(segment.transcript + '\n')
            print(f"Transcript generated at {transcript_filename}")

if __name__ == "__main__":
    abspath = os.path.abspath(__file__)
    dname = os.path.dirname(abspath)
    os.chdir(dname)

    sys.stdout = Logger("output.log")

    print(f"Set cwd to {os.getcwd()}", file=sys.stderr)

    with gr.Blocks() as demo:
        mic_choices = MicStream.getMicDevices()
        mic_device = gr.Dropdown(choices=mic_choices, label="Microphone")
        min_volume = gr.Number(label="Minimum volume", value=-1.3)
        max_volume = gr.Number(label="Maximum volume", value=-0.8)
        record_audio = gr.Button("Record audio")
        stop_recording = gr.Button("Stop recording")
        transcribe_audio = gr.Button("Transcribe audio")
        concatenated_path = gr.Text(label="Combined audio filename", value="combined.wav")
        min_length = gr.Number(label="Minimum length (seconds)", value=3.0)
        concatenate_audio = gr.Button("Combine audio files")

        dbg_output = gr.Text(label="Output")

        record_audio.click(recordAudio, [mic_device, min_volume, max_volume],
                dbg_output)
        stop_recording.click(stopApp, [], dbg_output)

        transcribe_audio.click(transcribeAudio, [concatenated_path], dbg_output)

        concatenate_audio.click(concatenateWavFiles, [concatenated_path],
                dbg_output)

        demo.load(getOutput, None, dbg_output, every=0.5)
    demo.launch()
    sys.exit(0)

    concatenateWavFiles("concatenated.wav")
    sys.exit(0)