Obsidian-Knowledge-Base/tools/extract.py

# /// 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()