Make writing output audio file optional when using microphone

Signed-off-by: makaveli10 <suryanvineet47@gmail.com>
This commit is contained in:
makaveli10
2024-05-27 17:28:48 +05:30
parent 03e30e1fed
commit 61d07edabb
+24 -16
View File
@@ -1,4 +1,5 @@
import os import os
import shutil
import wave import wave
import numpy as np import numpy as np
@@ -272,7 +273,7 @@ class TranscriptionTeeClient:
Attributes: Attributes:
clients (list): the underlying Client instances responsible for handling WebSocket connections. clients (list): the underlying Client instances responsible for handling WebSocket connections.
""" """
def __init__(self, clients): def __init__(self, clients, save_output_recording=False, output_recording_filename="./output_recording.wav"):
self.clients = clients self.clients = clients
if not self.clients: if not self.clients:
raise Exception("At least one client is required.") raise Exception("At least one client is required.")
@@ -281,6 +282,8 @@ class TranscriptionTeeClient:
self.channels = 1 self.channels = 1
self.rate = 16000 self.rate = 16000
self.record_seconds = 60000 self.record_seconds = 60000
self.save_output_recording = save_output_recording
self.output_recording_filename = output_recording_filename
self.frames = b"" self.frames = b""
self.p = pyaudio.PyAudio() self.p = pyaudio.PyAudio()
try: try:
@@ -473,7 +476,7 @@ class TranscriptionTeeClient:
return process return process
def record(self, out_file="output_recording.wav"): def record(self):
""" """
Record audio data from the input stream and save it to a WAV file. Record audio data from the input stream and save it to a WAV file.
@@ -485,15 +488,12 @@ class TranscriptionTeeClient:
The recording will continue until the specified duration is reached or until the `RECORDING` flag is set to `False`. The recording will continue until the specified duration is reached or until the `RECORDING` flag is set to `False`.
The recording process can be interrupted by sending a KeyboardInterrupt (e.g., pressing Ctrl+C). After recording, The recording process can be interrupted by sending a KeyboardInterrupt (e.g., pressing Ctrl+C). After recording,
the method combines all the saved audio chunks into the specified `out_file`. the method combines all the saved audio chunks into the specified `out_file`.
Args:
out_file (str, optional): The name of the output WAV file to save the entire recording.
Default is "output_recording.wav".
""" """
n_audio_file = 0 n_audio_file = 0
if not os.path.exists("chunks"): if self.save_output_recording:
os.makedirs("chunks", exist_ok=True) if os.path.exists("chunks"):
shutil.rmtree("chunks")
os.makedirs("chunks")
try: try:
for _ in range(0, int(self.rate / self.chunk * self.record_seconds)): for _ in range(0, int(self.rate / self.chunk * self.record_seconds)):
if not any(client.recording for client in self.clients): if not any(client.recording for client in self.clients):
@@ -507,6 +507,7 @@ class TranscriptionTeeClient:
# save frames if more than a minute # save frames if more than a minute
if len(self.frames) > 60 * self.rate: if len(self.frames) > 60 * self.rate:
if self.save_output_recording:
t = threading.Thread( t = threading.Thread(
target=self.write_audio_frames_to_file, target=self.write_audio_frames_to_file,
args=( args=(
@@ -520,7 +521,7 @@ class TranscriptionTeeClient:
self.write_all_clients_srt() self.write_all_clients_srt()
except KeyboardInterrupt: except KeyboardInterrupt:
if len(self.frames): if self.save_output_recording and len(self.frames):
self.write_audio_frames_to_file( self.write_audio_frames_to_file(
self.frames[:], f"chunks/{n_audio_file}.wav" self.frames[:], f"chunks/{n_audio_file}.wav"
) )
@@ -529,8 +530,8 @@ class TranscriptionTeeClient:
self.stream.close() self.stream.close()
self.p.terminate() self.p.terminate()
self.close_all_clients() self.close_all_clients()
if self.save_output_recording:
self.write_output_recording(n_audio_file, out_file) self.write_output_recording(n_audio_file)
self.write_all_clients_srt() self.write_all_clients_srt()
def write_audio_frames_to_file(self, frames, file_name): def write_audio_frames_to_file(self, frames, file_name):
@@ -552,7 +553,7 @@ class TranscriptionTeeClient:
wavfile.setframerate(self.rate) wavfile.setframerate(self.rate)
wavfile.writeframes(frames) wavfile.writeframes(frames)
def write_output_recording(self, n_audio_file, out_file): def write_output_recording(self, n_audio_file):
""" """
Combine and save recorded audio chunks into a single WAV file. Combine and save recorded audio chunks into a single WAV file.
@@ -571,7 +572,7 @@ class TranscriptionTeeClient:
for i in range(n_audio_file) for i in range(n_audio_file)
if os.path.exists(f"chunks/{i}.wav") if os.path.exists(f"chunks/{i}.wav")
] ]
with wave.open(out_file, "wb") as wavfile: with wave.open(self.output_recording_filename, "wb") as wavfile:
wavfile: wave.Wave_write wavfile: wave.Wave_write
wavfile.setnchannels(self.channels) wavfile.setnchannels(self.channels)
wavfile.setsampwidth(2) wavfile.setsampwidth(2)
@@ -586,6 +587,9 @@ class TranscriptionTeeClient:
# remove this file # remove this file
os.remove(in_file) os.remove(in_file)
wavfile.close() wavfile.close()
# clean up temporary directory to store chunks
if os.path.exists("chunks"):
shutil.rmtree("chunks")
@staticmethod @staticmethod
def bytes_to_float_array(audio_bytes): def bytes_to_float_array(audio_bytes):
@@ -616,6 +620,8 @@ class TranscriptionClient(TranscriptionTeeClient):
port (int): The port number to connect to on the server. port (int): The port number to connect to on the server.
lang (str, optional): The primary language for transcription. Default is None, which defaults to English ('en'). lang (str, optional): The primary language for transcription. Default is None, which defaults to English ('en').
translate (bool, optional): Indicates whether translation tasks are required (default is False). translate (bool, optional): Indicates whether translation tasks are required (default is False).
save_output_recording (bool, optional): Indicates whether to save recording from microphone.
output_recording_filename (str, optional): File to save the output recording.
Attributes: Attributes:
client (Client): An instance of the underlying Client class responsible for handling the WebSocket connection. client (Client): An instance of the underlying Client class responsible for handling the WebSocket connection.
@@ -627,6 +633,8 @@ class TranscriptionClient(TranscriptionTeeClient):
transcription_client() transcription_client()
``` ```
""" """
def __init__(self, host, port, lang=None, translate=False, model="small", use_vad=True): def __init__(self, host, port, lang=None, translate=False, model="small", use_vad=True, save_output_recording=False, output_recording_filename="./output_recording.wav"):
self.client = Client(host, port, lang, translate, model, srt_file_path="output.srt", use_vad=use_vad) self.client = Client(host, port, lang, translate, model, srt_file_path="output.srt", use_vad=use_vad)
TranscriptionTeeClient.__init__(self, [self.client]) if save_output_recording and not output_recording_filename.endswith(".wav"):
raise ValueError(f"Please provide a valid `output_recording_filename`: {output_recording_filename}")
TranscriptionTeeClient.__init__(self, [self.client], save_output_recording=save_output_recording, output_recording_filename=output_recording_filename)