61 lines
1.8 KiB
Python
61 lines
1.8 KiB
Python
import torch
|
|
from pyannote.audio import Pipeline
|
|
import whisper
|
|
|
|
AUDIO_FILE = "2024-07-29_audio.wav"
|
|
filename = (AUDIO_FILE[::-1].split(".")[1].split("/")[0].split("\\")[0])[::-1]
|
|
|
|
|
|
torch.backends.cuda.matmul.allow_tf32 = True
|
|
torch.backends.cudnn.allow_tf32 = True
|
|
|
|
# Load Whisper model for transcription
|
|
whisper_model = whisper.load_model("large")
|
|
|
|
|
|
# Transcribe audio using Whisper
|
|
def transcribe_audio(audio_path):
|
|
result = whisper_model.transcribe(audio_path)
|
|
segments = result["segments"]
|
|
return [(segment["start"], segment["end"], segment["text"]) for segment in segments]
|
|
|
|
|
|
# Initialize Pyannote Pipeline for diarization
|
|
pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization-3.1")
|
|
pipeline.to(torch.device("cuda"))
|
|
|
|
|
|
# Perform diarization
|
|
def perform_diarization(audio_path) -> Pipeline:
|
|
diarization = pipeline(audio_path, min_speakers=5, max_speakers=7)
|
|
|
|
# dump the diarization output to disk using RTTM format
|
|
with open(f"diarization_{filename}.rttm", "w") as rttm:
|
|
diarization.write_rttm(rttm)
|
|
|
|
print("Finished diarization")
|
|
|
|
return diarization
|
|
|
|
|
|
# Load audio and perform both transcription and diarization
|
|
transcription_segments = transcribe_audio(AUDIO_FILE)
|
|
diarization: Pipeline = perform_diarization(AUDIO_FILE)
|
|
|
|
# Print speaker and corresponding text
|
|
print("\nSpeaker and Text Segments:")
|
|
diarization_with_text = []
|
|
for segment in transcription_segments:
|
|
start, end, text = segment
|
|
for spk_segment, _, speaker_label in diarization.itertracks(yield_label=True):
|
|
if spk_segment.start < end and spk_segment.end > start:
|
|
diarization_with_text.append(f"Speaker {speaker_label}: {text}")
|
|
break
|
|
|
|
|
|
with open(f"transcript-diarization_{filename}.txt", "w") as fp:
|
|
fp.writelines(f"{l}\n" for l in diarization_with_text)
|
|
|
|
print("\n".join(diarization_with_text))
|
|
|
|
print("Saved diarization")
|