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
+33 -25
View File
@@ -1,4 +1,5 @@
import os
import shutil
import wave
import numpy as np
@@ -272,7 +273,7 @@ class TranscriptionTeeClient:
Attributes:
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
if not self.clients:
raise Exception("At least one client is required.")
@@ -281,6 +282,8 @@ class TranscriptionTeeClient:
self.channels = 1
self.rate = 16000
self.record_seconds = 60000
self.save_output_recording = save_output_recording
self.output_recording_filename = output_recording_filename
self.frames = b""
self.p = pyaudio.PyAudio()
try:
@@ -473,7 +476,7 @@ class TranscriptionTeeClient:
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.
@@ -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 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`.
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
if not os.path.exists("chunks"):
os.makedirs("chunks", exist_ok=True)
if self.save_output_recording:
if os.path.exists("chunks"):
shutil.rmtree("chunks")
os.makedirs("chunks")
try:
for _ in range(0, int(self.rate / self.chunk * self.record_seconds)):
if not any(client.recording for client in self.clients):
@@ -507,20 +507,21 @@ class TranscriptionTeeClient:
# save frames if more than a minute
if len(self.frames) > 60 * self.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
if self.save_output_recording:
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""
self.write_all_clients_srt()
except KeyboardInterrupt:
if len(self.frames):
if self.save_output_recording and len(self.frames):
self.write_audio_frames_to_file(
self.frames[:], f"chunks/{n_audio_file}.wav"
)
@@ -529,8 +530,8 @@ class TranscriptionTeeClient:
self.stream.close()
self.p.terminate()
self.close_all_clients()
self.write_output_recording(n_audio_file, out_file)
if self.save_output_recording:
self.write_output_recording(n_audio_file)
self.write_all_clients_srt()
def write_audio_frames_to_file(self, frames, file_name):
@@ -552,7 +553,7 @@ class TranscriptionTeeClient:
wavfile.setframerate(self.rate)
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.
@@ -571,7 +572,7 @@ class TranscriptionTeeClient:
for i in range(n_audio_file)
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.setnchannels(self.channels)
wavfile.setsampwidth(2)
@@ -586,6 +587,9 @@ class TranscriptionTeeClient:
# remove this file
os.remove(in_file)
wavfile.close()
# clean up temporary directory to store chunks
if os.path.exists("chunks"):
shutil.rmtree("chunks")
@staticmethod
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.
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).
save_output_recording (bool, optional): Indicates whether to save recording from microphone.
output_recording_filename (str, optional): File to save the output recording.
Attributes:
client (Client): An instance of the underlying Client class responsible for handling the WebSocket connection.
@@ -627,6 +633,8 @@ class TranscriptionClient(TranscriptionTeeClient):
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)
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)