Upload postprocess.py

pipeline/postprocess.py

@@ -0,0 +1,253 @@
+# ─── Post-Processing Engine v9 ─────────────────────────────────────────────────
+
+MIN_DISPLAY     = 1.5
+MAX_DISPLAY     = 7.0
+MERGE_MAX_CHARS = 72
+MAX_GAP         = 1.5
+MAX_LINE_CHARS  = 60    # threshold برای split به دو subtitle جداگانه
+WRAP_CHARS      = 42    # استاندارد Netflix برای wrap به دو خط
+CPS_LIMIT       = 17    # حداکثر کاراکتر قابل‌خواندن در ثانیه - برای تشخیص جمله‌ی «به‌سرعت رد شده»
+
+
+def format_timestamp(seconds: float) -> str:
+    ms = round((seconds % 1) * 1000)
+    s  = int(seconds) % 60
+    m  = int(seconds // 60) % 60
+    h  = int(seconds // 3600)
+    return f"{h:02d}:{m:02d}:{s:02d},{ms:03d}"
+
+
+def merge_segments(segments: list) -> list:
+    if not segments:
+        return []
+
+    merged  = []
+    current = None
+
+    for seg in segments:
+        if current is None:
+            current = {
+                "text":  seg["text"].strip(),
+                "start": seg["start"],
+                "end":   seg["end"],
+                "words": list(seg.get("words", [])),
+                "speaker": seg.get("speaker"),
+            }
+            continue
+
+        gap               = seg["start"] - current["end"]
+        combined_text     = current["text"] + " " + seg["text"].strip()
+        combined_duration = seg["end"] - current["start"]
+        ends_sentence     = current["text"].rstrip().endswith((".", "!", "?"))
+        current_duration  = current["end"] - current["start"]
+        is_short          = current_duration < 1.5
+        same_speaker      = current.get("speaker") is None or current.get("speaker") == seg.get("speaker")
+        too_fast          = (len(current["text"]) / current_duration) > CPS_LIMIT if current_duration > 0 else True
+        fits_dialogue     = len(current["text"]) <= WRAP_CHARS and len(seg["text"].strip()) <= WRAP_CHARS
+        already_dialogue  = current.get("is_dialogue_block", False)
+
+        should_merge = (
+            not ends_sentence
+            and combined_duration <= MAX_DISPLAY
+            and same_speaker
+            and gap <= MAX_GAP
+            and (
+                is_short
+                or len(combined_text) <= MERGE_MAX_CHARS
+            )
+        )
+
+        should_pair_as_dialogue = (
+            not already_dialogue
+            and not same_speaker
+            and too_fast
+            and gap <= MAX_GAP
+            and fits_dialogue
+        )
+
+        if should_merge:
+            current["text"]   = combined_text
+            current["end"]    = seg["end"]
+            current["words"] += list(seg.get("words", []))
+        elif should_pair_as_dialogue:
+            current["text"]   = f"- {current['text']}\n- {seg['text'].strip()}"
+            current["end"]    = seg["end"]
+            current["words"] += list(seg.get("words", []))
+            current["is_dialogue_block"] = True
+        else:
+            merged.append(current)
+            current = {
+                "text":  seg["text"].strip(),
+                "start": seg["start"],
+                "end":   seg["end"],
+                "words": list(seg.get("words", [])),
+                "speaker": seg.get("speaker"),
+            }
+
+    if current:
+        merged.append(current)
+
+    return merged
+
+
+def remove_duplicates(segments: list) -> list:
+    if not segments:
+        return []
+    result = [segments[0]]
+    for seg in segments[1:]:
+        if seg["text"].strip() != result[-1]["text"].strip():
+            result.append(seg)
+    return result
+
+
+def fix_overlaps(segments: list) -> list:
+    fixed = []
+    for seg in segments:
+        seg = dict(seg)
+        if fixed and seg["start"] < fixed[-1]["end"]:
+            seg["start"] = fixed[-1]["end"] + 0.001
+        if seg["start"] >= seg["end"]:
+            seg["end"] = seg["start"] + MIN_DISPLAY
+        fixed.append(seg)
+    return fixed
+
+
+def apply_min_display(segments: list) -> list:
+    result = []
+    for i, seg in enumerate(segments):
+        seg = dict(seg)
+        duration = seg["end"] - seg["start"]
+        if duration < MIN_DISPLAY:
+            if i + 1 < len(segments):
+                next_start = segments[i + 1]["start"]
+                max_end    = next_start - 0.080
+            else:
+                max_end = seg["end"] + MAX_DISPLAY
+            new_end    = seg["start"] + MIN_DISPLAY
+            seg["end"] = min(new_end, max(max_end, seg["end"]))
+        result.append(seg)
+    return result
+
+
+def split_long_segment(seg: dict) -> list:
+    """
+    فقط اگه punctuation نزدیک وسط جمله بود، به دو subtitle جداگانه بشکن.
+    اگه نقطه مناسبی نبود، نشکن — wrap_text در build_srt دو خطش می‌کنه.
+    بلوک‌های دیالوگِ دو-گوینده‌ای (is_dialogue_block) هیچ‌وقت شکسته نمی‌شن.
+    """
+    if seg.get("is_dialogue_block"):
+        return [seg]
+
+    words = seg.get("words", [])
+    text  = seg["text"].strip()
+
+    if len(text) <= MAX_LINE_CHARS or len(words) < 4:
+        return [seg]
+
+    mid         = len(words) // 2
+    punct_chars = {'.', ',', '!', '?', ';', ':'}
+    best_punct  = -1
+    best_dist   = float('inf')
+
+    for i in range(1, len(words) - 1):
+        word = words[i]["word"].strip()
+        if word and word[-1] in punct_chars:
+            dist = abs(i - mid)
+            if dist < best_dist:
+                best_dist  = dist
+                best_punct = i + 1
+
+    # اگه punctuation نزدیک وسط نبود → نشکن، wrap_text کافیه
+    if best_punct == -1 or best_dist > len(words) // 3:
+        return [seg]
+
+    words1 = words[:best_punct]
+    words2 = words[best_punct:]
+
+    if len(words1) < 3 or len(words2) < 3:
+        return [seg]
+
+    start2 = words2[0]["start"]
+    end1   = max(start2 - 0.080, words1[-1]["end"])
+
+    if end1 - words1[0]["start"] < MIN_DISPLAY:
+        end1 = words1[0]["start"] + MIN_DISPLAY
+        if end1 >= start2:
+            end1 = start2 - 0.040
+
+    seg1 = {
+        "text":  " ".join(w["word"].strip() for w in words1),
+        "start": words1[0]["start"],
+        "end":   end1,
+        "words": words1,
+    }
+    seg2 = {
+        "text":  " ".join(w["word"].strip() for w in words2),
+        "start": start2,
+        "end":   words2[-1]["end"],
+        "words": words2,
+    }
+
+    return split_long_segment(seg1) + split_long_segment(seg2)
+
+
+def wrap_text(text: str) -> str:
+    """اگه متن بیشتر از WRAP_CHARS بود، به دو خط تقسیم کن.
+    اگه متن از قبل یه بلوک دیالوگِ دو-گوینده‌ای (با - و خط جدید) بود، دست‌نخورده برگردون."""
+    if text.startswith("- ") and "\n" in text:
+        return text
+
+    if len(text) <= WRAP_CHARS:
+        return text
+
+    words      = text.split()
+    best_break = -1
+    best_diff  = float('inf')
+    current_len = 0
+
+    for i, word in enumerate(words):
+        current_len += len(word) + (1 if i > 0 else 0)
+        diff = abs(current_len - len(text) / 2)
+        if diff < best_diff:
+            best_diff  = diff
+            best_break = i + 1
+
+    if best_break <= 0 or best_break >= len(words):
+        return text
+
+    line1 = " ".join(words[:best_break])
+    line2 = " ".join(words[best_break:])
+    return f"{line1}\n{line2}"
+
+
+def build_srt(segments: list) -> str:
+    lines = []
+    index = 1
+    for seg in segments:
+        text = seg["text"].strip()
+        if not text:
+            continue
+        start_tc = format_timestamp(seg["start"])
+        end_tc   = format_timestamp(seg["end"])
+        wrapped  = wrap_text(text)
+        lines.append(f"{index}\n{start_tc} --> {end_tc}\n{wrapped}")
+        index += 1
+    return "\n\n".join(lines) + "\n"
+
+
+def process(segments: list) -> str:
+    if not segments:
+        return ""
+
+    merged = merge_segments(segments)
+    merged = remove_duplicates(merged)
+    fixed  = fix_overlaps(merged)
+
+    split = []
+    for seg in fixed:
+        split.extend(split_long_segment(seg))
+
+    split = fix_overlaps(split)
+    final = apply_min_display(split)
+
+    return build_srt(final)