WhisperLive/client.py

import io
import os
import argparse
import wave

import numpy as np
import scipy
import ffmpeg
import torch
import pyaudio
import threading
import textwrap
import json
import torchaudio
import websocket


CHUNK = 1024
FORMAT = pyaudio.paInt16
CHANNELS = 1
RATE = 16000
RECORD_SECONDS = 60000



def on_message(ws, message):
    message = json.loads(message)
    text = []
    if len(message):
        for seg in message:
            if len(text):
                if text[-1] != seg["text"]:
                    text.append(seg["text"])
            else:
                text.append(seg["text"])
    if len(text) > 3:
        text = text[-3:]
    wrapper = textwrap.TextWrapper(width=60)
    word_list = wrapper.wrap(text="".join(text))
    # Print each line.
    os.system('clear')
    for element in word_list:
        print(element)

def on_error(ws, error):
    print(error)

def on_close(ws, close_status_code, close_msg):
    print("### websocket connection closed ###")

def on_open(ws):
    print("Opened connection")
    

class Client:
    def __init__(self, host=None, port=None):
        self.timestamp_offset = 0.0
        self.audio_bytes = None
        self.p = pyaudio.PyAudio()
        self.stream = self.p.open(format=FORMAT,
                        channels=CHANNELS,
                        rate=RATE,
                        input=True,
                        frames_per_buffer=CHUNK)
        print(self.p.get_sample_size(FORMAT))

        # create websocket connection
        if host is not None and port is not None:
            socket_url = f"ws://{host}:{port}"    
            self.client_socket = websocket.WebSocketApp(socket_url,
                                  on_open=on_open,
                                  on_message=on_message,
                                  on_error=on_error,
                                  on_close=on_close)
        else:
            print("No host or port specified.")
            return

        # start websocket client in a thread
        self.ws_thread = threading.Thread(target=self.client_socket.run_forever)
        self.ws_thread.setDaemon(True)
        self.ws_thread.start()

        # voice activity detection model
        self.vad_model, _ = torch.hub.load(repo_or_dir='snakers4/silero-vad',
                              model='silero_vad',
                              force_reload=True,
                              onnx=True)
        self.window_size = 1024
        self.vad_threshold = 0.4

        self.frames = b""
        print("* recording")

    def send_packet_to_server(self, message):
        try:
            self.client_socket.send(message, websocket.ABNF.OPCODE_BINARY)
        except Exception as e:
            print(e)

    @staticmethod
    def bytes_to_audio_tensor(audio_bytes):
        bytes_io = io.BytesIO()
        raw_data = np.frombuffer(
            buffer=audio_bytes, dtype=np.int16
        )
        scipy.io.wavfile.write(bytes_io, RATE, raw_data)
        audio, _ = torchaudio.load(bytes_io)
        return audio.squeeze(0), raw_data.astype(np.float32) / 32768.0
    
    def play_file(self, filename):
        # read audio and create pyaudio stream
        self.wf = wave.open(filename, 'rb')
        self.stream = self.p.open(format=self.p.get_format_from_width(self.wf.getsampwidth()),
                channels=self.wf.getnchannels(),
                rate=self.wf.getframerate(),
                input=True,
                output=True,
                frames_per_buffer=CHUNK)
        try:
            while True:
                data = self.wf.readframes(CHUNK)
                if data==b'': break

                # voice activity detection
                chunk_tensor, audio_array = Client.bytes_to_audio_tensor(data)
                try:
                    speech_prob = self.vad_model(chunk_tensor, RATE).item()
                except ValueError:
                    break   # input audio chunk is too short
                if speech_prob > self.vad_threshold:
                    self.send_packet_to_server(audio_array.tobytes())
                self.stream.write(data)

            self.wf.close()
            self.stream.close()

        except KeyboardInterrupt:
            print("Keyboard interrupt.")


    def get_client_socket(self):
        return self.client_socket
    
    def write_audio_frames_to_file(self, frames, file_name):
        wf = wave.open(file_name, 'wb')
        wf.setnchannels(CHANNELS)
        wf.setsampwidth(2)
        wf.setframerate(RATE)
        wf.writeframes(frames)
        wf.close()

    def record(self, out_file="output_recording.wav"):
        n_audio_file = 0
        # create dir for saving audio chunks
        if not os.path.exists("chunks"):
            os.makedirs("chunks", exist_ok=True)
        try:
            for _ in range(0, int(RATE / CHUNK * RECORD_SECONDS)):
                data = self.stream.read(CHUNK)
                self.frames += data

                # voice activity detection
                chunk_tensor , audio_array = Client.bytes_to_audio_tensor(data)
                
                speech_prob = self.vad_model(chunk_tensor, RATE).item()
                if speech_prob > self.vad_threshold:
                    self.send_packet_to_server(audio_array.tobytes())

                # save frames if more than a minute
                if len(self.frames) > 60*RATE:
                    t = threading.Thread(
                        target=self.write_audio_frames_to_file,
                        args=(self.frames[:], f"chunks/{n_audio_file}.wav", )
                    )
                    t.start()
                    n_audio_file += 1
                    self.frames = b""

        except KeyboardInterrupt:
            if len(self.frames):
                self.write_audio_frames_to_file(
                    self.frames[:], f"chunks/{n_audio_file}.wav")
                n_audio_file += 1
            self.stream.stop_stream()
            self.stream.close()
            self.p.terminate()


            # combine all the audio files
            self.write_output_recording(n_audio_file, out_file)
    
    def write_output_recording(self, n_audio_file, out_file):
        input_files = [f"chunks/{i}.wav" for i in range(n_audio_file) if os.path.exists(f"chunks/{i}.wav")]
        wf = wave.open(out_file, 'wb')
        wf.setnchannels(CHANNELS)
        wf.setsampwidth(2)
        wf.setframerate(RATE)
        for in_file in input_files:
            w = wave.open(in_file, 'rb')
            while True:
                data = w.readframes(CHUNK)
                if data==b'': break
                wf.writeframes(data)
            w.close()
            # remove this file
            os.remove(in_file)
        wf.close()


def resample(file: str, sr: int = 16000):
    """
    # https://github.com/openai/whisper/blob/7858aa9c08d98f75575035ecd6481f462d66ca27/whisper/audio.py#L22
    Open an audio file and read as mono waveform, resampling as necessary,
    save the resampled audio
    Parameters
    ----------
    file: str
        The audio file to open
    sr: int
        The sample rate to resample the audio if necessary
    """
    try:
        # This launches a subprocess to decode audio while down-mixing and resampling as necessary.
        # Requires the ffmpeg CLI and `ffmpeg-python` package to be installed.
        out, _ = (
            ffmpeg.input(file, threads=0)
            .output("-", format="s16le", acodec="pcm_s16le", ac=1, ar=sr)
            .run(cmd=["ffmpeg", "-nostdin"], capture_stdout=True, capture_stderr=True)
        )
    except ffmpeg.Error as e:
        raise RuntimeError(f"Failed to load audio: {e.stderr.decode()}") from e
    np_buffer = np.frombuffer(out, dtype=np.int16)

    resampled_file = f"{file.split('.')[0]}_resampled.wav"
    scipy.io.wavfile.write(resampled_file, sr, np_buffer.astype(np.int16))
    return resampled_file


if __name__=="__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--audio', type=str, help='audio file to transcribe')
    parser.add_argument('--host', default=None, type=str, help='websocket server address to connect to')
    parser.add_argument('--port', default=None, type=str, help='websocket server port to connect to')
    opt = parser.parse_args()
    c = Client(host=opt.host, port=opt.port)

    if opt.audio is not None:
        resampled_file = resample(opt.audio)
        c.play_file(resampled_file)
    else:
        c.record()