Update app.py

app.py

@@ -275,70 +275,70 @@ def transcribe_video_with_speakers(video_path):
     return transcript_with_speakers, detected_language
 
 # Function to get the appropriate translation model based on target language
-def get_translation_model(source_language, target_language):
-    """
-    Get the translation model based on the source and target language.
+# def get_translation_model(source_language, target_language):
+#     """
+#     Get the translation model based on the source and target language.
 
-    Parameters:
-    - target_language (str): The language to translate the content into (e.g., 'es', 'fr').
-    - source_language (str): The language of the input content (default is 'en' for English).
+#     Parameters:
+#     - target_language (str): The language to translate the content into (e.g., 'es', 'fr').
+#     - source_language (str): The language of the input content (default is 'en' for English).
     
-    Returns:
-    - str: The translation model identifier.
-    """
-    # List of allowable languages
-    allowable_languages = ["en", "es", "fr", "zh", "de", "it", "pt", "ja", "ko", "ru", "hi", "tr"]
-
-    # Validate source and target languages
-    if source_language not in allowable_languages:
-        logger.debug(f"Invalid source language '{source_language}'. Supported languages are: {', '.join(allowable_languages)}")
-        # Return a default model if source language is invalid
-        source_language = "en"  # Default to 'en'
-
-    if target_language not in allowable_languages:
-        logger.debug(f"Invalid target language '{target_language}'. Supported languages are: {', '.join(allowable_languages)}")
-        # Return a default model if target language is invalid
-        target_language = "zh"  # Default to 'zh'
-
-    if source_language == target_language:
-        source_language = "en"  # Default to 'en'
-        target_language = "zh"  # Default to 'zh'
-
-    # Return the model using string concatenation
-    return f"Helsinki-NLP/opus-mt-{source_language}-{target_language}"
-
-def translate_single_entry(entry, translator):
-    original_text = entry["text"]
-    translated_text = translator(original_text)[0]['translation_text']
-    return {
-        "start": entry["start"],
-        "original": original_text,
-        "translated": translated_text,
-        "end": entry["end"],
-        "speaker": entry["speaker"]
-    }
-
-def translate_text(transcription_json, source_language, target_language):
-    # Load the translation model for the specified target language
-    translation_model_id = get_translation_model(source_language, target_language)
-    logger.debug(f"Translation model: {translation_model_id}")
-    translator = pipeline("translation", model=translation_model_id)
-
-    # Use ThreadPoolExecutor to parallelize translations
-    with concurrent.futures.ThreadPoolExecutor() as executor:
-        # Submit all translation tasks and collect results
-        translate_func = lambda entry: translate_single_entry(entry, translator)
-        translated_json = list(executor.map(translate_func, transcription_json))
-
-    # Sort the translated_json by start time
-    translated_json.sort(key=lambda x: x["start"])
-
-    # Log the components being added to translated_json
-    for entry in translated_json:
-        logger.debug("Added to translated_json: start=%s, original=%s, translated=%s, end=%s, speaker=%s",
-                     entry["start"], entry["original"], entry["translated"], entry["end"], entry["speaker"])
-
-    return translated_json
+#     Returns:
+#     - str: The translation model identifier.
+#     """
+#     # List of allowable languages
+#     allowable_languages = ["en", "es", "fr", "zh", "de", "it", "pt", "ja", "ko", "ru", "hi", "tr"]
+
+#     # Validate source and target languages
+#     if source_language not in allowable_languages:
+#         logger.debug(f"Invalid source language '{source_language}'. Supported languages are: {', '.join(allowable_languages)}")
+#         # Return a default model if source language is invalid
+#         source_language = "en"  # Default to 'en'
+
+#     if target_language not in allowable_languages:
+#         logger.debug(f"Invalid target language '{target_language}'. Supported languages are: {', '.join(allowable_languages)}")
+#         # Return a default model if target language is invalid
+#         target_language = "zh"  # Default to 'zh'
+
+#     if source_language == target_language:
+#         source_language = "en"  # Default to 'en'
+#         target_language = "zh"  # Default to 'zh'
+
+#     # Return the model using string concatenation
+#     return f"Helsinki-NLP/opus-mt-{source_language}-{target_language}"
+
+# def translate_single_entry(entry, translator):
+#     original_text = entry["text"]
+#     translated_text = translator(original_text)[0]['translation_text']
+#     return {
+#         "start": entry["start"],
+#         "original": original_text,
+#         "translated": translated_text,
+#         "end": entry["end"],
+#         "speaker": entry["speaker"]
+#     }
+
+# def translate_text(transcription_json, source_language, target_language):
+#     # Load the translation model for the specified target language
+#     translation_model_id = get_translation_model(source_language, target_language)
+#     logger.debug(f"Translation model: {translation_model_id}")
+#     translator = pipeline("translation", model=translation_model_id)
+
+#     # Use ThreadPoolExecutor to parallelize translations
+#     with concurrent.futures.ThreadPoolExecutor() as executor:
+#         # Submit all translation tasks and collect results
+#         translate_func = lambda entry: translate_single_entry(entry, translator)
+#         translated_json = list(executor.map(translate_func, transcription_json))
+
+#     # Sort the translated_json by start time
+#     translated_json.sort(key=lambda x: x["start"])
+
+#     # Log the components being added to translated_json
+#     for entry in translated_json:
+#         logger.debug("Added to translated_json: start=%s, original=%s, translated=%s, end=%s, speaker=%s",
+#                      entry["start"], entry["original"], entry["translated"], entry["end"], entry["speaker"])
+
+#     return translated_json
 
 def update_translations(file, edited_table, process_mode):
     """
@@ -518,220 +518,265 @@ def solve_optimal_alignment(original_segments, generated_durations, total_durati
 
     return original_segments
 
-ocr_model = None
-ocr_lock = threading.Lock()
-
-def init_ocr_model():
-    global ocr_model
-    with ocr_lock:
-        if ocr_model is None:
-            ocr_model = PaddleOCR(use_angle_cls=True, lang="ch")
-
-def find_best_subtitle_region(frame, ocr_model, region_height_ratio=0.35, num_strips=5, min_conf=0.5):
-    """
-    Automatically identifies the best subtitle region in a video frame using OCR confidence.
-
-    Parameters:
-    - frame: full video frame (BGR np.ndarray)
-    - ocr_model: a loaded PaddleOCR model
-    - region_height_ratio: portion of image height to scan (from bottom up)
-    - num_strips: how many horizontal strips to evaluate
-    - min_conf: minimum average confidence to consider a region valid
-
-    Returns:
-    - crop_region: the cropped image region with highest OCR confidence
-    - region_box: (y_start, y_end) of the region in the original frame
-    """
-    height, width, _ = frame.shape
-    region_height = int(height * region_height_ratio)
-    base_y_start = height - region_height
-    strip_height = region_height // num_strips
-
-    best_score = -1
-    best_crop = None
-    best_bounds = (0, height)
-
-    for i in range(num_strips):
-        y_start = base_y_start + i * strip_height
-        y_end = y_start + strip_height
-        strip = frame[y_start:y_end, :]
-
-        try:
-            result = ocr_model.ocr(strip, cls=True)
-            if not result or not result[0]:
-                continue
+# ocr_model = None
+# ocr_lock = threading.Lock()
+
+# def init_ocr_model():
+#     global ocr_model
+#     with ocr_lock:
+#         if ocr_model is None:
+#             ocr_model = PaddleOCR(use_angle_cls=True, lang="ch")
+
+# def find_best_subtitle_region(frame, ocr_model, region_height_ratio=0.35, num_strips=5, min_conf=0.5):
+#     """
+#     Automatically identifies the best subtitle region in a video frame using OCR confidence.
+
+#     Parameters:
+#     - frame: full video frame (BGR np.ndarray)
+#     - ocr_model: a loaded PaddleOCR model
+#     - region_height_ratio: portion of image height to scan (from bottom up)
+#     - num_strips: how many horizontal strips to evaluate
+#     - min_conf: minimum average confidence to consider a region valid
+
+#     Returns:
+#     - crop_region: the cropped image region with highest OCR confidence
+#     - region_box: (y_start, y_end) of the region in the original frame
+#     """
+#     height, width, _ = frame.shape
+#     region_height = int(height * region_height_ratio)
+#     base_y_start = height - region_height
+#     strip_height = region_height // num_strips
+
+#     best_score = -1
+#     best_crop = None
+#     best_bounds = (0, height)
+
+#     for i in range(num_strips):
+#         y_start = base_y_start + i * strip_height
+#         y_end = y_start + strip_height
+#         strip = frame[y_start:y_end, :]
+
+#         try:
+#             result = ocr_model.ocr(strip, cls=True)
+#             if not result or not result[0]:
+#                 continue
 
-            total_score = sum(line[1][1] for line in result[0])
-            avg_score = total_score / len(result[0])
+#             total_score = sum(line[1][1] for line in result[0])
+#             avg_score = total_score / len(result[0])
 
-            if avg_score > best_score:
-                best_score = avg_score
-                best_crop = strip
-                best_bounds = (y_start, y_end)
+#             if avg_score > best_score:
+#                 best_score = avg_score
+#                 best_crop = strip
+#                 best_bounds = (y_start, y_end)
 
-        except Exception as e:
-            continue  # Fail silently on OCR issues
+#         except Exception as e:
+#             continue  # Fail silently on OCR issues
 
-    if best_score >= min_conf and best_crop is not None:
-        return best_crop, best_bounds
-    else:
-        # Fallback to center-bottom strip
-        fallback_y = height - int(height * 0.2)
-        return frame[fallback_y:, :], (fallback_y, height)
+#     if best_score >= min_conf and best_crop is not None:
+#         return best_crop, best_bounds
+#     else:
+#         # Fallback to center-bottom strip
+#         fallback_y = height - int(height * 0.2)
+#         return frame[fallback_y:, :], (fallback_y, height)
         
-def ocr_frame_worker(args, min_confidence=0.7):
-    frame_idx, frame_time, frame = args
+# def ocr_frame_worker(args, min_confidence=0.7):
+#     frame_idx, frame_time, frame = args
+
+#     init_ocr_model()  # Load model in thread-safe way
+
+#     if frame is None or frame.size == 0 or not isinstance(frame, np.ndarray):
+#         return {"time": frame_time, "text": ""}
+
+#     if frame.dtype != np.uint8:
+#         frame = frame.astype(np.uint8)
+
+#     try:
+#         result = ocr_model.ocr(frame, cls=True)
+#         lines = result[0] if result else []
+#         texts = [line[1][0] for line in lines if line[1][1] >= min_confidence]
+#         combined_text = " ".join(texts).strip()
+#         return {"time": frame_time, "text": combined_text}
+#     except Exception as e:
+#         print(f"⚠️ OCR failed at {frame_time:.2f}s: {e}")
+#         return {"time": frame_time, "text": ""}
+
+# def frame_is_in_audio_segments(frame_time, audio_segments, tolerance=0.2):
+#     for segment in audio_segments:
+#         start, end = segment["start"], segment["end"]
+#         if (start - tolerance) <= frame_time <= (end + tolerance):
+#             return True
+#     return False
+
+# def extract_ocr_subtitles_parallel(video_path, transcription_json, 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_idx % int(fps * interval_sec) == 0:
+#             frame_time = frame_idx / fps
+#             if frame_is_in_audio_segments(frame_time, transcription_json):
+#                 frames.append((frame_idx, frame_time, frame.copy()))
+#         success, frame = cap.read()
+#         frame_idx += 1
+#     cap.release()
+
+#     ocr_results = []
+#     ocr_failures = 0  # Count OCR failures
+#     with concurrent.futures.ThreadPoolExecutor(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:
+#                 ocr_failures += 1
+
+#     logger.info(f"✅ OCR extraction completed: {len(ocr_results)} frames successful, {ocr_failures} frames failed.")
+#     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
+#             logger.debug(f"MERGED: Current end extended to {time:.2f}s for text: '{current['text'][:50]}...' (Similarity: {sim})")
+#         else:
+#             logger.debug(f"NOT MERGING (Similarity: {sim} < Threshold: {text_similarity_threshold}):")
+#             logger.debug(f"  Previous segment: {current['start']:.2f}s - {current['end']:.2f}s: '{current['text'][:50]}...'")
+#             logger.debug(f"  New segment: {time:.2f}s: '{text[:50]}...'")
+#             collapsed.append(current)
+#             current = {"start": time, "end": time, "text": text}
+#     if current:
+#         collapsed.append(current)
+
+#     logger.info(f"✅ OCR subtitles collapsed into {len(collapsed)} segments.")
+#     for idx, seg in enumerate(collapsed):
+#         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,
+#     replace_threshold=90,
+#     time_tolerance=1.0
+# ):
+#     merged = []
+#     used_whisperx = set()
+#     whisperx_used_flags = [False] * len(whisperx_json)
 
-    init_ocr_model()  # Load model in thread-safe way
+#     # Step 1: Attempt to match each OCR entry to a WhisperX entry
+#     for ocr in ocr_json:
+#         ocr_start, ocr_end = ocr["start"], ocr["end"]
+#         ocr_text = ocr["text"]
 
-    if frame is None or frame.size == 0 or not isinstance(frame, np.ndarray):
-        return {"time": frame_time, "text": ""}
+#         best_match = None
+#         best_score = -1
+#         best_idx = None
 
-    if frame.dtype != np.uint8:
-        frame = frame.astype(np.uint8)
+#         for idx, wx in enumerate(whisperx_json):
+#             wx_start, wx_end = wx["start"], wx["end"]
+#             wx_text = wx["text"]
 
-    try:
-        result = ocr_model.ocr(frame, cls=True)
-        lines = result[0] if result else []
-        texts = [line[1][0] for line in lines if line[1][1] >= min_confidence]
-        combined_text = " ".join(texts).strip()
-        return {"time": frame_time, "text": combined_text}
-    except Exception as e:
-        print(f"⚠️ OCR failed at {frame_time:.2f}s: {e}")
-        return {"time": frame_time, "text": ""}
-
-def frame_is_in_audio_segments(frame_time, audio_segments, tolerance=0.2):
-    for segment in audio_segments:
-        start, end = segment["start"], segment["end"]
-        if (start - tolerance) <= frame_time <= (end + tolerance):
-            return True
-    return False
-
-def extract_ocr_subtitles_parallel(video_path, transcription_json, 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_idx % int(fps * interval_sec) == 0:
-            frame_time = frame_idx / fps
-            if frame_is_in_audio_segments(frame_time, transcription_json):
-                frames.append((frame_idx, frame_time, frame.copy()))
-        success, frame = cap.read()
-        frame_idx += 1
-    cap.release()
-
-    ocr_results = []
-    ocr_failures = 0  # Count OCR failures
-    with concurrent.futures.ThreadPoolExecutor(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:
-                ocr_failures += 1
+#             # Check for time overlap
+#             overlap = not (ocr_end < wx_start - time_tolerance or ocr_start > wx_end + time_tolerance)
+#             if not overlap:
+#                 continue
 
-    logger.info(f"✅ OCR extraction completed: {len(ocr_results)} frames successful, {ocr_failures} frames failed.")
-    return ocr_results
+#             sim = fuzz.ratio(ocr_text, wx_text)
+#             if sim > best_score:
+#                 best_score = sim
+#                 best_match = wx
+#                 best_idx = idx
 
-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 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)
+
+#     # Step 2: Add untouched WhisperX segments
+#     for idx, wx in enumerate(whisperx_json):
+#         if not whisperx_used_flags[idx]:
+#             merged.append(wx)
+
+#     # Step 3: Sort all merged segments
+#     merged = sorted(merged, key=lambda x: x["start"])
 
-        if not current:
-            current = {"start": time, "end": time, "text": text}
-            continue
+#     return merged
 
-        sim = fuzz.ratio(current["text"], text)
-        if sim >= text_similarity_threshold:
-            current["end"] = time
-            logger.debug(f"MERGED: Current end extended to {time:.2f}s for text: '{current['text'][:50]}...' (Similarity: {sim})")
-        else:
-            logger.debug(f"NOT MERGING (Similarity: {sim} < Threshold: {text_similarity_threshold}):")
-            logger.debug(f"  Previous segment: {current['start']:.2f}s - {current['end']:.2f}s: '{current['text'][:50]}...'")
-            logger.debug(f"  New segment: {time:.2f}s: '{text[:50]}...'")
-            collapsed.append(current)
-            current = {"start": time, "end": time, "text": text}
-    if current:
-        collapsed.append(current)
-
-    logger.info(f"✅ OCR subtitles collapsed into {len(collapsed)} segments.")
-    for idx, seg in enumerate(collapsed):
-        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,
-    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_end = ocr["start"], 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"]
-
-            # 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)
-            if sim > best_score:
-                best_score = sim
-                best_match = wx
-                best_idx = idx
-
-        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)
+def process_segment_with_gpt(segment, source_lang, target_lang, model="gpt-4"):
+    original_text = segment["text"]
+    prompt = (
+        f"You are a multilingual assistant. Given the following text in {source_lang}, "
+        f"1) restore punctuation, and 2) translate it into {target_lang}.\n\n"
+        f"Text:\n{original_text}\n\n"
+        f"Return in JSON format:\n"
+        f'{{"punctuated": "...", "translated": "..."}}'
+    )
 
-    # Step 2: Add untouched WhisperX segments
-    for idx, wx in enumerate(whisperx_json):
-        if not whisperx_used_flags[idx]:
-            merged.append(wx)
+    try:
+        response = openai.ChatCompletion.create(
+            model=model,
+            messages=[{"role": "user", "content": prompt}],
+            temperature=0.3
+        )
+        content = response.choices[0].message.content.strip()
+        result_json = eval(content) if content.startswith("{") else {}
 
-    # Step 3: Sort all merged segments
-    merged = sorted(merged, key=lambda x: x["start"])
+        return {
+            "start": segment["start"],
+            "end": segment["end"],
+            "speaker": segment.get("speaker", "SPEAKER_00"),
+            "original": result_json.get("punctuated", original_text),
+            "translated": result_json.get("translated", "")
+        }
 
-    return merged
+    except Exception as e:
+        print(f"❌ Error for segment {segment['start']}-{segment['end']}: {e}")
+        return {
+            "start": segment["start"],
+            "end": segment["end"],
+            "speaker": segment.get("speaker", "SPEAKER_00"),
+            "original": original_text,
+            "translated": ""
+        }
 
+def punctuate_and_translate_parallel(transcription_json, source_lang="zh", target_lang="en", model="gpt-4o", max_workers=5):
+    with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
+        futures = [
+            executor.submit(process_segment_with_gpt, seg, source_lang, target_lang, model)
+            for seg in transcription_json
+        ]
+        return [f.result() for f in concurrent.futures.as_completed(futures)]
+        
 # def merge_speaker_and_time_from_whisperx(ocr_json, whisperx_json, text_sim_threshold=80, replace_threshold=90):
 #     merged = []
 #     used_whisperx = set()
@@ -1132,10 +1177,10 @@ 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_transcribed_segments(transcription_json, file.name)
+        translated_json_raw = punctuate_and_translate_parallel(transcription_json, source_language, target_language)
         # Step 2: Translate the transcription
-        logger.info(f"Translating transcription from {source_language} to {target_language}...")
-        translated_json_raw = translate_text(transcription_json_merged, source_language, target_language)
+        # logger.info(f"Translating transcription from {source_language} to {target_language}...")
+        # translated_json_raw = translate_text(transcription_json_merged, )
         logger.info(f"Translation completed. Number of translated segments: {len(translated_json_raw)}")
 
         translated_json = apply_adaptive_speed(translated_json_raw, source_language, target_language)