Update app.py

app.py

@@ -133,7 +133,7 @@ def handle_feedback(feedback):
             conn.commit()
         return "Thank you for your feedback!", None
 
-def segment_background_audio(audio_path, background_audio_path="background_segments.wav"):
+def segment_background_audio(audio_path, background_audio_path="background_segments.wav", speech_audio_path="speech_segment.wav"):
     """
     Uses Demucs to separate audio and extract background (non-vocal) parts.
     Merges drums, bass, and other stems into a single background track.
@@ -150,6 +150,7 @@ def segment_background_audio(audio_path, background_audio_path="background_segme
     stem_dir = os.path.join("separated", "htdemucs", filename)
 
     # Step 3: Load and merge background stems
+    vocals = AudioSegment.from_wav(os.path.join(stem_dir, "vocals.wav"))
     drums = AudioSegment.from_wav(os.path.join(stem_dir, "drums.wav"))
     bass = AudioSegment.from_wav(os.path.join(stem_dir, "bass.wav"))
     other = AudioSegment.from_wav(os.path.join(stem_dir, "other.wav"))
@@ -158,7 +159,8 @@ def segment_background_audio(audio_path, background_audio_path="background_segme
 
     # Step 4: Export the merged background
     background.export(background_audio_path, format="wav")
-    return background_audio_path
+    vocals.export(speech_audio_path, format="wav")
+    return background_audio_path, speech_audio_path
 
 def transcribe_video_with_speakers(video_path):
     # Extract audio from video
@@ -167,7 +169,7 @@ def transcribe_video_with_speakers(video_path):
     video.audio.write_audiofile(audio_path)
     logger.info(f"Audio extracted from video: {audio_path}")
 
-    segment_result = segment_background_audio(audio_path)
+    segment_result, speech_audio_path = segment_background_audio(audio_path)
     print(f"Saved non-speech (background) audio to local")
     
     # Set up device
@@ -180,7 +182,7 @@ def transcribe_video_with_speakers(video_path):
         logger.info("WhisperX model loaded")
     
         # Transcribe
-        result = model.transcribe(audio_path, chunk_size=6, print_progress = True)
+        result = model.transcribe(speech_audio_path, chunk_size=6, print_progress = True)
         logger.info("Audio transcription completed")
 
         # Get the detected language
@@ -188,12 +190,12 @@ def transcribe_video_with_speakers(video_path):
         logger.debug(f"Detected language: {detected_language}")
         # Alignment
         # model_a, metadata = whisperx.load_align_model(language_code=result["language"], device=device)
-        # result = whisperx.align(result["segments"], model_a, metadata, audio_path, device)
+        # result = whisperx.align(result["segments"], model_a, metadata, speech_audio_path, device)
         # logger.info("Transcription alignment completed")
     
         # Diarization (works independently of Whisper model size)
         diarize_model = whisperx.DiarizationPipeline(use_auth_token=hf_api_key, device=device)
-        diarize_segments = diarize_model(audio_path)
+        diarize_segments = diarize_model(speech_audio_path)
         logger.info("Speaker diarization completed")
     
         # Assign speakers
@@ -213,31 +215,62 @@ def transcribe_video_with_speakers(video_path):
         }
         for segment in result["segments"]
     ]
-
+        
     # Collect audio for each speaker
     speaker_audio = {}
-    for segment in result["segments"]:
-        speaker = segment["speaker"]
-        if speaker not in speaker_audio:
-            speaker_audio[speaker] = []
-        speaker_audio[speaker].append((segment["start"], segment["end"]))
-
+    logger.info("🔎 Start collecting valid audio segments per speaker...")
+    
+    for idx, segment in enumerate(result["segments"]):
+        speaker = segment.get("speaker", "SPEAKER_00")
+        start = segment["start"]
+        end = segment["end"]
+    
+        if end > start and (end - start) > 0.05:  # Require >50ms duration
+            if speaker not in speaker_audio:
+                speaker_audio[speaker] = [(start, end)]
+            else:
+                speaker_audio[speaker].append((start, end))
+            
+            logger.debug(f"Segment {idx}: Added to speaker {speaker} [{start:.2f}s → {end:.2f}s]")
+        else:
+            logger.warning(f"⚠️ Segment {idx} discarded: invalid duration ({start:.2f}s → {end:.2f}s)")
+    
     # Collapse and truncate speaker audio
     speaker_sample_paths = {}
-    audio_clip = AudioFileClip(audio_path)
+    audio_clip = AudioFileClip(speech_audio_path)
+    
+    logger.info(f"🔎 Found {len(speaker_audio)} speakers with valid segments. Start creating speaker samples...")
+    
     for speaker, segments in speaker_audio.items():
+        logger.info(f"🔹 Speaker {speaker}: {len(segments)} valid segments")
+    
         speaker_clips = [audio_clip.subclip(start, end) for start, end in segments]
-        combined_clip = concatenate_audioclips(speaker_clips)
+        if not speaker_clips:
+            logger.warning(f"⚠️ No valid audio clips for speaker {speaker}. Skipping sample creation.")
+            continue
+    
+        if len(speaker_clips) == 1:
+            logger.debug(f"Speaker {speaker}: Only one clip, skipping concatenation.")
+            combined_clip = speaker_clips[0]
+        else:
+            logger.debug(f"Speaker {speaker}: Concatenating {len(speaker_clips)} clips.")
+            combined_clip = concatenate_audioclips(speaker_clips)
+    
         truncated_clip = combined_clip.subclip(0, min(30, combined_clip.duration))
+        logger.debug(f"Speaker {speaker}: Truncated to {truncated_clip.duration:.2f} seconds.")
+
+        # Step 4: Save the final result
         sample_path = f"speaker_{speaker}_sample.wav"
         truncated_clip.write_audiofile(sample_path)
         speaker_sample_paths[speaker] = sample_path
-        logger.info(f"Created sample for {speaker}: {sample_path}")
+        logger.info(f"✅ Created and saved sample for {speaker}: {sample_path}")
 
-    # Clean up
+    # Cleanup
+    logger.info("🧹 Closing audio clip and removing temporary files...")
     video.close()
     audio_clip.close()
-    os.remove(audio_path)
+    os.remove(speech_audio_path)
+    logger.info("✅ Finished processing all speaker samples.")
 
     return transcript_with_speakers, detected_language
 
@@ -637,6 +670,74 @@ def collapse_ocr_subtitles(ocr_json, text_similarity_threshold=90):
 
     return collapsed
 
+def merge_speaker_and_time_from_whisperx(ocr_json, whisperx_json, text_sim_threshold=80, replace_threshold=90):
+    """
+    Given OCR and WhisperX segments, merge speaker ID and optionally replace time.
+    """
+    merged = []
+
+    for ocr in ocr_json:
+        ocr_start = ocr["start"]
+        ocr_end = ocr["end"]
+        ocr_text = ocr["text"]
+        
+        best_match = None
+        best_score = -1
+
+        for wx in whisperx_json:
+            wx_start, wx_end = wx["start"], wx["end"]
+            wx_text = wx["text"]
+            
+            # Time overlap (soft constraint)
+            time_center_diff = abs((ocr_start + ocr_end)/2 - (wx_start + wx_end)/2)
+            if time_center_diff > 3:  # skip if too far
+                continue
+
+            # Text similarity
+            sim = fuzz.ratio(ocr_text, wx_text)
+            if sim > best_score:
+                best_score = sim
+                best_match = wx
+
+        new_entry = copy.deepcopy(ocr)
+        if best_match:
+            new_entry["speaker"] = best_match.get("speaker", "UNKNOWN")
+            new_entry["ocr_similarity"] = best_score
+
+            if best_score >= replace_threshold:
+                new_entry["start"] = best_match["start"]
+                new_entry["end"] = best_match["end"]
+
+        else:
+            new_entry["speaker"] = "UNKNOWN"
+            new_entry["ocr_similarity"] = None
+
+        merged.append(new_entry)
+
+    return merged
+
+
+def realign_ocr_segments(merged_ocr_json, min_gap=0.2):
+    """
+    Realign OCR segments to avoid overlaps using midpoint-based adjustment.
+    """
+    merged_ocr_json = sorted(merged_ocr_json, key=lambda x: x["start"])
+
+    for i in range(1, len(merged_ocr_json)):
+        prev = merged_ocr_json[i - 1]
+        curr = merged_ocr_json[i]
+
+        # If current overlaps with previous, adjust
+        if curr["start"] < prev["end"] + min_gap:
+            midpoint = (prev["end"] + curr["start"]) / 2
+            prev["end"] = round(midpoint - min_gap / 2, 3)
+            curr["start"] = round(midpoint + min_gap / 2, 3)
+
+            # Prevent negative durations
+            if curr["start"] >= curr["end"]:
+                curr["end"] = round(curr["start"] + 0.3, 3)
+
+    return merged_ocr_json
 def post_edit_transcribed_segments(transcription_json, video_path,
                                    interval_sec=0.5,
                                    text_similarity_threshold=80,
@@ -658,80 +759,12 @@ def post_edit_transcribed_segments(transcription_json, video_path,
     # Step 2: Collapse repetitive OCR
     collapsed_ocr = collapse_ocr_subtitles(ocr_json, text_similarity_threshold=90)
 
-    # Step 3: Refine existing WhisperX segments (Phase 1)
-    merged_segments = []
-    used_ocr_indices = set()
-
-    for entry_idx, entry in enumerate(transcription_json):
-        start = entry.get("start", 0)
-        end = entry.get("end", 0)
-        base_text = entry.get("text", "")
-
-        best_match_idx = None
-        best_score = -1
-
-        for ocr_idx, ocr in enumerate(collapsed_ocr):
-            time_overlap = not (ocr["end"] < start - time_tolerance or ocr["start"] > end + time_tolerance)
-            if not time_overlap:
-                continue
-
-            sim = fuzz.ratio(ocr["text"], base_text)
-            if sim > best_score:
-                best_score = sim
-                best_match_idx = ocr_idx
-
-        updated_entry = entry.copy()
-        if best_match_idx is not None and best_score >= text_similarity_threshold:
-            updated_entry["text"] = collapsed_ocr[best_match_idx]["text"]
-            updated_entry["ocr_matched"] = True
-            updated_entry["ocr_similarity"] = best_score
-            used_ocr_indices.add(best_match_idx)
-        else:
-            updated_entry["ocr_matched"] = False
-            updated_entry["ocr_similarity"] = best_score if best_score >= 0 else None
-
-        merged_segments.append(updated_entry)
-    
-    # Step 4: Insert unused OCR segments (Phase 2)
-    inserted_segments = []
-    for ocr_idx, ocr in enumerate(collapsed_ocr):
-        if ocr_idx in used_ocr_indices:
-            continue
-    
-        # Check for fuzzy duplicates in WhisperX
-        duplicate = False
-        for whisper_seg in transcription_json:
-            if abs(ocr["start"] - whisper_seg["start"]) < time_tolerance or abs(ocr["end"] - whisper_seg["end"]) < time_tolerance:
-                sim = fuzz.ratio(ocr["text"], whisper_seg["text"])
-                if sim >= text_similarity_threshold:
-                    duplicate = True
-                    break
-    
-        if duplicate:
-            logger.debug(f"🟡 Skipping near-duplicate OCR: '{ocr['text']}'")
-            continue
-    
-        # Infer speaker from nearest WhisperX entry
-        nearby = sorted(transcription_json, key=lambda x: abs(x["start"] - ocr["start"]))
-        speaker_guess = nearby[0].get("speaker", "unknown") if nearby else "unknown"
-    
-        inserted_segment = {
-            "start": ocr["start"],
-            "end": ocr["end"],
-            "text": ocr["text"],
-            "speaker": speaker_guess
-        }
-        inserted_segments.append(inserted_segment)
-
-    # Step 5: Combine and sort
-    final_segments = merged_segments + inserted_segments
-    final_segments = sorted(final_segments, key=lambda x: x["start"])
-
-    print(f"✅ Post-editing completed: {len(final_segments)} total segments "
-          f"({len(inserted_segments)} OCR-inserted segments)")
-
-    return final_segments
+    # Step 3: Merge and realign OCR segments.
+    ocr_merged = merge_speaker_and_time_from_whisperx(ocr_json, whisperx_json)
+    ocr_realigned = realign_ocr_segments(ocr_merged)
 
+    logger.info(f"✅ Final merged and realigned OCR: {len(ocr_realigned)} segments")
+    return ocr_realigned
 
 def process_entry(entry, i, tts_model, video_width, video_height, process_mode, target_language, font_path, speaker_sample_paths=None):
     logger.debug(f"Processing entry {i}: {entry}")