Create pipeline/postprocess.py

pipeline/postprocess.py

@@ -0,0 +1,167 @@
+# ─── Post-Processing Engine ────────────────────────────────────────────────────
+# این لایه خروجی خام Whisper رو به SRT حرفه‌ای تبدیل می‌کنه
+
+def format_timestamp(seconds: float) -> str:
+    """تبدیل ثانیه به فرمت SRT: HH:MM:SS,mmm"""
+    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 detect_cps_threshold(segments: list) -> float:
+    """
+    تشخیص سرعت گفتار از روی کل ویدیو و برگرداندن CPS threshold مناسب.
+    slow/documentary → 17 | normal → 20 | fast/vlog → 25
+    """
+    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 find_split_point(words: list) -> int:
+    """
+    پیدا کردن بهترین نقطه split در لیست کلمات.
+    اولویت ۱: مکث طبیعی (gap > 150ms)
+    اولویت ۲: مرز معنایی
+    اولویت ۳: نزدیک‌ترین نقطه به وسط
+    """
+    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
+
+
+def split_segment(seg: dict, cps_threshold: float) -> list:
+    """
+    بررسی اینکه آیا segment نیاز به split داره.
+    اگه CPS بالاتر از threshold بود → split.
+    """
+    words = seg.get("words", [])
+    if not words:
+        return [seg]
+
+    duration = words[-1]["end"] - words[0]["start"]
+    if duration <= 0:
+        return [seg]
+
+    char_count = len(seg["text"].replace(" ", ""))
+    cps = char_count / duration
+
+    # اگه CPS پایین‌تر از threshold بود → نیازی به split نیست
+    if cps <= cps_threshold:
+        return [seg]
+
+    # split لازمه — پیدا کردن نقطه split
+    split_idx = find_split_point(words)
+
+    words_part1 = words[:split_idx]
+    words_part2 = words[split_idx:]
+
+    if not words_part1 or not words_part2:
+        return [seg]
+
+    # ─── محاسبه timestamp (ترتیب مهمه!) ───────────────────────────────────────
+    # ۱. اول start[subtitle دوم] رو از word timestamp می‌گیریم
+    start2 = words_part2[0]["start"]
+    # ۲. بعد end[subtitle اول] رو مشتق می‌کنیم
+    end1 = max(start2 - 0.080, words_part1[-1]["end"])
+
+    # safety: حداقل 800ms نمایش
+    if end1 - words_part1[0]["start"] < 0.800:
+        end1 = words_part1[0]["start"] + 0.800
+
+    seg1 = {
+        "words": words_part1,
+        "text": " ".join(w["word"].strip() for w in words_part1),
+        "start": words_part1[0]["start"],
+        "end": end1,
+    }
+
+    seg2 = {
+        "words": words_part2,
+        "text": " ".join(w["word"].strip() for w in words_part2),
+        "start": start2,
+        "end": words_part2[-1]["end"],
+    }
+
+    # recursive split اگه هنوز طولانیه
+    return split_segment(seg1, cps_threshold) + split_segment(seg2, cps_threshold)
+
+
+def build_srt(segments: list) -> str:
+    """ساخت فایل SRT نهایی از لیست segment‌ها"""
+    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
+        lines.append(f"{index}\n{start_tc} --> {end_tc}\n{text}")
+        index += 1
+
+    return "\n\n".join(lines) + "\n"
+
+
+def process(segments: list) -> str:
+    """
+    تابع اصلی — خروجی Whisper ر�� می‌گیره و SRT برمی‌گردونه.
+    segments: لیستی از dict با کلیدهای text, words, start, end
+    """
+    if not segments:
+        return ""
+
+    # تشخیص CPS threshold بر اساس سرعت کل ویدیو
+    cps_threshold = detect_cps_threshold(segments)
+
+    # split کردن segment‌های طولانی
+    final_segments = []
+    for seg in segments:
+        final_segments.extend(split_segment(seg, cps_threshold))
+
+    # ساخت SRT نهایی
+    return build_srt(final_segments)