feat: add tools/extract.py for subtitle text extraction and frame grabbing

tools/extract.py

@@ -0,0 +1,191 @@
+# /// script
+# dependencies = []
+# ///
+"""Extract timed text from video subtitles; grab single frames.
+
+Commands:
+  parse  <input> [--outdir <dir>]  extract timed text from subtitles
+  frame  <video> <timestamp> --outpath <path>  extract single frame
+"""
+
+import argparse
+import os
+import re
+import subprocess
+import sys
+from pathlib import Path
+
+SRT_PATTERN = re.compile(
+    r"(\d+)\s+(\d{2}:\d{2}:\d{2}[.,]\d{3})\s*-->\s*(\d{2}:\d{2}:\d{2}[.,]\d{3})\s*((?:(?!\n\n).)*)",
+    re.DOTALL,
+)
+ASS_DIALOGUE = re.compile(
+    r"^Dialogue:\s*\d+,(\d:\d{2}:\d{2}\.\d{2}),(\d:\d{2}:\d{2}\.\d{2}),.*?,(.*?)$",
+    re.MULTILINE,
+)
+ASS_TAG = re.compile(r"\{[^}]*\}")
+ASS_LINEBREAK = re.compile(r"\\[Nn]")
+
+
+def find_subtitle(video_path: Path) -> Path | None:
+    """Return matching .srt (prefer Chinese) or .ass for video_path, or None."""
+
+    def _candidates(base: str) -> list[Path]:
+        parent = video_path.parent
+        out: list[Path] = []
+        for f in parent.iterdir():
+            if not f.is_file():
+                continue
+            stem = f.stem
+            if not stem.startswith(base):
+                continue
+            if f.suffix == ".srt" and "_英文" not in stem and "_English" not in stem:
+                out.append(f)
+            elif f.suffix == ".ass":
+                out.append(f)
+        return out
+
+    base = video_path.stem
+    cands = _candidates(base)
+    srt = [c for c in cands if c.suffix == ".srt"]
+    ass = [c for c in cands if c.suffix == ".ass"]
+    chosen = srt[0] if srt else (ass[0] if ass else None)
+
+    if chosen is None:
+        shorter = base.split("-")[0] if "-" in base else base
+        if shorter != base:
+            cands = _candidates(shorter)
+            srt = [c for c in cands if c.suffix == ".srt"]
+            ass = [c for c in cands if c.suffix == ".ass"]
+            chosen = srt[0] if srt else (ass[0] if ass else None)
+    return chosen
+
+
+def _ts_to_sec(ts: str) -> float:
+    ts = ts.replace(",", ".")
+    parts = ts.split(":")
+    h, m = int(parts[0]), int(parts[1])
+    s = float(parts[2]) if len(parts) == 3 else 0.0
+    return h * 3600 + m * 60 + s
+
+
+def sec_to_ts(sec: float) -> str:
+    h = int(sec // 3600)
+    m = int((sec % 3600) // 60)
+    s = sec % 60
+    return f"{h:02d}:{m:02d}:{s:06.3f}"
+
+
+def parse_srt(path: Path) -> list[tuple[float, str]]:
+    text = path.read_text(encoding="utf-8-sig")
+    segs: list[tuple[float, str]] = []
+    for m in SRT_PATTERN.finditer(text):
+        start = _ts_to_sec(m.group(2))
+        raw = m.group(4).strip().replace("\n", " ")
+        if raw:
+            segs.append((start, raw))
+    return segs
+
+
+def parse_ass(path: Path) -> list[tuple[float, str]]:
+    text = path.read_text(encoding="utf-8-sig")
+    segs: list[tuple[float, str]] = []
+    for m in ASS_DIALOGUE.finditer(text):
+        start = _ts_to_sec(m.group(1))
+        raw = m.group(3).strip()
+        raw = ASS_TAG.sub("", raw)
+        raw = ASS_LINEBREAK.sub(" ", raw)
+        if raw:
+            segs.append((start, raw))
+    return segs
+
+
+def collect_videos(input_path: Path) -> list[Path]:
+    if input_path.is_file():
+        return [input_path]
+    return sorted(input_path.rglob("*.mp4"))
+
+
+def cmd_parse(args: argparse.Namespace) -> None:
+    input_path = Path(args.input)
+    outdir = Path(args.outdir)
+    videos = collect_videos(input_path)
+
+    for vp in videos:
+        sub = find_subtitle(vp)
+        if sub is None:
+            print(f"[WARN] No subtitle for: {vp.name}", file=sys.stderr)
+            continue
+        segs: list[tuple[float, str]] = []
+        if sub.suffix == ".srt":
+            segs = parse_srt(sub)
+        elif sub.suffix == ".ass":
+            segs = parse_ass(sub)
+
+        if input_path.is_dir():
+            rel = vp.relative_to(input_path)
+            out = outdir / rel.with_suffix(".txt")
+        else:
+            out = outdir / vp.with_suffix(".txt").name
+        out.parent.mkdir(parents=True, exist_ok=True)
+
+        lines = [f"[{sec_to_ts(s)}] {t}" for s, t in segs]
+        out.write_text("\n".join(lines), encoding="utf-8")
+        print(f"[OK] {vp.name} -> {out}  ({len(segs)} segments)")
+
+
+def parse_ts(s: str) -> float:
+    s = s.strip()
+    if ":" in s:
+        parts = s.split(":")
+        h, m = int(parts[0]), int(parts[1])
+        sec = float(parts[2]) if len(parts) == 3 else 0.0
+        return h * 3600 + m * 60 + sec
+    return float(s)
+
+
+def cmd_frame(args: argparse.Namespace) -> None:
+    video = Path(args.video)
+    if not video.exists():
+        print(f"[ERR] Video not found: {video}", file=sys.stderr)
+        sys.exit(1)
+    ts = parse_ts(args.timestamp)
+    out = Path(args.outpath)
+    out.parent.mkdir(parents=True, exist_ok=True)
+
+    cmd = [
+        "ffmpeg", "-ss", str(ts), "-i", str(video),
+        "-vframes", "1", "-q:v", "2", "-y", str(out),
+    ]
+    result = subprocess.run(cmd, capture_output=True, text=True)
+    if result.returncode != 0:
+        print(f"[ERR] ffmpeg failed: {result.stderr.strip()}", file=sys.stderr)
+        sys.exit(1)
+    print(f"[OK] Frame saved -> {out}")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Video subtitle extractor & frame grabber")
+    sub = parser.add_subparsers(dest="command")
+
+    p = sub.add_parser("parse", help="Extract timed text from subtitles")
+    p.add_argument("input", help="Video file or directory")
+    p.add_argument("--outdir", default="./transcripts", help="Output directory")
+
+    p = sub.add_parser("frame", help="Extract single frame as JPEG")
+    p.add_argument("video", help="Path to video file")
+    p.add_argument("timestamp", help="Timestamp (HH:MM:SS or seconds)")
+    p.add_argument("--outpath", required=True, help="Output image path")
+
+    args = parser.parse_args()
+    if args.command == "parse":
+        cmd_parse(args)
+    elif args.command == "frame":
+        cmd_frame(args)
+    else:
+        parser.print_help()
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()