Update pipeline/postprocess.py

pipeline/postprocess.py

@@ -1,8 +1,6 @@
-# ─── Post-Processing Engine ────────────────────────────────────────────────────
-# این لایه خروجی خام Whisper رو به SRT حرفه‌ای تبدیل می‌کنه
+# ─── Post-Processing Engine v2 ─────────────────────────────────────────────────
 
 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
@@ -11,13 +9,8 @@ def format_timestamp(seconds: float) -> str:
 
 
 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:
@@ -26,12 +19,9 @@ def detect_cps_threshold(segments: list) -> float:
         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:
@@ -40,13 +30,47 @@ def detect_cps_threshold(segments: list) -> float:
         return 25.0
 
 
-def find_split_point(words: list) -> int:
+def merge_segments(segments: list, max_chars: int = 80, max_gap: float = 1.5) -> list:
     """
-    پیدا کردن بهترین نقطه split در لیست کلمات.
-    اولویت ۱: مکث طبیعی (gap > 150ms)
-    اولویت ۲: مرز معنایی
-    اولویت ۳: نزدیک‌ترین نقطه به وسط
+    segment‌های کوتاه رو با هم ادغام کن تا جمله‌های کامل بسازه.
+    max_chars: حداکثر کاراکتر برای یه subtitle
+    max_gap: اگه gap بین دو segment بیشتر از این بود، merge نکن
     """
+    if not segments:
+        return []
+
+    merged = []
+    current = segments[0].copy()
+    current["words"] = list(current.get("words", []))
+
+    for seg in segments[1:]:
+        current_text = current["text"].strip()
+        next_text = seg["text"].strip()
+        combined_chars = len((current_text + " " + next_text).replace(" ", ""))
+
+        # 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
+            current["end"] = seg["end"]
+            current["words"].extend(seg.get("words", []))
+        else:
+            merged.append(current)
+            current = seg.copy()
+            current["words"] = list(current.get("words", []))
+
+    merged.append(current)
+    return merged
+
+
+def find_split_point(words: list) -> int:
     mid = len(words) // 2
 
     # اولویت ۱ — مکث طبیعی (gap > 150ms)
@@ -57,11 +81,10 @@ def find_split_point(words: list) -> int:
         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):
@@ -69,15 +92,11 @@ def find_split_point(words: list) -> int:
             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]
@@ -89,26 +108,19 @@ def split_segment(seg: dict, cps_threshold: float) -> list:
     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
 
@@ -118,7 +130,6 @@ def split_segment(seg: dict, cps_threshold: float) -> list:
         "start": words_part1[0]["start"],
         "end": end1,
     }
-
     seg2 = {
         "words": words_part2,
         "text": " ".join(w["word"].strip() for w in words_part2),
@@ -126,15 +137,12 @@ def split_segment(seg: dict, cps_threshold: float) -> list:
         "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"])
@@ -143,25 +151,20 @@ def build_srt(segments: list) -> str:
             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‌های طولانی
+    # اول merge کن، بعد split
+    merged = merge_segments(segments)
+
     final_segments = []
-    for seg in segments:
+    for seg in merged:
         final_segments.extend(split_segment(seg, cps_threshold))
 
-    # ساخت SRT نهایی
     return build_srt(final_segments)