scripts/broadcast/extract_shots.py

"""Bước 1 P0 (BRIEF 11/08/2026): cắt frame từng cú từ video broadcast.

Đọc ``datasets/bb9_pilot/shots.json`` (Danh chọn tay 10 cú full-table), cắt
frame bằng OpenCV theo [start, end] của mỗi cú, GIỮ TIMESTAMP THẬT của từng
frame (``CAP_PROP_POS_MSEC`` — PTS container, không phải số frame × 1/fps):
backend §6 sau này ăn RoPE theo timestamp thật để chịu VFR YouTube, nên lưu
ngay từ P0 khỏi phải trích lại (BRIEF bối cảnh #7).

Frame ra ``datasets/bb9_pilot/shot_XX/frames/*.jpg`` + ``frames_meta.csv``
(frame_file, frame_idx_video, t_video_s) + ``meta.json`` per shot. Thư mục
``datasets/`` nằm ngoài git (repo public, bản quyền video) — chỉ commit
script + shots.json KHÔNG commit frame.

``shots.json`` (Danh đưa timestamp, format tự quyết P0):
    {"video": "pFV7sdBJ42M.mp4",            # tuỳ chọn, mặc định dò *.mp4
     "shots": [{"id": 1, "start": "1:23:45.5", "end": "1:23:58",
                "note": "ghi chu tuy y"}, ...]}
start/end nhận giây (số) hoặc chuỗi "HH:MM:SS(.f)" / "MM:SS(.f)".

Chạy (venv CV — chỉ cần cv2/numpy, nhưng nếp P0 dồn hết vào cv-env):
    D:\\Khoa luan\\poolcoach-cv-env\\Scripts\\python.exe \
        scripts/broadcast/extract_shots.py
"""

from __future__ import annotations

import argparse
import csv
import json
import sys
from datetime import datetime, timezone
from pathlib import Path

ROOT = Path(__file__).resolve().parents[2]  # poolcoach-rl/
sys.path.insert(0, str(ROOT / "src"))

import cv2  # noqa: E402

PILOT_ROOT = ROOT / "datasets" / "bb9_pilot"
# Cảnh báo VFR: dt giữa 2 frame lệch quá 25% so với 1/fps danh định thì đếm.
DT_JITTER_TOL = 0.25


def parse_ts(v) -> float:
    """Giây (float) hoặc "HH:MM:SS(.f)" / "MM:SS(.f)" -> giây."""
    if isinstance(v, (int, float)):
        return float(v)
    parts = str(v).strip().split(":")
    if not 1 <= len(parts) <= 3 or any(p == "" for p in parts):
        raise ValueError(f"Timestamp khong hop le: {v!r}")
    sec = 0.0
    for p in parts:
        sec = sec * 60.0 + float(p)
    return sec


def extract_shot(cap: cv2.VideoCapture, shot: dict, out_dir: Path,
                 fps_nominal: float, jpg_quality: int) -> dict:
    """Cắt 1 cú: seek tới start, đọc tới end, lưu jpg + timestamp thật."""
    t0, t1 = parse_ts(shot["start"]), parse_ts(shot["end"])
    if t1 <= t0:
        raise ValueError(f"Shot {shot['id']}: end ({t1}) <= start ({t0}).")

    frames_dir = out_dir / "frames"
    frames_dir.mkdir(parents=True, exist_ok=True)

    cap.set(cv2.CAP_PROP_POS_MSEC, t0 * 1000.0)
    rows: list[tuple[str, int, float]] = []
    n_jitter = 0
    prev_t = None
    while True:
        ok = cap.grab()
        if not ok:
            break
        t = cap.get(cv2.CAP_PROP_POS_MSEC) / 1000.0
        # Seek theo msec có thể rơi về keyframe trước start — bỏ qua tới t0.
        if t < t0 - 1e-6:
            continue
        if t > t1 + 1e-6:
            break
        ok, frame = cap.retrieve()
        if not ok:
            break
        idx_video = int(cap.get(cv2.CAP_PROP_POS_FRAMES)) - 1
        name = f"{len(rows):05d}.jpg"
        cv2.imwrite(str(frames_dir / name), frame,
                    [cv2.IMWRITE_JPEG_QUALITY, jpg_quality])
        rows.append((name, idx_video, round(t, 6)))
        if prev_t is not None and fps_nominal > 0:
            dt_nom = 1.0 / fps_nominal
            if abs((t - prev_t) - dt_nom) > DT_JITTER_TOL * dt_nom:
                n_jitter += 1
        prev_t = t

    if not rows:
        raise RuntimeError(
            f"Shot {shot['id']}: khong doc duoc frame nao trong "
            f"[{t0}, {t1}]s — kiem lai timestamp/video.")

    with open(out_dir / "frames_meta.csv", "w", newline="", encoding="utf-8") as f:
        w = csv.writer(f)
        w.writerow(["frame_file", "frame_idx_video", "t_video_s"])
        w.writerows(rows)

    meta = {
        "shot_id": shot["id"],
        "note": shot.get("note", ""),
        "start_requested_s": t0,
        "end_requested_s": t1,
        "n_frames": len(rows),
        "t_first_s": rows[0][2],
        "t_last_s": rows[-1][2],
        "fps_nominal": fps_nominal,
        "n_dt_jitter": n_jitter,
        "extracted_at": datetime.now(timezone.utc).isoformat(timespec="seconds"),
    }
    with open(out_dir / "meta.json", "w", encoding="utf-8") as f:
        json.dump(meta, f, indent=2, ensure_ascii=False)
    return meta


def main() -> None:
    ap = argparse.ArgumentParser(description="Cat frame tung cu theo shots.json (P0)")
    ap.add_argument("--root", type=Path, default=PILOT_ROOT)
    ap.add_argument("--shots", type=Path, default=None,
                    help="mac dinh <root>/shots.json")
    ap.add_argument("--video", type=Path, default=None,
                    help="mac dinh: 'video' trong shots.json, roi *.mp4 duy nhat trong root")
    ap.add_argument("--only", type=int, nargs="*", default=None,
                    help="chi cat cac shot id nay")
    ap.add_argument("--jpg-quality", type=int, default=95)
    args = ap.parse_args()

    shots_path = args.shots or args.root / "shots.json"
    spec = json.loads(shots_path.read_text(encoding="utf-8"))
    shots = spec["shots"]
    if args.only:
        shots = [s for s in shots if s["id"] in set(args.only)]
    if not shots:
        sys.exit("[ERROR] Khong co shot nao de cat.")

    video = args.video
    if video is None and spec.get("video"):
        video = args.root / spec["video"]
    if video is None:
        mp4s = sorted(args.root.glob("*.mp4"))
        if len(mp4s) != 1:
            sys.exit(f"[ERROR] Can --video: thay {len(mp4s)} file *.mp4 trong {args.root}")
        video = mp4s[0]

    cap = cv2.VideoCapture(str(video))
    if not cap.isOpened():
        sys.exit(f"[ERROR] Khong mo duoc video: {video}")
    fps_nominal = cap.get(cv2.CAP_PROP_FPS)
    print(f"video: {video.name}  fps_nominal={fps_nominal:.3f}")

    summaries = []
    for shot in shots:
        out_dir = args.root / f"shot_{shot['id']:02d}"
        m = extract_shot(cap, shot, out_dir, fps_nominal, args.jpg_quality)
        summaries.append(m)
        print(f"shot_{shot['id']:02d}: {m['n_frames']} frames  "
              f"t=[{m['t_first_s']:.3f}, {m['t_last_s']:.3f}]s  "
              f"jitter={m['n_dt_jitter']}")
    cap.release()

    total = sum(m["n_frames"] for m in summaries)
    jit = sum(m["n_dt_jitter"] for m in summaries)
    print(f"\n[DONE] {len(summaries)} shots, {total} frames, dt-jitter {jit} frame-gaps")
    if jit:
        print("[WARN] Co dt jitter — video VFR that, timestamp da luu la PTS container.")


if __name__ == "__main__":
    main()