Added revisions to the translation app

video_transcription/ai_transcriber_v1/diarizer.py

@@ -0,0 +1,100 @@
+# 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