101 lines
3.4 KiB
Python
101 lines
3.4 KiB
Python
# Note: This feature requires pyannote.audio and a HuggingFace token.
|
|
# If these are not present, this module will likely fail or raise errors.
|
|
# Due to the complexity and weight of pyannote.audio, this is a placeholder
|
|
# for where the logic would sit. Implementing full diarization requires
|
|
# downloading models and handling complex segment merging.
|
|
|
|
import os
|
|
import sys
|
|
|
|
def diarize_audio(audio_path, num_speakers=None, hf_token=None):
|
|
"""
|
|
Performs speaker diarization using pyannote.audio.
|
|
|
|
Args:
|
|
audio_path (str): Path to audio file.
|
|
num_speakers (int, optional): Number of speakers if known.
|
|
hf_token (str): HuggingFace Auth Token.
|
|
|
|
Returns:
|
|
list: List of segments with speaker labels [(start, end, speaker), ...].
|
|
"""
|
|
try:
|
|
from pyannote.audio import Pipeline
|
|
except ImportError:
|
|
print("Error: pyannote.audio not installed. Diarization skipped.")
|
|
return None
|
|
|
|
if not hf_token:
|
|
print("Error: HuggingFace Token (HF_TOKEN) not found. Diarization skipped.")
|
|
return None
|
|
|
|
print(f"Loading Diarization Pipeline (pyannote/speaker-diarization-3.1)...")
|
|
try:
|
|
# Note: 'use_auth_token' was deprecated in favor of 'token' in recent versions
|
|
pipeline = Pipeline.from_pretrained(
|
|
"pyannote/speaker-diarization-3.1",
|
|
token=hf_token
|
|
)
|
|
|
|
# Move to GPU if available
|
|
import torch
|
|
if torch.cuda.is_available():
|
|
pipeline.to(torch.device("cuda"))
|
|
|
|
print(f"Diarizing {audio_path}...")
|
|
diarization = pipeline(audio_path, num_speakers=num_speakers)
|
|
|
|
results = []
|
|
for turn, _, speaker in diarization.itertracks(yield_label=True):
|
|
results.append({
|
|
"start": turn.start,
|
|
"end": turn.end,
|
|
"speaker": speaker
|
|
})
|
|
|
|
return results
|
|
|
|
except Exception as e:
|
|
print(f"Error during diarization: {e}")
|
|
return None
|
|
|
|
def merge_diarization_with_transcript(transcript_segments, diarization_segments):
|
|
"""
|
|
Merges Whisper segments with Diarization speaker labels based on time overlap.
|
|
|
|
Args:
|
|
transcript_segments (list): Whisper segments [{'start': 0.0, 'end': 1.0, 'text': 'Hi'}, ...]
|
|
diarization_segments (list): Diarization segments [{'start': 0.1, 'end': 0.9, 'speaker': 'SPEAKER_00'}]
|
|
|
|
Returns:
|
|
list: Enhanced transcript segments with 'speaker' key.
|
|
"""
|
|
if not diarization_segments:
|
|
return transcript_segments
|
|
|
|
# Simple overlap matching logic
|
|
for t_seg in transcript_segments:
|
|
# Find diarization segment with max overlap
|
|
t_start = t_seg['start']
|
|
t_end = t_seg['end']
|
|
|
|
best_speaker = "Unknown"
|
|
max_overlap = 0
|
|
|
|
for d_seg in diarization_segments:
|
|
d_start = d_seg['start']
|
|
d_end = d_seg['end']
|
|
|
|
# Calculate intersection
|
|
overlap_start = max(t_start, d_start)
|
|
overlap_end = min(t_end, d_end)
|
|
overlap_duration = max(0, overlap_end - overlap_start)
|
|
|
|
if overlap_duration > max_overlap:
|
|
max_overlap = overlap_duration
|
|
best_speaker = d_seg['speaker']
|
|
|
|
t_seg['speaker'] = best_speaker
|
|
|
|
return transcript_segments
|