Update app.py

app.py

@@ -41,6 +41,7 @@ import soundfile as sf
 from paddleocr import PaddleOCR
 import cv2
 from rapidfuzz import fuzz
+from tqdm import tqdm
 
 logger = logging.getLogger(__name__)
 
@@ -513,77 +514,117 @@ def solve_optimal_alignment(original_segments, generated_durations, total_durati
 
     return original_segments
 
-def extract_subtitles_with_ocr(video_path):
-    ocr = PaddleOCR(use_angle_cls=True, lang="ch")  # Change `lang` as needed
-    vidcap = cv2.VideoCapture(video_path)
-    fps = vidcap.get(cv2.CAP_PROP_FPS)
-
-    subtitles = []
-    frame_id = 0
-    success, image = vidcap.read()
+def ocr_frame_worker(args):
+    frame_idx, frame_time, frame = args
+    ocr = PaddleOCR(use_angle_cls=True, lang="ch")  # Initialize OCR inside worker
+    result = ocr.ocr(frame, cls=True)
+    texts = [line[1][0] for line in result[0]] if result[0] else []
+    combined_text = " ".join(texts).strip()
+    return {"time": frame_time, "text": combined_text}
+
+def extract_ocr_subtitles_parallel(video_path, interval_sec=0.5, num_workers=4):
+    cap = cv2.VideoCapture(video_path)
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    frames = []
+    frame_idx = 0
+    success, frame = cap.read()
 
     while success:
-        if frame_id % int(fps) == 0:  # OCR 1 frame per second (adjust if needed)
-            result = ocr.ocr(image, cls=True)
-            texts = [line[1][0] for line in result[0]]  # Get text parts
-            combined_text = " ".join(texts).strip()
-            if combined_text:
-                subtitles.append({
-                    "time": frame_id / fps,
-                    "text": combined_text
-                })
-
-        frame_id += 1
-        success, image = vidcap.read()
-
-    vidcap.release()
-    return subtitles
-
-def align_subtitles_to_transcripts(ocr_subtitles, whisperx_segments):
-    aligned_pairs = []
-
-    for ocr_entry in ocr_subtitles:
-        ocr_time = ocr_entry["time"]
-        best_score = -1
-        best_segment = None
-
-        for seg in whisperx_segments:
-            # Only consider segments close in time (within +/- 2s)
-            if abs(seg["start"] - ocr_time) < 2.0 or abs(seg["end"] - ocr_time) < 2.0:
-                score = fuzz.ratio(seg["text"], ocr_entry["text"])
-                if score > best_score:
-                    best_score = score
-                    best_segment = seg
-
-        if best_segment:
-            aligned_pairs.append({
-                "whisper_text": best_segment["text"],
-                "ocr_text": ocr_entry["text"],
-                "start": best_segment["start"],
-                "end": best_segment["end"],
-                "similarity": best_score
-            })
-
-    return aligned_pairs
-
-def correct_transcripts_with_ocr(aligned_pairs):
-    corrected_segments = []
-
-    for pair in aligned_pairs:
-        if pair["similarity"] > 80:
-            # Trust OCR more if they are close
-            corrected_text = pair["ocr_text"]
+        if frame_idx % int(fps * interval_sec) == 0:
+            frame_time = frame_idx / fps
+            frames.append((frame_idx, frame_time, frame.copy()))
+        success, frame = cap.read()
+        frame_idx += 1
+    cap.release()
+
+    ocr_results = []
+    with concurrent.futures.ProcessPoolExecutor(max_workers=num_workers) as executor:
+        futures = [executor.submit(ocr_frame_worker, frame) for frame in frames]
+        
+        for f in tqdm(concurrent.futures.as_completed(futures), total=len(futures)):
+            try:
+                result = f.result()
+                if result["text"]:
+                    ocr_results.append(result)
+            except Exception as e:
+                print(f"⚠️ OCR failed for a frame: {e}")
+    return ocr_results
+
+def collapse_ocr_subtitles(ocr_json, text_similarity_threshold=90):
+    collapsed = []
+    current = None
+    for entry in ocr_json:
+        time = entry["time"]
+        text = entry["text"]
+
+        if not current:
+            current = {"start": time, "end": time, "text": text}
+            continue
+
+        sim = fuzz.ratio(current["text"], text)
+        if sim >= text_similarity_threshold:
+            current["end"] = time
         else:
-            corrected_text = pair["whisper_text"]
+            collapsed.append(current)
+            current = {"start": time, "end": time, "text": text}
+    if current:
+        collapsed.append(current)
+    return collapsed
+
+def post_edit_transcribed_segments(transcription_json, video_path,
+                                   interval_sec=0.5,
+                                   text_similarity_threshold=80,
+                                   time_tolerance=1.0,
+                                   num_workers=4):
+    """
+    Given WhisperX transcription (transcription_json) and video,
+    use OCR subtitles to post-correct and merge the transcriptions.
+    """
+
+    # Step 1: Extract OCR subtitles
+    ocr_json = extract_ocr_subtitles_parallel(video_path, interval_sec=interval_sec, num_workers=num_workers)
+
+    # Step 2: Collapse repetitive OCR
+    collapsed_ocr = collapse_ocr_subtitles(ocr_json, text_similarity_threshold=90)
+
+    # Step 3: Merge OCR with WhisperX
+    merged_segments = []
 
-        corrected_segments.append({
-            "start": pair["start"],
-            "end": pair["end"],
-            "text": corrected_text
-        })
+    for entry in transcription_json:
+        start = entry.get("start", 0)
+        end = entry.get("end", 0)
+        base_text = entry.get("text", "")
 
-    return corrected_segments
+        best_match = None
+        best_score = -1
+
+        for ocr in collapsed_ocr:
+            # Check time overlap
+            time_overlap = not (ocr["end"] < start - time_tolerance or ocr["start"] > end + time_tolerance)
+            if not time_overlap:
+                continue
+
+            # Text similarity
+            sim = fuzz.ratio(ocr["text"], base_text)
+            if sim > best_score:
+                best_score = sim
+                best_match = ocr
+
+        # If good match found, replace the original text
+        updated_entry = entry.copy()
+        if best_match and best_score >= text_similarity_threshold:
+            updated_entry["text"] = best_match["text"]
+            updated_entry["ocr_matched"] = True
+            updated_entry["ocr_similarity"] = best_score
+        else:
+            updated_entry["ocr_matched"] = False
+            updated_entry["ocr_similarity"] = best_score if best_score >= 0 else None
 
+        merged_segments.append(updated_entry)
+
+    print(f"✅ Post-editing completed: {len(merged_segments)} segments")
+    return merged_segments
+    
 # def get_frame_image_bytes(video, t):
 #     frame = video.get_frame(t)
 #     img = Image.fromarray(frame)
@@ -634,21 +675,6 @@ def correct_transcripts_with_ocr(aligned_pairs):
 #         return entry
 
 
-# def post_edit_translated_segments(translated_json, video_path):
-#     video = VideoFileClip(video_path)
-
-#     def process(entry):
-#         mid_time = (entry['start'] + entry['end']) / 2
-#         image_bytes = get_frame_image_bytes(video, mid_time)
-#         entry = post_edit_segment(entry, image_bytes)
-#         return entry
-
-#     with concurrent.futures.ThreadPoolExecutor() as executor:
-#         edited = list(executor.map(process, translated_json))
-
-#     video.close()
-#     return edited
-
 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}")
     error_message = None
@@ -953,12 +979,12 @@ def upload_and_manage(file, target_language, process_mode):
         transcription_json, source_language = transcribe_video_with_speakers(file.name)
         logger.info(f"Transcription completed. Detected source language: {source_language}")
 
+        transcription_json_merged = post_edit_translated_segments(transcription_json, file.name)
         # Step 2: Translate the transcription
         logger.info(f"Translating transcription from {source_language} to {target_language}...")
-        translated_json_raw = translate_text(transcription_json, source_language, target_language)
+        translated_json_raw = translate_text(transcription_json_merged, source_language, target_language)
         logger.info(f"Translation completed. Number of translated segments: {len(translated_json_raw)}")
 
-        # translated_json = post_edit_translated_segments(translated_json, file.name)
         translated_json = apply_adaptive_speed(translated_json_raw, source_language, target_language)
         
         # Step 3: Add transcript to video based on timestamps