Update app.py

app.py

@@ -383,10 +383,11 @@ def create_subtitle_clip_pil(text, start_time, end_time, video_width, video_heig
         logger.error(f"\u274c Failed to create subtitle clip: {e}")
         return None
 
+
 def solve_optimal_alignment(original_segments, generated_durations, total_duration):
     """
-    Robust version: Aligns generated speech segments, falls back to greedy allocation if solver fails.
-    Modifies and returns the translated_json with updated 'start' and 'end'.
+    Aligns speech segments using quadratic programming. If optimization fails,
+    applies greedy fallback: center shorter segments, stretch longer ones.
     """
     N = len(original_segments)
     d = np.array(generated_durations)
@@ -414,13 +415,105 @@ def solve_optimal_alignment(original_segments, generated_durations, total_durati
     except Exception as e:
         print(f"⚠️ Optimization failed: {e}, falling back to greedy alignment.")
 
-        current_time = 0.0
         for i in range(N):
-            original_segments[i]['start'] = round(current_time, 3)
-            original_segments[i]['end'] = round(current_time + generated_durations[i], 3)
-            current_time += generated_durations[i]
+            orig_start = original_segments[i]['start']
+            orig_end = original_segments[i]['end']
+            orig_mid = (orig_start + orig_end) / 2
+            gen_duration = generated_durations[i]
+            orig_duration = orig_end - orig_start
+
+            if gen_duration <= orig_duration:
+                new_start = orig_mid - gen_duration / 2
+                new_end = orig_mid + gen_duration / 2
+            else:
+                extra = (gen_duration - orig_duration) / 2
+                new_start = orig_start - extra
+                new_end = orig_end + extra
+
+                # Prevent overlap with previous
+                if i > 0:
+                    prev_end = original_segments[i - 1]['end']
+                    new_start = max(new_start, prev_end + 0.01)
+
+                # Prevent overlap with next
+                if i < N - 1:
+                    next_start = original_segments[i + 1]['start']
+                    new_end = min(new_end, next_start - 0.01)
+
+                if new_end <= new_start:
+                    new_start = orig_start
+                    new_end = orig_start + gen_duration
+
+            original_segments[i]['start'] = round(new_start, 3)
+            original_segments[i]['end'] = round(new_end, 3)
 
     return original_segments
+
+def get_frame_image_bytes(video, t):
+    frame = video.get_frame(t)
+    img = Image.fromarray(frame)
+    buf = io.BytesIO()
+    img.save(buf, format='JPEG')
+    return buf.getvalue()
+
+def post_edit_segment(entry, image_bytes):
+    try:
+        system_prompt = """You are a multilingual assistant helping polish subtitles and voiceover content.
+Your job is to fix punctuation, validate meaning, improve tone, and ensure the translation matches the speaker's intended message."""
+
+        user_prompt = f"""
+Original (source) transcript: {entry.get("original", "")}
+Translated version: {entry.get("translated", "")}
+Speaker ID: {entry.get("speaker", "")}
+Time: {entry.get("start")} - {entry.get("end")}
+
+Please:
+1. Add correct punctuation and sentence boundaries.
+2. Improve fluency and tone of the translated text.
+3. Ensure the meaning is preserved from the original.
+4. Use the attached image frame to infer emotion or setting.
+
+Return the revised original and translated texts in the following format:
+Original: <edited original>
+Translated: <edited translation>
+"""
+        response = ChatCompletion.create(
+            model="gpt-4o",
+            messages=[
+                {"role": "system", "content": system_prompt},
+                {"role": "user", "content": user_prompt, "image": image_bytes}
+            ]
+        )
+
+        output = response.choices[0].message.content.strip()
+        lines = output.splitlines()
+        for line in lines:
+            if line.startswith("Original:"):
+                entry['original'] = line[len("Original:"):].strip()
+            elif line.startswith("Translated:"):
+                entry['translated'] = line[len("Translated:"):].strip()
+
+        return entry
+    except Exception as e:
+        print(f"Post-editing failed for segment: {e}")
+        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
@@ -644,6 +737,8 @@ def upload_and_manage(file, target_language, process_mode):
         translated_json = translate_text(transcription_json, source_language, target_language)
         logger.info(f"Translation completed. Number of translated segments: {len(translated_json)}")
 
+        translated_json = post_edit_translated_segments(translated_json, file.name)
+
         # Step 3: Add transcript to video based on timestamps
         logger.info("Adding translated transcript to video...")
         add_transcript_voiceover(file.name, translated_json, output_video_path, process_mode, target_language)