Update pipeline/postprocess.py

pipeline/postprocess.py

@@ -1,4 +1,4 @@
-# ─── Post-Processing Engine v2 ─────────────────────────────────────────────────
+# ─── Post-Processing Engine v3 ─────────────────────────────────────────────────
 
 def format_timestamp(seconds: float) -> str:
     ms = round((seconds % 1) * 1000)
@@ -8,147 +8,134 @@ def format_timestamp(seconds: float) -> str:
     return f"{h:02d}:{m:02d}:{s:02d},{ms:03d}"
 
 
-def detect_cps_threshold(segments: list) -> float:
-    total_chars = 0
-    total_duration = 0.0
-    for seg in segments:
-        words = seg.get("words", [])
-        if not words:
-            continue
-        duration = words[-1]["end"] - words[0]["start"]
-        if duration > 0:
-            total_chars += len(seg["text"].replace(" ", ""))
-            total_duration += duration
-    if total_duration == 0:
-        return 20.0
-    avg_cps = total_chars / total_duration
-    if avg_cps < 12:
-        return 17.0
-    elif avg_cps <= 18:
-        return 20.0
-    else:
-        return 25.0
-
-
-def merge_segments(segments: list, max_chars: int = 80, max_gap: float = 1.5) -> list:
+def merge_segments(segments: list) -> list:
     """
-    segment‌های کوتاه رو با هم ادغام کن تا جمله‌های کامل بسازه.
-    max_chars: حداکثر کاراکتر برای یه subtitle
-    max_gap: اگه gap بین دو segment بیشتر از این بود، merge نکن
+    ادغام segment‌های کوتاه و تک‌کلمه‌ای به جملات کامل.
+    هر جمله حداکثر ۸۰ کاراکتر و حداکثر ۷ ثانیه.
     """
     if not segments:
         return []
 
+    MAX_CHARS = 80
+    MAX_DURATION = 7.0
+    MAX_GAP = 1.5  # اگه gap بین دو segment بیشتر از این بود، جدا بمون
+
     merged = []
-    current = segments[0].copy()
-    current["words"] = list(current.get("words", []))
+    current = None
 
-    for seg in segments[1:]:
-        current_text = current["text"].strip()
-        next_text = seg["text"].strip()
-        combined_chars = len((current_text + " " + next_text).replace(" ", ""))
+    for seg in segments:
+        if current is None:
+            current = {
+                "text": seg["text"].strip(),
+                "start": seg["start"],
+                "end": seg["end"],
+                "words": seg.get("words", []),
+            }
+            continue
 
-        # gap بین دو segment
         gap = seg["start"] - current["end"]
-
-        # شرایط merge:
-        # ۱. حروف مجموع زیر max_chars باشه
-        # ۲. gap زیاد نباشه
-        # ۳. جمله فعلی با نقطه تموم نشده باشه
-        ends_with_punct = current_text.endswith((".", "!", "?"))
-
-        if combined_chars <= max_chars and gap <= max_gap and not ends_with_punct:
-            current["text"] = current_text + " " + next_text
+        combined_text = current["text"] + " " + seg["text"].strip()
+        combined_duration = seg["end"] - current["start"]
+        ends_sentence = current["text"].rstrip().endswith((".", "!", "?"))
+
+        should_merge = (
+            len(combined_text.replace(" ", "")) <= MAX_CHARS
+            and combined_duration <= MAX_DURATION
+            and gap <= MAX_GAP
+            and not ends_sentence
+        )
+
+        if should_merge:
+            current["text"] = combined_text
             current["end"] = seg["end"]
-            current["words"].extend(seg.get("words", []))
+            current["words"] = current["words"] + seg.get("words", [])
         else:
             merged.append(current)
-            current = seg.copy()
-            current["words"] = list(current.get("words", []))
-
-    merged.append(current)
-    return merged
+            current = {
+                "text": seg["text"].strip(),
+                "start": seg["start"],
+                "end": seg["end"],
+                "words": seg.get("words", []),
+            }
 
+    if current:
+        merged.append(current)
 
-def find_split_point(words: list) -> int:
-    mid = len(words) // 2
-
-    # اولویت ۱ — مکث طبیعی (gap > 150ms)
-    best_gap_idx = -1
-    best_gap = 0.0
-    for i in range(1, len(words)):
-        gap = words[i]["start"] - words[i - 1]["end"]
-        if gap > 0.150 and gap > best_gap:
-            best_gap = gap
-            best_gap_idx = i
-    if best_gap_idx != -1:
-        return best_gap_idx
-
-    # اولویت ۲ — مرز معنایی
-    conjunctions = {"and", "but", "or", "so", "because", "that", "which", "when", "while", "if"}
-    for i in range(mid - 2, mid + 3):
-        if 1 <= i < len(words):
-            word_clean = words[i]["word"].strip().lower().strip(",.!?")
-            if word_clean in conjunctions:
-                return i
-
-    # اولویت ۳ — وسط
-    return mid
+    return merged
 
 
-def split_segment(seg: dict, cps_threshold: float) -> list:
+def fix_overlaps(segments: list) -> list:
+    """
+    رفع timestamp‌های overlapping.
+    هیچ subtitle نباید قبل از شروع subtitle قبلی تموم بشه.
+    """
+    fixed = []
+    for i, seg in enumerate(segments):
+        if fixed and seg["start"] < fixed[-1]["end"]:
+            # overlap داریم — start رو تنظیم کن
+            seg = dict(seg)
+            seg["start"] = fixed[-1]["end"] + 0.001
+            if seg["start"] >= seg["end"]:
+                seg["end"] = seg["start"] + 0.5
+        fixed.append(seg)
+    return fixed
+
+
+def split_long_segment(seg: dict, max_chars: int = 80) -> list:
+    """
+    شکستن segment‌های خیلی طولانی به دو بخش.
+    """
     words = seg.get("words", [])
-    if not words:
-        return [seg]
+    text = seg["text"].strip()
 
-    duration = words[-1]["end"] - words[0]["start"]
-    if duration <= 0:
+    if len(text.replace(" ", "")) <= max_chars or len(words) < 4:
         return [seg]
 
-    char_count = len(seg["text"].replace(" ", ""))
-    cps = char_count / duration
+    mid = len(words) // 2
 
-    if cps <= cps_threshold:
-        return [seg]
+    # پیدا کردن بهترین نقطه split
+    best_idx = mid
+    best_gap = -1
+    for i in range(1, len(words)):
+        gap = words[i]["start"] - words[i-1]["end"]
+        if gap > best_gap:
+            best_gap = gap
+            best_idx = i
 
-    split_idx = find_split_point(words)
-    words_part1 = words[:split_idx]
-    words_part2 = words[split_idx:]
+    words1 = words[:best_idx]
+    words2 = words[best_idx:]
 
-    if not words_part1 or not words_part2:
+    if not words1 or not words2:
         return [seg]
 
-    start2 = words_part2[0]["start"]
-    end1 = max(start2 - 0.080, words_part1[-1]["end"])
-
-    if end1 - words_part1[0]["start"] < 0.800:
-        end1 = words_part1[0]["start"] + 0.800
+    start2 = words2[0]["start"]
+    end1 = max(start2 - 0.08, words1[-1]["end"])
 
     seg1 = {
-        "words": words_part1,
-        "text": " ".join(w["word"].strip() for w in words_part1),
-        "start": words_part1[0]["start"],
+        "text": " ".join(w["word"].strip() for w in words1),
+        "start": words1[0]["start"],
         "end": end1,
+        "words": words1,
     }
     seg2 = {
-        "words": words_part2,
-        "text": " ".join(w["word"].strip() for w in words_part2),
+        "text": " ".join(w["word"].strip() for w in words2),
         "start": start2,
-        "end": words_part2[-1]["end"],
+        "end": words2[-1]["end"],
+        "words": words2,
     }
 
-    return split_segment(seg1, cps_threshold) + split_segment(seg2, cps_threshold)
+    return split_long_segment(seg1, max_chars) + split_long_segment(seg2, max_chars)
 
 
 def build_srt(segments: list) -> str:
     lines = []
     index = 1
     for seg in segments:
-        start_tc = format_timestamp(seg["start"])
-        end_tc = format_timestamp(seg["end"])
         text = seg["text"].strip()
         if not text:
             continue
+        start_tc = format_timestamp(seg["start"])
+        end_tc = format_timestamp(seg["end"])
         lines.append(f"{index}\n{start_tc} --> {end_tc}\n{text}")
         index += 1
     return "\n\n".join(lines) + "\n"
@@ -158,13 +145,18 @@ def process(segments: list) -> str:
     if not segments:
         return ""
 
-    cps_threshold = detect_cps_threshold(segments)
-
-    # اول merge کن، بعد split
+    # مرحله ۱ — ادغام segment‌های کوتاه
     merged = merge_segments(segments)
 
-    final_segments = []
-    for seg in merged:
-        final_segments.extend(split_segment(seg, cps_threshold))
+    # مرحله ۲ — رفع overlap
+    fixed = fix_overlaps(merged)
+
+    # مرحله ۳ — شکستن segment‌های خیلی طولانی
+    final = []
+    for seg in fixed:
+        final.extend(split_long_segment(seg))
+
+    # مرحله ۴ — رفع overlap بعد از split
+    final = fix_overlaps(final)
 
-    return build_srt(final_segments)
+    return build_srt(final)