vrsub/scripts/extract_reference_srt.py

"""从烧录字幕视频自动提取**参考时间轴**（供时间对齐集成测试使用）。

用户选择的参考数据途径：烧录字幕视频自动提取——复用仓库自带的生产链路
（frame-extract 抽帧 → subtitle-ocr 逐帧 OCR → 汇总 SRT），把烧录字幕的
出现/消失帧时间换算成秒，产出严格符合测试契约的参考字幕：

    tests/shared/data/alignment/<name>.reference.srt  （标准 SRT）

为什么烧录字幕的时间是可信的 ground truth：

- 视频里烧录的字幕是**画面的一部分**，它在哪一帧出现/消失是客观事实；
  frame-extract 按帧号取帧，时间 = 帧号 / fps，零累计偏差；
- subtitle-ocr 对逐帧 OCR 后合并相同字幕，结束 = 最后可见帧 + 采样间隔，
  与视频烧录时间对齐（见 nodes/subtitle_ocr.py 的 _assemble_srt）。

因此提取出的参考时间轴可直接作为"说话真实发生的时间"与 whisper 产物比对。
**提取结果请人工核对无误后再使用**：脚本只负责自动化，不负责保证正确。

用法（在 vrsub 根目录，需 Ollama glm-ocr 服务可达、有 ffmpeg）：

    uv run python scripts/extract_reference_srt.py <视频路径> [--out tests/shared/data/alignment <name>] [--interval 0.5] [--crop 0,0.75,1,0.25]

- --interval：抽帧间隔秒（默认 0.5；字幕时长 2s 时约取 4 帧，够稳定合并）
- --crop：字幕区域相对比例 x,y,w,h（默认底部 1/4，与生产默认一致）
- 输出文件名默认 <视频名>.reference.srt；用 --out-dir/--name 可指定目录与
  前缀，保证与 tests/realdata_contract.alignment_candidates() 的发现规则一致
  （<name>.<ext> 与 <name>.reference.srt 同目录同名）。
"""

from __future__ import annotations

import argparse
import json
import sys
import urllib.request
from pathlib import Path

from wov_sdk.models import InvokeRequest

PROJECT_ROOT = Path(__file__).resolve().parent.parent

# 与 tests/test_integration_subtitle_ocr.py、tests/test_integration_vlm.py 相同约定。
OLLAMA_HOST = "http://192.168.123.70:11434"
MODEL = "glm-ocr:latest"


def _ollama_reachable(host: str = OLLAMA_HOST, model: str = MODEL) -> bool:
    """探测 Ollama 服务与目标模型是否可用。"""
    try:
        req = urllib.request.Request(
            f"{host}/api/show",
            data=('{"model": "%s"}' % model).encode(),
            headers={"Content-Type": "application/json"},
            method="POST",
        )
        with urllib.request.urlopen(req, timeout=5) as resp:
            return resp.status == 200
    except (urllib.error.URLError, OSError):
        return False


def _register_nodes() -> None:
    """注册全部内置节点，供 subtitle-ocr 内部调 vlm-ocr 使用。"""
    import sys as _sys

    _sys.path.insert(0, str(PROJECT_ROOT))
    from wov_app import registry

    registry.register_all()


def _parse_crop(raw: str) -> list[float]:
    """解析 crop 参数：'x,y,w,h' -> [x, y, w, h]。"""
    parts = [float(p) for p in raw.split(",")]
    if len(parts) != 4:
        raise SystemExit(f"crop 参数格式错误，应为 x,y,w,h，得到: {raw}")
    return parts


def extract_reference_srt(
    video: Path,
    out_path: Path,
    interval_seconds: float = 0.5,
    crop: list[float] | None = None,
    tmp_root: Path | None = None,
) -> Path:
    """跑 frame-extract → subtitle-ocr，把汇总 SRT 写为参考字幕文件。

    返回写出的参考 SRT 路径。crop 缺省用画面底部 1/4（[0, 0.75, 1, 0.25]）。
    tmp_root 供测试注入临时目录，缺省用视频同名临时目录（用完清理）。
    """
    from nodes.frame_extract import invoke as frame_invoke
    from nodes.subtitle_ocr import invoke as ocr_invoke

    if crop is None:
        crop = [0.0, 0.75, 1.0, 0.25]
    tmp = tmp_root or (video.parent / f"_ref_{video.stem}")
    tmp.mkdir(parents=True, exist_ok=True)

    # ① 抽帧：按帧号精确取帧（时间 = 帧号/fps，无累计偏差）。
    frames_resp = frame_invoke(
        InvokeRequest(
            run_id=f"ref_{video.stem}",
            node_instance_id="",
            inputs={"video_uri": str(video)},
            params={"interval_seconds": interval_seconds, "crop": crop},
            output_dir=str(tmp / "frames"),
        )
    )
    if frames_resp.status != "completed":
        raise SystemExit(f"抽帧失败: {frames_resp.error}")

    # ② 逐帧 OCR → 汇总 SRT（消失时间 = 最后可见帧 + 采样间隔）。
    ocr_resp = ocr_invoke(
        InvokeRequest(
            run_id=f"ref_{video.stem}",
            node_instance_id="",
            inputs={"frames_manifest": str(frames_resp.outputs["frames_manifest"])},
            params={"model": MODEL, "ollama_host": OLLAMA_HOST},
            output_dir=str(tmp / "ocr"),
        )
    )
    if ocr_resp.status != "completed":
        raise SystemExit(f"OCR 失败: {ocr_resp.error}")

    srt = Path(ocr_resp.outputs["srt_uri"])
    out_path.parent.mkdir(parents=True, exist_ok=True)
    out_path.write_text(srt.read_text(encoding="utf-8"), encoding="utf-8")
    return out_path


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="从烧录字幕视频自动提取参考时间轴")
    parser.add_argument("video", type=Path, help="烧录字幕视频路径（mp4 等）")
    parser.add_argument("--out-dir", type=Path, default=PROJECT_ROOT / "tests" / "shared" / "data" / "alignment")
    parser.add_argument("--name", type=str, default=None, help="输出名前缀（默认=视频文件名）")
    parser.add_argument("--interval", type=float, default=0.5, help="抽帧间隔秒（默认 0.5）")
    parser.add_argument("--crop", type=str, default="0,0.75,1,0.25", help="字幕区域 x,y,w,h")
    args = parser.parse_args(argv)

    video = args.video.resolve()
    if not video.is_file():
        print(f"视频不存在: {video}")
        return 1
    if not _ollama_reachable():
        print(f"Ollama 不可用（{OLLAMA_HOST} / {MODEL}），无法提取参考字幕。")
        return 2

    _register_nodes()
    name = args.name or video.stem
    out_path = args.out_dir / f"{name}.reference.srt"
    print(f"提取参考字幕 -> {out_path}", flush=True)
    extract_reference_srt(
        video,
        out_path,
        interval_seconds=args.interval,
        crop=_parse_crop(args.crop),
    )
    print("完成。请人工核对参考时间轴后再用于时间对齐测试。", flush=True)
    return 0


if __name__ == "__main__":
    sys.exit(main())