-
Notifications
You must be signed in to change notification settings - Fork 0
/
non_real_time_asr.py
206 lines (172 loc) · 7.84 KB
/
non_real_time_asr.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
import os
import json
from pydub import AudioSegment
import wave
from pyannote.audio import Pipeline
from transformers import pipeline
import torch
from huggingface_hub import login
from multiprocessing import Pool, cpu_count
import numpy as np
import signal
# Set environment variable to enable CPU fallback for MPS
os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
# Timeout handler
class TimeoutException(Exception):
pass
def timeout_handler(signum, frame):
raise TimeoutException("Diarization process timed out")
# Register the timeout handler
signal.signal(signal.SIGALRM, timeout_handler)
def convert_m4a_to_wav(input_path, output_path):
try:
print(f"Converting {input_path} to {output_path}")
audio = AudioSegment.from_file(input_path, format="m4a")
audio.export(output_path, format="wav")
print(f"Conversion complete: {output_path}")
except Exception as e:
print(f"Error during conversion: {e}")
def transcribe_audio_file(filepath, whisper_model):
try:
# Load audio data
with wave.open(filepath, "rb") as wf:
data = wf.readframes(wf.getnframes())
rate = wf.getframerate()
array = np.frombuffer(data, dtype=np.int16).astype(np.float32) / 32768.0 # Convert to float32 and normalize
input_data = {"array": array, "sampling_rate": rate}
# Perform transcription
transcription = whisper_model(input_data)["text"]
words = [{"word": w, "start": i, "end": i + 1} for i, w in enumerate(transcription.split())] # Mock word timings
return transcription, words
except Exception as e:
print(f"Error processing file {filepath}: {e}")
return None, None
def diarize_audio_file(filepath, pipeline, device, timeout_duration=600): # Increased timeout duration
try:
print(f"Diarizing file: {filepath}")
signal.alarm(timeout_duration) # Set the timeout
pipeline.to(device) # Ensure the pipeline is using the correct device
diarization = pipeline({'audio': filepath})
signal.alarm(0) # Cancel the timeout
return diarization
except TimeoutException:
print(f"Diarization for {filepath} timed out.")
return None
except Exception as e:
print(f"Error during diarization: {e}")
return None
def process_file(file_info):
filename, directory, hf_token, device = file_info
transcriptions = {}
diarizations = {}
if filename.endswith(".m4a"):
filepath = os.path.join(directory, filename)
wav_filename = filename.replace(".m4a", ".wav")
wav_filepath = os.path.join(directory, wav_filename)
print(f"Found m4a file: {filepath}")
convert_m4a_to_wav(filepath, wav_filepath)
print(f"Processing file: {wav_filepath}")
# Initialize Whisper model
whisper_model = pipeline("automatic-speech-recognition", model="openai/whisper-base", device=device)
transcription, word_list = transcribe_audio_file(wav_filepath, whisper_model)
if transcription:
transcriptions[wav_filename] = (transcription, word_list)
print(f"Transcription for {wav_filepath}: {transcription}")
# Initialize the Hugging Face login and diarization pipeline within each process
login(token=hf_token)
try:
diarization_pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization", use_auth_token=hf_token)
diarization = diarize_audio_file(wav_filepath, diarization_pipeline, device)
if diarization:
diarizations[wav_filename] = diarization
print(f"Diarization for {wav_filepath}: {diarization}")
except Exception as e:
print(f"Error during diarization: {e}")
return transcriptions, diarizations
def transcribe_and_diarize_audio_files(directory, hf_token):
transcriptions = {}
diarizations = {}
print(f"Directory to scan: {directory}")
device = torch.device('mps' if torch.backends.mps.is_available() else 'cpu')
file_info = [(filename, directory, hf_token, device) for filename in os.listdir(directory)]
with Pool(processes=cpu_count()) as pool:
results = pool.map(process_file, file_info)
for result in results:
trans, diar = result
transcriptions.update(trans)
diarizations.update(diar)
return transcriptions, diarizations
def label_speakers(diarization):
speakers = {}
for segment, _, speaker in diarization.itertracks(yield_label=True):
if speaker not in speakers:
speakers[speaker] = "Doctor" if len(speakers) == 0 else "Patient"
return speakers
def format_output(transcription, word_list, diarization, speakers):
diarized_transcript = []
diarized_segments = []
segments = [(segment, speakers[speaker]) for segment, _, speaker in diarization.itertracks(yield_label=True)]
labeled_segments = [(segment.start, segment.end, label) for segment, label in segments]
labeled_segments.sort() # Ensure segments are sorted by start time
current_speaker = None
current_text = []
for segment in labeled_segments:
start, end, label = segment
words_in_segment = [word['word'] for word in word_list if start <= word['start'] < end]
if words_in_segment:
if current_speaker != label:
if current_speaker is not None:
diarized_transcript.append(f"[{current_speaker}] {' '.join(current_text)}")
current_speaker = label
current_text = words_in_segment
else:
current_text.extend(words_in_segment)
diarized_segments.append({
"start": start,
"end": end,
"label": label,
"words": ' '.join(words_in_segment)
})
if current_text:
diarized_transcript.append(f"[{current_speaker}] {' '.join(current_text)}")
return '\n'.join(diarized_transcript), diarized_segments
def save_to_json(transcriptions, diarizations, output_directory):
output = {}
diarized_output = {}
for filename in transcriptions:
transcription, word_list = transcriptions[filename]
diarization = diarizations.get(filename)
if diarization:
speakers = label_speakers(diarization)
formatted_output, diarized_segments = format_output(transcription, word_list, diarization, speakers)
output[filename] = {
"transcription": transcription,
"diarization": formatted_output
}
diarized_output[filename] = {
"diarized_text": formatted_output,
"segments": diarized_segments
}
else:
output[filename] = {
"transcription": transcription
}
if not os.path.exists(output_directory):
os.makedirs(output_directory)
output_filepath = os.path.join(output_directory, "output.json")
with open(output_filepath, 'w') as f:
json.dump(output, f, indent=4)
print(f"Results saved to {output_filepath}")
diarized_output_filepath = os.path.join(output_directory, "diarized_output.json")
with open(diarized_output_filepath, 'w') as f:
json.dump(diarized_output, f, indent=4)
print(f"Diarized segments saved to {diarized_output_filepath}")
if __name__ == "__main__":
# Authenticate with Hugging Face
hf_token = "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"
login(token=hf_token)
audio_directory = os.path.join(os.path.dirname(__file__), '..', 'MediNotes', 'Recordings')
output_directory = os.path.join(os.path.dirname(__file__), '..', 'MediNotes', 'Output')
print(f"Audio directory: {audio_directory}")
transcriptions, diarizations = transcribe_and_diarize_audio_files(audio_directory, hf_token)
save_to_json(transcriptions, diarizations, output_directory)