diff --git a/client.py b/client.py index 8c908bd..4dbdeab 100644 --- a/client.py +++ b/client.py @@ -1,113 +1,98 @@ -import io import os import argparse import wave -import uuid -import hashlib -import base64 -import time import numpy as np import scipy import ffmpeg -import torch -import socket, pickle, pyaudio, struct +import pyaudio import threading import textwrap import json -import torchaudio -from dataclasses import dataclass +import websocket + CHUNK = 1024 FORMAT = pyaudio.paInt16 CHANNELS = 1 RATE = 16000 RECORD_SECONDS = 60000 -all_segments = [] -@dataclass(frozen=True) -class Constants: - ACK = b"acknowledged" - RECORDING_OVER = b"audio_data_over" - RECEIVED_AUDIO_FILE = b"audio_file_sent" - RECEIVING_AUDIO_FILE = b"sending_audio_file" + +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, topic=None, host=None, port=None): + def __init__(self, host=None, port=None): self.timestamp_offset = 0.0 self.audio_bytes = None self.p = pyaudio.PyAudio() - self.payload_size = struct.calcsize("Q") self.stream = self.p.open(format=FORMAT, channels=CHANNELS, rate=RATE, input=True, frames_per_buffer=CHUNK) print(self.p.get_sample_size(FORMAT)) - self.client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - host_ip = 'localhost' if host is None else host - port = 5901 if port is None else port - socket_address = (host_ip, port) - self.client_socket.connect(socket_address) - print("CLIENT CONNECTED TO", socket_address) - - # 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 - - # subscribing to the correct topic - if topic is not None: - self.topic = topic + # 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: - self.topic = self.get_mac_address().decode() + 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() self.frames = b"" - data = b"" - while True: - while len(data) < self.payload_size: - packet = self.client_socket.recv(4*1024) #4K - if not packet: break - data+=packet - packed_msg_size = data[:self.payload_size] - data = data[self.payload_size:] - try: - msg_size = struct.unpack("Q",packed_msg_size)[0] - except struct.error: - break - while len(data) < msg_size: - data += self.client_socket.recv(4*1024) - frame_data = data[:msg_size] - frame_data = pickle.loads(frame_data) - if Constants.ACK in frame_data: - print("Server is ready. Sending audio ...") - break print("* recording") def send_packet_to_server(self, message): - a = pickle.dumps(message) - message = struct.pack("Q",len(a))+a - self.client_socket.sendall(message) - - def get_mac_address(self): - mac = hex(uuid.getnode()) - hasher = hashlib.sha1(mac.encode()) - return base64.urlsafe_b64encode(hasher.digest()[:5]) + 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() + def bytes_to_float_array(audio_bytes): 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) + return raw_data.astype(np.float32) / 32768.0 def play_file(self, filename): # read audio and create pyaudio stream @@ -123,35 +108,15 @@ class Client: data = self.wf.readframes(CHUNK) if data==b'': break - # voice activity detection - chunk_tensor = 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: - data_dict = { - "topic": self.topic, - "audio": data - } - self.send_packet_to_server(data_dict) + audio_array = Client.bytes_to_float_array(data) + self.send_packet_to_server(audio_array.tobytes()) self.stream.write(data) self.wf.close() self.stream.close() - # let the server know that we're done - data = Constants.RECORDING_OVER - self.send_packet_to_server(data) - with open("results.json", "w") as f: - json_dict = json.dumps(all_segments, indent=2) - f.write(json_dict) - except KeyboardInterrupt: - # write all segments to a file - with open("results.json", "w") as f: - json_dict = json.dumps(all_segments, indent=2) - f.write(json_dict) + print("Keyboard interrupt.") def get_client_socket(self): @@ -175,16 +140,9 @@ class Client: data = self.stream.read(CHUNK) self.frames += data - # voice activity detection - chunk_tensor = Client.bytes_to_audio_tensor(data) + audio_array = Client.bytes_to_float_array(data) - speech_prob = self.vad_model(chunk_tensor, RATE).item() - if speech_prob > self.vad_threshold: - data_dict = { - "topic": self.topic, - "audio": data - } - self.send_packet_to_server(data_dict) + self.send_packet_to_server(audio_array.tobytes()) # save frames if more than a minute if len(self.frames) > 60*RATE: @@ -205,16 +163,9 @@ class Client: self.stream.close() self.p.terminate() - # let the server know that we're done - data = Constants.RECORDING_OVER - self.send_packet_to_server(data) # combine all the audio files self.write_output_recording(n_audio_file, out_file) - # write all segments to a file - with open("results.json", "w") as f: - json_dict = json.dumps(all_segments, indent=2) - f.write(json_dict) 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")] @@ -234,41 +185,6 @@ class Client: wf.close() -def recieve_response(client_socket): - data = b"" - payload_size = struct.calcsize("Q") - - while True: - while len(data) < payload_size: - packet = client_socket.recv(4*1024) # 4K - if not packet: break - data+=packet - packed_msg_size = data[:payload_size] - data = data[payload_size:] - try: - msg_size = struct.unpack("Q",packed_msg_size)[0] - except struct.error: - break - while len(data) < msg_size: - data += client_socket.recv(4*1024) - frame_data = data[:msg_size] - data = data[msg_size:] - response = pickle.loads(frame_data) - - if response is not None and isinstance(response, dict): - os.system('clear') - text = response['text'] - segments = response['segments'] - if len(segments): - for seg in segments: - all_segments.append(seg) - wrapper = textwrap.TextWrapper(width=50) - word_list = wrapper.wrap(text=text) - # Print each line. - for element in word_list: - print(element) - - def resample(file: str, sr: int = 16000): """ # https://github.com/openai/whisper/blob/7858aa9c08d98f75575035ecd6481f462d66ca27/whisper/audio.py#L22 @@ -301,20 +217,13 @@ def resample(file: str, sr: int = 16000): if __name__=="__main__": parser = argparse.ArgumentParser() parser.add_argument('--audio', type=str, help='audio file to transcribe') - parser.add_argument('--topic', default=None, type=str, help='topic to subscribe for results') - parser.add_argument('--host', default=None, type=str, help='server address to connect to') - parser.add_argument('--port', default=None, type=str, help='server port to connect to') + 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(topic=opt.topic, host=opt.host, port=opt.port) - while True: - if c.get_client_socket() is not None: - break - client_socket = c.get_client_socket() - t2 = threading.Thread(target=recieve_response, args=(client_socket, )) - t2.start() + 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() - t2.join() diff --git a/requirements/client.txt b/requirements/client.txt index ce9121d..3befca2 100644 --- a/requirements/client.txt +++ b/requirements/client.txt @@ -1,5 +1,5 @@ PyAudio ffmpeg-python scipy -torch==1.12.1 -torchaudio==0.12.1 \ No newline at end of file +websocket-client +onnxruntime \ No newline at end of file diff --git a/server.py b/server.py index b3b1555..6d291ca 100644 --- a/server.py +++ b/server.py @@ -1,144 +1,81 @@ -import socket, pickle, struct, time, pyaudio +# import asyncio +import websockets +import pickle, struct, time, pyaudio import threading -import os +import os, json import wave import textwrap + +import logging +logging.basicConfig(level = logging.INFO) + from collections import deque from dataclasses import dataclass -import torch import numpy as np -import paho.mqtt.client as mqtt - +from websockets.sync import server +from websockets.sync.server import serve from transcriber import WhisperModel -def on_connect(mqttc, obj, flags, rc): - pass +clients = {} -def on_message(mqttc, obj, msg): - pass - -def on_publish(mqttc, obj, mid): - pass - -def on_subscribe(mqttc, obj, mid, granted_qos): - pass - -def on_log(mqttc, obj, level, string): - pass - - -@dataclass(frozen=True) -class Constants: - AUDIO_OVER = b"audio_data_over" - ACK = b"acknowledged" - SENDING_FILE = b"sending_audio_file" - FILE_SENT = b"audio_file_sent" +def recv_audio(websocket): + """ + Receive audio chunks from client in an infinite loop. + """ + global clients + client = ServeClient(websocket) + clients[websocket] = client + while True: + try: + frame_data = websocket.recv() + if isinstance(frame_data, str): + logging.info(frame_data) + continue + else: + frame_np = np.frombuffer(frame_data, np.float32) + clients[websocket].add_frames(frame_np) + + except Exception as e: + clients[websocket].cleanup() + clients.pop(websocket) + logging.info("Connection Closed.") + break class ServeClient: - CHUNK = 1024 - FORMAT = pyaudio.paInt16 - CHANNELS = 1 RATE = 16000 - def __init__(self, client_socket, device=None, verbose=True): + def __init__(self, websocket, topic=None, device=None): self.payload_size = struct.calcsize("Q") self.data = b"" self.frames = b"" - self.frames_np = None - self.transcriber = WhisperModel("medium.en", device="cuda", compute_type="float16") + self.transcriber = WhisperModel("small.en", compute_type="float16", local_files_only=False) self.timestamp_offset = 0.0 + self.frames_np = None self.frames_offset = 0.0 self.text = [] self.current_out = '' self.prev_out = '' self.t_start=None - self.client_socket = client_socket - self.verbose = verbose self.exit = False self.same_output_threshold = 0 self.show_prev_out_thresh = 5 # if pause(no output from whisper) show previous output for 5 seconds self.add_pause_thresh = 3 # add a blank to segment list as a pause(no speech) for 3 seconds + self.transcript = [] + self.send_last_n_segments = 10 # text formatting self.wrapper = textwrap.TextWrapper(width=50) self.pick_previous_segments = 2 # setup mqtt - self.topic = None - self.mqttc = mqtt.Client() - self.mqttc.on_message = on_message - self.mqttc.on_connect = on_connect - self.mqttc.on_publish = on_publish - self.mqttc.on_subscribe = on_subscribe - self.mqttc = mqtt.Client() - self.mqttc.connect("mqtt.kurg.org", 1883, 60) - self.mqttc.loop_start() - - # send response to client; server is ready - self.send_response_to_client(Constants.ACK) + self.topic = topic # threading - self.recv_thread = threading.Thread(target=self.recv_audio) + self.websocket = websocket self.trans_thread = threading.Thread(target=self.speech_to_text) - self.recv_thread.start() self.trans_thread.start() - - def recv_audio(self): - """ - Receive audio chunks from client in an infinite loop. - """ - if self.client_socket: - try: - while True: - while len(self.data) < self.payload_size: - packet = self.client_socket.recv(4*1024) # 4K - if not packet: break - self.data+=packet - - packed_msg_size = self.data[:self.payload_size] - self.data = self.data[self.payload_size:] - msg_size = struct.unpack("Q",packed_msg_size)[0] - - while len(self.data) < msg_size: - self.data += self.client_socket.recv(4*1024) - frame_data = self.data[:msg_size] - self.data = self.data[msg_size:] - frame_data = pickle.loads(frame_data) - if self.topic is None: - self.topic = frame_data["topic"] - - frame = frame_data["audio"] - - # client says audio over - if Constants.AUDIO_OVER in frame: - break - - frame_np = np.frombuffer(frame, dtype=np.int16) - if self.frames_np is not None and self.frames_np.shape[0] > 60*self.RATE: - self.frames_offset += 45.0 - self.frames_np = self.frames_np[int(45*self.RATE):] - - if self.frames_np is None: - self.frames_np = frame_np.copy() - else: - self.frames_np = np.concatenate((self.frames_np, frame_np), axis=0) - - # set frames np to None so to stop translation for this client - self.frames_np = None - self.exit = True - except Exception as e: - if self.verbose: print(f"[ERROR]: {e}") - self.exit = True - - def send_response_to_client(self, message): - """ - Send serialized response to client. - """ - a = pickle.dumps(message) - message = struct.pack("Q",len(a))+a - self.client_socket.sendall(message) def fill_output(self, output): """ @@ -159,81 +96,86 @@ class ServeClient: text = '' else: text += seg - wrapped = self.wrapper.wrap( - text="".join(text + output))[-2:] - return " ".join(wrapped) + wrapped = "".join(text + output) + return wrapped + + def add_frames(self, frame_np): + if self.frames_np is not None and self.frames_np.shape[0] > 45*self.RATE: + self.frames_offset += 45.0 + self.frames_np = self.frames_np[int(30*self.RATE):] + if self.frames_np is None: + self.frames_np = frame_np.copy() + else: + self.frames_np = np.concatenate((self.frames_np, frame_np), axis=0) def speech_to_text(self): """ Process audio stream in an infinite loop. """ while True: - if self.exit: - self.mqttc.disconnect() - self.client_socket.close() - self.transcriber.destroy() + if self.exit: + logging.info("Exiting speech to text thread") break + + if self.frames_np is None: + continue - if self.frames_np is None: continue - - # clip audio if the current chunk exceeds 25 seconds, this basically implies that - # no valid segment for the last 25 seconds from whisper + # clip audio if the current chunk exceeds 30 seconds, this basically implies that + # no valid segment for the last 30 seconds from whisper if self.frames_np[int((self.timestamp_offset - self.frames_offset)*self.RATE):].shape[0] > 25 * self.RATE: duration = self.frames_np.shape[0] / self.RATE self.timestamp_offset = self.frames_offset + duration - 5 - - # add 200 ms from the last chunk if available - if len(self.text) and self.frames_np[:-int((self.timestamp_offset - self.frames_offset)*self.RATE)].shape[0]: - samples_take = max(0, (self.timestamp_offset - self.frames_offset)*self.RATE - 0.2*self.RATE) - else: - samples_take = max(0, (self.timestamp_offset - self.frames_offset)*self.RATE) + + samples_take = max(0, (self.timestamp_offset - self.frames_offset)*self.RATE) input_bytes = self.frames_np[int(samples_take):].copy() duration = input_bytes.shape[0] / self.RATE - if duration<1.0: continue - + if duration<1.0: + continue try: - input_sample = input_bytes.astype(np.float32) / 32768.0 + input_sample = input_bytes.copy() # set previous complete segment as initial prompt if len(self.text) and self.text[-1] != '': initial_prompt = self.text[-1] else: initial_prompt = None - + # whisper transcribe with prompt result = self.transcriber.transcribe(input_sample, initial_prompt=initial_prompt) if len(result): self.t_start = None - output, segments = self.update_segments(result, duration) - out_dict = { - 'text': output, - 'segments': segments - } - if self.topic is not None: - self.mqttc.publish(self.topic, payload=str(out_dict)) - self.send_response_to_client(out_dict) + last_segment = self.update_segments(result, duration) + if len(self.transcript) < self.send_last_n_segments: + segments = self.transcript + else: + segments = self.transcript[-self.send_last_n_segments:] + if last_segment is not None: + segments = segments + [last_segment] + + try: + self.websocket.send(json.dumps(segments)) + except Exception as e: + logging.info(f"[ERROR]: {e}") else: # show previous output if there is pause i.e. no output from whisper - output = '' + segments = [] if self.t_start is None: self.t_start = time.time() - if time.time() - self.t_start < self.show_prev_out_thresh: - output = self.fill_output('') - + if len(self.transcript) < self.send_last_n_segments: + segments = self.transcript + else: + segments = self.transcript[-self.send_last_n_segments:] + # add a blank if there is no speech for 3 seconds if len(self.text) and self.text[-1] != '': if time.time() - self.t_start > self.add_pause_thresh: self.text.append('') - # publish outputs - out_dict = { - 'text': output, - 'segments': [] - } - if self.topic is not None: - self.mqttc.publish(self.topic, payload=str(out_dict)) - self.send_response_to_client(out_dict) + try: + self.websocket.send(json.dumps(segments)) + except Exception as e: + logging.info(f"[INFO]: {e}") except Exception as e: - if self.verbose: print(f"[ERROR]: {e}") + logging.info(f"[INFO]: {e}") time.sleep(0.01) def update_segments(self, segments, duration): @@ -249,15 +191,15 @@ class ServeClient: transcription for the current chunk """ offset = None - transcript = [] self.current_out = '' + last_segment = None # process complete segments if len(segments) > 1: for i, s in enumerate(segments[:-1]): text_ = s.text self.text.append(text_) start, end = self.timestamp_offset + s.start, self.timestamp_offset + min(duration, s.end) - transcript.append( + self.transcript.append( { 'start': start, 'end': end, @@ -268,6 +210,11 @@ class ServeClient: offset = min(duration, s.end) self.current_out += segments[-1].text + last_segment = { + 'start': self.timestamp_offset + segments[-1].start, + 'end': self.timestamp_offset + min(duration, segments[-1].end), + 'text': self.current_out + } # if same incomplete segment is seen multiple times then update the offset # and append the segment to the list @@ -279,7 +226,7 @@ class ServeClient: if self.same_output_threshold > 5: if not len(self.text) or self.text[-1].strip().lower()!=self.current_out.strip().lower(): self.text.append(self.current_out) - transcript.append( + self.transcript.append( { 'start': self.timestamp_offset, 'end': self.timestamp_offset + duration, @@ -289,6 +236,7 @@ class ServeClient: self.current_out = '' offset = duration self.same_output_threshold = 0 + last_segment = None else: self.prev_out = self.current_out @@ -296,36 +244,15 @@ class ServeClient: if offset is not None: self.timestamp_offset += offset - # format and return output - output = self.current_out - return self.fill_output(output), transcript + return last_segment + + def cleanup(self): + logging.info("Cleaning up.") + self.exit = True + self.transcriber.destroy() + -if __name__=="__main__": - # create socket - server_socket = socket.socket(socket.AF_INET,socket.SOCK_STREAM) - server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) - host='127.0.0.1' - port=5901 - backlog=5 - socket_address = (host, port) - print('STARTING SERVER AT',socket_address,'...') - server_socket.bind(socket_address) - server_socket.listen(backlog) - client_sockets = [] - device = 0 - try: - while True: - client_socket, addr = server_socket.accept() - print('GOT CONNECTION FROM:', addr) - client = ServeClient(client_socket, device=f'cuda:{device}') - client_sockets.append(client_socket) - print("waiting for new connection") - except Exception as e: - print(f"[ERROR main]: {e}") - for sock in client_sockets: - try: - sock.close() - except: - pass - +if __name__ == "__main__": + with serve(recv_audio, "127.0.0.1", 9090) as server: + server.serve_forever() \ No newline at end of file diff --git a/websocket_server.py b/websocket_server.py deleted file mode 100644 index 2978109..0000000 --- a/websocket_server.py +++ /dev/null @@ -1,258 +0,0 @@ -# import asyncio -import websockets -import pickle, struct, time, pyaudio -import threading -import os, json -import wave -import textwrap - -import logging -logging.basicConfig(level = logging.INFO) - -from collections import deque -from dataclasses import dataclass - -import torch -import numpy as np -from websockets.sync import server -from websockets.sync.server import serve -from transcriber import WhisperModel - - -clients = {} - -def recv_audio(websocket): - """ - Receive audio chunks from client in an infinite loop. - """ - global clients - client = ServeClient(websocket) - clients[websocket] = client - while True: - try: - frame_data = websocket.recv() - if isinstance(frame_data, str): - logging.info(frame_data) - continue - frame_np = np.frombuffer(frame_data, np.float32) - clients[websocket].add_frames(frame_np) - - except websockets.ConnectionClosedOK: - clients[websocket].cleanup() - clients.pop(websocket) - logging.info("Connection Closed.") - break - - -class ServeClient: - RATE = 16000 - def __init__(self, websocket, topic=None, device=None): - self.payload_size = struct.calcsize("Q") - self.data = b"" - self.frames = b"" - self.transcriber = WhisperModel("small.en", compute_type="float16", local_files_only=False) - self.timestamp_offset = 0.0 - self.frames_np = None - self.frames_offset = 0.0 - self.text = [] - self.current_out = '' - self.prev_out = '' - self.t_start=None - self.exit = False - self.same_output_threshold = 0 - self.show_prev_out_thresh = 5 # if pause(no output from whisper) show previous output for 5 seconds - self.add_pause_thresh = 3 # add a blank to segment list as a pause(no speech) for 3 seconds - self.transcript = [] - self.send_last_n_segments = 10 - - # text formatting - self.wrapper = textwrap.TextWrapper(width=50) - self.pick_previous_segments = 2 - - # setup mqtt - self.topic = topic - - # threading - self.websocket = websocket - self.trans_thread = threading.Thread(target=self.speech_to_text) - self.trans_thread.start() - - def fill_output(self, output): - """ - Format output with current and previous complete segments - into two lines of 50 characters. - - Args: - output(str): current incomplete segment - - Returns: - transcription wrapped in two lines - """ - text = '' - pick_prev = min(len(self.text), self.pick_previous_segments) - for seg in self.text[-pick_prev:]: - # discard everything before a 3 second pause - if seg == '': - text = '' - else: - text += seg - wrapped = "".join(text + output) - return wrapped - - def add_frames(self, frame_np): - if self.frames_np is not None and self.frames_np.shape[0] > 45*self.RATE: - self.frames_offset += 45.0 - self.frames_np = self.frames_np[int(30*self.RATE):] - if self.frames_np is None: - self.frames_np = frame_np.copy() - else: - self.frames_np = np.concatenate((self.frames_np, frame_np), axis=0) - - def speech_to_text(self): - """ - Process audio stream in an infinite loop. - """ - while True: - if self.exit: - logging.info("Exiting speech to text thread") - break - - if self.frames_np is None: - continue - - # clip audio if the current chunk exceeds 30 seconds, this basically implies that - # no valid segment for the last 30 seconds from whisper - if self.frames_np[int((self.timestamp_offset - self.frames_offset)*self.RATE):].shape[0] > 25 * self.RATE: - duration = self.frames_np.shape[0] / self.RATE - self.timestamp_offset = self.frames_offset + duration - 5 - - samples_take = max(0, (self.timestamp_offset - self.frames_offset)*self.RATE) - input_bytes = self.frames_np[int(samples_take):].copy() - duration = input_bytes.shape[0] / self.RATE - if duration<1.0: - continue - try: - input_sample = input_bytes.copy() - # set previous complete segment as initial prompt - if len(self.text) and self.text[-1] != '': - initial_prompt = self.text[-1] - else: - initial_prompt = None - - # whisper transcribe with prompt - result = self.transcriber.transcribe(input_sample, initial_prompt=initial_prompt) - if len(result): - self.t_start = None - last_segment = self.update_segments(result, duration) - if len(self.transcript) < self.send_last_n_segments: - segments = self.transcript - else: - segments = self.transcript[-self.send_last_n_segments:] - if last_segment is not None: - segments = segments + [last_segment] - - try: - self.websocket.send(json.dumps(segments)) - except Exception as e: - logging.info(f"[ERROR]: {e}") - else: - # show previous output if there is pause i.e. no output from whisper - segments = [] - if self.t_start is None: self.t_start = time.time() - if time.time() - self.t_start < self.show_prev_out_thresh: - if len(self.transcript) < self.send_last_n_segments: - segments = self.transcript - else: - segments = self.transcript[-self.send_last_n_segments:] - - # add a blank if there is no speech for 3 seconds - if len(self.text) and self.text[-1] != '': - if time.time() - self.t_start > self.add_pause_thresh: - self.text.append('') - - try: - self.websocket.send(json.dumps(segments)) - except Exception as e: - logging.info(f"[INFO]: {e}") - except Exception as e: - logging.info(f"[INFO]: {e}") - time.sleep(0.01) - - def update_segments(self, segments, duration): - """ - Processes the segments from whisper. Appends all the segments to the list - except for the last segment assuming that it is incomplete. - - Args: - segments(dict) : dictionary of segments as returned by whisper - duration(float): duration of the current chunk - - Returns: - transcription for the current chunk - """ - offset = None - self.current_out = '' - last_segment = None - # process complete segments - if len(segments) > 1: - for i, s in enumerate(segments[:-1]): - text_ = s.text - self.text.append(text_) - start, end = self.timestamp_offset + s.start, self.timestamp_offset + min(duration, s.end) - self.transcript.append( - { - 'start': start, - 'end': end, - 'text': text_ - } - ) - - offset = min(duration, s.end) - - self.current_out += segments[-1].text - last_segment = { - 'start': self.timestamp_offset + segments[-1].start, - 'end': self.timestamp_offset + min(duration, segments[-1].end), - 'text': self.current_out - } - - # if same incomplete segment is seen multiple times then update the offset - # and append the segment to the list - if self.current_out.strip() == self.prev_out.strip() and self.current_out != '': - self.same_output_threshold += 1 - else: - self.same_output_threshold = 0 - - if self.same_output_threshold > 5: - if not len(self.text) or self.text[-1].strip().lower()!=self.current_out.strip().lower(): - self.text.append(self.current_out) - self.transcript.append( - { - 'start': self.timestamp_offset, - 'end': self.timestamp_offset + duration, - 'text': self.current_out - } - ) - self.current_out = '' - offset = duration - self.same_output_threshold = 0 - last_segment = None - else: - self.prev_out = self.current_out - - # update offset - if offset is not None: - self.timestamp_offset += offset - - return last_segment - - def cleanup(self): - logging.info("Cleaning up.") - self.exit = True - self.transcriber.destroy() - - - -if __name__ == "__main__": - with serve(recv_audio, "127.0.0.1", 9090) as server: - server.serve_forever() \ No newline at end of file