Upload postprocess.py

pipeline/postprocess.py

@@ -0,0 +1,341 @@
+# ─── Post-Processing Engine v10 ────────────────────────────────────────────
+
+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    # حداکثر کاراکتر قابل‌خواندن در ثانیه - برای تشخیص جمله‌ی «به‌سرعت رد شده»
+
+# کلمه‌های کوچیک و به‌هم‌چسبیده‌ای که نباید تنها و جدا از کلمه‌ی بعدشون
+# ته یه segment (چه raw، چه بعد از split) بمونن.
+# هم توی split_long_segment استفاده میشه، هم توی rebalance_boundaries.
+WEAK_ENDINGS = {
+    'in', 'my', 'your', 'his', 'her', 'our', 'their', 'its', 'the', 'a', 'an',
+    'of', 'to', 'at', 'on', 'by', 'for', 'with', 'is', 'are', 'was', 'were',
+    'and', 'or', 'but', 'so', 'if', 'that', 'this', 'we', 'you', 'he', 'she', 'they',
+}
+
+
+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 rebalance_boundaries(segments: list) -> list:
+    """
+    مشکلی که این تابع حل می‌کنه: خودِ Whisper بعضی‌وقتا وسط یه عبارت به‌هم‌چسبیده
+    (مثل «in my | opinion» یا «of this | video») مرز segment می‌ذاره — نه سر یه
+    مکث طبیعی. چون این دو تیکه معمولاً segment‌های RAW جدا از هم می‌مونن (مثلاً چون
+    مجموع مدت‌زمانشون از MAX_DISPLAY بیشتره و merge_segments نمی‌تونه merge‌شون کنه)،
+    این مرز اشتباه هیچ‌وقت به merge_segments یا split_long_segment نمی‌رسه که اصلاح بشه.
+
+    این تابع، قبل از هر پردازش دیگه‌ای، مرز بین segmentهای پشت‌سرهم رو چک می‌کنه:
+    اگه segment اول با یه یا چند کلمه‌ی «ضعیف» و بدون علامت نگارشی (in/my/this/your/...)
+    تموم بشه، اون کلمه(ها) با همون زمان‌بندی اصلی‌شون به ابتدای segment بعدی منتقل میشن.
+    فقط بین segmentهای هم‌گوینده (یا وقتی گوینده مشخص نیست) و با فاصله‌ی کوتاه انجام میشه.
+    """
+    segments = [dict(s, words=list(s.get("words", []))) for s in segments]
+
+    for i in range(len(segments) - 1):
+        cur = segments[i]
+        nxt = segments[i + 1]
+
+        cur_words = cur["words"]
+        nxt_words = nxt["words"]
+        if not cur_words or not nxt_words:
+            continue
+
+        # فقط بین segmentهای هم‌گوینده جابه‌جا کن (اگه گوینده مشخص نبود، اجازه بده)
+        if cur.get("speaker") is not None and nxt.get("speaker") is not None \
+                and cur.get("speaker") != nxt.get("speaker"):
+            continue
+
+        # فقط اگه فاصله کوتاهه، یعنی احتمالاً گفتار پیوسته‌ست نه یه مکث واقعی
+        gap = nxt_words[0]["start"] - cur_words[-1]["end"]
+        if gap > MAX_GAP:
+            continue
+
+        moved_any = False
+        while len(cur_words) > 1:
+            last     = cur_words[-1]
+            stripped = last["word"].strip()
+            bare     = stripped.strip(".,!?;:").lower()
+            has_punct = bool(stripped) and stripped[-1] in {'.', '!', '?', ',', ';', ':'}
+
+            if bare in WEAK_ENDINGS and not has_punct:
+                nxt_words.insert(0, cur_words.pop())
+                moved_any = True
+            else:
+                break
+
+        # فقط وقتی واقعاً کلمه‌ای جابه‌جا شد، text/start/end رو بازسازی کن —
+        # مرزهایی که دست‌نخورده می‌مونن نباید حتی به‌خاطر گرد‌شدگی جزئی تغییر کنن
+        if moved_any:
+            cur["text"]  = " ".join(w["word"].strip() for w in cur_words)
+            cur["end"]   = cur_words[-1]["end"]
+            nxt["text"]  = " ".join(w["word"].strip() for w in nxt_words)
+            nxt["start"] = nxt_words[0]["start"]
+
+    return segments
+
+
+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]
+
+    punct_chars = {'.', '!', '?', ',', ';', ':'}
+    sentence_enders = {'.', '!', '?'}
+    min_part_chars = 10
+
+    full_text = text
+    mid_char = len(full_text) / 2
+
+    best_idx = -1
+    best_score = float('-inf')
+    char_pos = 0
+
+    for i, w in enumerate(words):
+        word_stripped = w["word"].strip()
+        this_word_end_char = char_pos + len(word_stripped)
+
+        if word_stripped and word_stripped[-1] in punct_chars and 0 < i < len(words) - 1:
+            # چک درست: خودِ کلمه‌ای که این علامت نگارشی روشه، نباید یه کلمه‌ی ضعیف باشه
+            bare_word = word_stripped.strip(".,!?;:").lower()
+            part1_chars = this_word_end_char
+            part2_chars = len(full_text) - this_word_end_char
+
+            if bare_word not in WEAK_ENDINGS and part1_chars >= min_part_chars and part2_chars >= min_part_chars:
+                dist_from_mid = abs(this_word_end_char - mid_char)
+                proximity_score = 1.0 - (dist_from_mid / len(full_text))
+                sentence_end_bonus = 0.5 if word_stripped[-1] in sentence_enders else 0.0
+                score = proximity_score + sentence_end_bonus
+
+                if score > best_score:
+                    best_score = score
+                    best_idx = i + 1
+
+        char_pos = this_word_end_char + 1
+
+    # هیچ نقطه‌ی شکستن مناسبی پیدا نشد → نشکن، wrap_text کافیه
+    if best_idx == -1:
+        return [seg]
+
+    words1 = words[:best_idx]
+    words2 = words[best_idx:]
+
+    if len(words1) < 2 or len(words2) < 2:
+        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 ""
+
+    segments = rebalance_boundaries(segments)
+    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)