Obsidian-Knowledge-Base/tools/extract.py

# /// script
# dependencies = ["requests", "silero-vad", "torch"]
# ///
"""Extract timed text from video subtitles; transcribe via ASR; grab frames.

Commands:
  parse       <input> [--outdir <dir>]           extract timed text from subtitles
  transcribe  <video> [--outpath <path>]          transcribe speech via VAD+ASR
  frame       <video> <timestamp> --outpath <p>   extract single frame
"""

import argparse
import os
import re
import subprocess
import sys
import tempfile
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 extract_audio(video_path: Path, out_wav: Path) -> None:
    """Extract 16kHz mono WAV from video."""
    cmd = [
        "ffmpeg", "-y", "-i", str(video_path),
        "-vn", "-ar", "16000", "-ac", "1",
        str(out_wav),
    ]
    subprocess.run(cmd, capture_output=True, check=True)


def read_audio_tensor(wav_path: Path) -> "torch.Tensor":
    """Read WAV file as mono float32 tensor (compat with silero-vad)."""
    import numpy as np
    import torch
    # Use ffmpeg to decode to raw PCM float32
    result = subprocess.run(
        ["ffmpeg", "-y", "-i", str(wav_path),
         "-f", "f32le", "-ac", "1", "-ar", "16000", "-"],
        capture_output=True, check=True,
    )
    arr = np.frombuffer(result.stdout, dtype=np.float32).copy()
    return torch.from_numpy(arr).squeeze()


def _find_model_dir(model_size: str) -> str | None:
    """Find cached faster-whisper model directory."""
    repo = f"Systran/faster-whisper-{model_size}"
    cache = Path(os.path.expanduser("~/.cache/huggingface/hub"))
    repo_dir = cache / f"models--{repo.replace('/', '--')}"
    if not repo_dir.exists():
        return None
    snapshots = repo_dir / "snapshots"
    if not snapshots.exists():
        return None
    snaps = list(snapshots.iterdir())
    return str(snaps[0]) if snaps else None


def cmd_transcribe(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)

    outpath = Path(args.outpath) if args.outpath else None
    model_size = args.model

    model_dir = _find_model_dir(model_size)
    if model_dir is None:
        print(f"[ERR] Model '{model_size}' not cached locally", file=sys.stderr)
        sys.exit(1)

    print(f"[...] Extracting audio from {video.name} ...", file=sys.stderr)
    wav = Path(tempfile.mktemp(suffix=".wav"))
    try:
        extract_audio(video, wav)

        print(f"[...] Loading faster-whisper ({model_size}) ...", file=sys.stderr)
        import os as _os
        _os.environ["HF_HUB_DISABLE_SYMLINKS_WARNING"] = "1"
        from faster_whisper import WhisperModel
        model = WhisperModel(model_dir, device="cpu", compute_type="int8")

        print(f"[...] Transcribing. This may take a few minutes ...", file=sys.stderr)
        segments, _info = model.transcribe(
            str(wav), beam_size=5, vad_filter=True,
            vad_parameters=dict(min_silence_duration_ms=300),
        )

        segs: list[tuple[float, str]] = []
        for seg in segments:
            text = seg.text.strip()
            if text:
                segs.append((seg.start, text))

        print(f"[...] {len(segs)} segments", file=sys.stderr)
        for s, t in segs:
            print(f"  [{sec_to_ts(s)}] {t[:80]}", file=sys.stderr)

    finally:
        wav.unlink(missing_ok=True)

    if not segs:
        print(f"[WARN] No speech transcribed", file=sys.stderr)
        return

    lines = [f"[{sec_to_ts(s)}] {t}" for s, t in segs]
    if outpath:
        outpath.parent.mkdir(parents=True, exist_ok=True)
        outpath.write_text("\n".join(lines), encoding="utf-8")
        print(f"[OK] Transcribed -> {outpath}  ({len(segs)} segments)")
    else:
        print("\n".join(lines))


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("transcribe", help="Transcribe speech via local faster-whisper")
    p.add_argument("video", help="Path to video file")
    p.add_argument("--outpath", help="Output transcript path (default: stdout)")
    p.add_argument("--model", default="base", help="Model size: tiny/base/small/medium/large")

    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 == "transcribe":
        cmd_transcribe(args)
    elif args.command == "frame":
        cmd_frame(args)
    else:
        parser.print_help()
        sys.exit(1)


if __name__ == "__main__":
    main()