Update app.py

app.py

@@ -668,13 +668,19 @@ def collapse_ocr_subtitles(ocr_json, text_similarity_threshold=90):
         logger.debug(f"[OCR Collapsed {idx}] {seg['start']:.2f}s - {seg['end']:.2f}s: {seg['text'][:50]}...")
     return collapsed
 
-def merge_speaker_and_time_from_whisperx(ocr_json, whisperx_json, text_sim_threshold=80, replace_threshold=90):
+def merge_speaker_and_time_from_whisperx(
+    ocr_json,
+    whisperx_json,
+    replace_threshold=90,
+    time_tolerance=1.0
+):
     merged = []
     used_whisperx = set()
+    whisperx_used_flags = [False] * len(whisperx_json)
 
+    # Step 1: Attempt to match each OCR entry to a WhisperX entry
     for ocr in ocr_json:
-        ocr_start = ocr["start"]
-        ocr_end = ocr["end"]
+        ocr_start, ocr_end = ocr["start"], ocr["end"]
         ocr_text = ocr["text"]
 
         best_match = None
@@ -685,11 +691,9 @@ def merge_speaker_and_time_from_whisperx(ocr_json, whisperx_json, text_sim_thres
             wx_start, wx_end = wx["start"], wx["end"]
             wx_text = wx["text"]
 
-            if idx in used_whisperx:
-                continue  # Already matched
-
-            time_center_diff = abs((ocr_start + ocr_end)/2 - (wx_start + wx_end)/2)
-            if time_center_diff > 3:
+            # Check for time overlap
+            overlap = not (ocr_end < wx_start - time_tolerance or ocr_start > wx_end + time_tolerance)
+            if not overlap:
                 continue
 
             sim = fuzz.ratio(ocr_text, wx_text)
@@ -698,23 +702,83 @@ def merge_speaker_and_time_from_whisperx(ocr_json, whisperx_json, text_sim_thres
                 best_match = wx
                 best_idx = idx
 
-        new_entry = copy.deepcopy(ocr)
-        if best_match:
-            new_entry["speaker"] = best_match.get("speaker", "UNKNOWN")
-            new_entry["ocr_similarity"] = best_score
+        if best_match and best_score >= replace_threshold:
+            # Replace WhisperX segment with higher quality OCR text
+            new_segment = copy.deepcopy(best_match)
+            new_segment["text"] = ocr_text
+            new_segment["ocr_replaced"] = True
+            new_segment["ocr_similarity"] = best_score
+            whisperx_used_flags[best_idx] = True
+            merged.append(new_segment)
+        else:
+            # No replacement, check if this OCR is outside WhisperX time coverage
+            covered = any(
+                abs((ocr_start + ocr_end)/2 - (wx["start"] + wx["end"])/2) < time_tolerance
+                for wx in whisperx_json
+            )
+            if not covered:
+                new_segment = copy.deepcopy(ocr)
+                new_segment["ocr_added"] = True
+                new_segment["speaker"] = "UNKNOWN"
+                merged.append(new_segment)
 
-            if best_score >= replace_threshold:
-                new_entry["start"] = best_match["start"]
-                new_entry["end"] = best_match["end"]
-                used_whisperx.add(best_idx)  # Mark used
+    # Step 2: Add untouched WhisperX segments
+    for idx, wx in enumerate(whisperx_json):
+        if not whisperx_used_flags[idx]:
+            merged.append(wx)
 
-        else:
-            new_entry["speaker"] = "UNKNOWN"
-            new_entry["ocr_similarity"] = None
+    # Step 3: Sort all merged segments
+    merged = sorted(merged, key=lambda x: x["start"])
 
-        merged.append(new_entry)
     return merged
 
+# def merge_speaker_and_time_from_whisperx(ocr_json, whisperx_json, text_sim_threshold=80, replace_threshold=90):
+#     merged = []
+#     used_whisperx = set()
+
+#     for ocr in ocr_json:
+#         ocr_start = ocr["start"]
+#         ocr_end = ocr["end"]
+#         ocr_text = ocr["text"]
+
+#         best_match = None
+#         best_score = -1
+#         best_idx = None
+
+#         for idx, wx in enumerate(whisperx_json):
+#             wx_start, wx_end = wx["start"], wx["end"]
+#             wx_text = wx["text"]
+
+#             if idx in used_whisperx:
+#                 continue  # Already matched
+
+#             time_center_diff = abs((ocr_start + ocr_end)/2 - (wx_start + wx_end)/2)
+#             if time_center_diff > 3:
+#                 continue
+
+#             sim = fuzz.ratio(ocr_text, wx_text)
+#             if sim > best_score:
+#                 best_score = sim
+#                 best_match = wx
+#                 best_idx = idx
+
+#         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"]
+#                 used_whisperx.add(best_idx)  # Mark used
+
+#         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.