scripts/broadcast/click_homography.py

"""Bước 2 P0 (BRIEF 11/08/2026): Danh click 4 pocket góc → homography metric.

Với mỗi ``shot_XX`` chưa có ``homography.json``: hiện frame ĐẦU của cú,
Danh click 4 pocket GÓC theo đúng quy ước ``table_corners()`` — đi vòng
quanh bàn, 2 click đầu nằm trên cùng một BĂNG NGẮN (vd: dưới-trái →
dưới-phải → trên-phải → trên-trái). Click ngược chiều cũng được —
``orient_corners`` tự đảo (06/08). Điểm click = giao 2 mép băng ở miệng
pocket (góc mặt chơi).

Mặt chơi bàn 9ft theo BRIEF: 2.54 × 1.27 m → truyền ``table_w=1.27,
table_l=2.54`` cho ``TableHomography`` (tái dụng module, KHÔNG sửa src;
mặc định pooltool 0.9906 × 1.9812 KHÔNG dùng ở đây).

Camera có thể đổi giữa các cú → mỗi cú một bộ click riêng; trong một cảnh
coi như camera tĩnh (design §4.3) nên chỉ cần frame đầu.

Phím: click trái = chấm điểm · u = xoá điểm cuối · r = chấm lại từ đầu ·
s = bỏ qua cú (pocket bị che → hỏi Danh chọn cú thay, đừng đoán) ·
Enter/Space = nhận (khi đủ 4 điểm, có preview top-view để mắt kiểm) ·
q/Esc = thoát.

PHẢI chạy foreground có màn hình (Danh thao tác) — đừng chạy từ shell nền
(bẫy BG21). Launcher: run_bb9_click.bat ở D:\\Khoa luan.

    D:\\Khoa luan\\poolcoach-cv-env\\Scripts\\python.exe \
        scripts/broadcast/click_homography.py
"""

from __future__ import annotations

import argparse
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
import numpy as np  # noqa: E402

from poolcoach_cv.homography import TableHomography, orient_corners  # noqa: E402

PILOT_ROOT = ROOT / "datasets" / "bb9_pilot"
# Mặt chơi 9ft chuẩn giải (BRIEF #8): W = chiều NGẮN, L = chiều DÀI.
TABLE_W_M = 1.27
TABLE_L_M = 2.54

WIN = "bb9 click homography"
ZOOM = 4          # magnifier phong to quanh con tro
ZOOM_R = 40       # ban kinh vung goc (px anh goc) dua vao magnifier
HELP = ("Click 4 pocket GOC vong quanh ban, 2 click dau = mot BANG NGAN | "
        "u=undo r=redo s=skip Enter=OK q=thoat")


class ClickState:
    def __init__(self) -> None:
        self.pts: list[tuple[float, float]] = []   # toa do ANH GOC
        self.cursor: tuple[int, int] = (0, 0)      # toa do anh HIEN THI


def _draw(base_disp: np.ndarray, img_full: np.ndarray, st: ClickState,
          scale: float, msg: str) -> np.ndarray:
    vis = base_disp.copy()
    for i, (x, y) in enumerate(st.pts):
        p = (int(round(x * scale)), int(round(y * scale)))
        cv2.circle(vis, p, 5, (0, 255, 255), -1)
        cv2.putText(vis, str(i + 1), (p[0] + 8, p[1] - 8),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 255), 2)
    if len(st.pts) >= 2:
        poly = np.asarray([[int(round(x * scale)), int(round(y * scale))]
                           for x, y in st.pts], np.int32)
        cv2.polylines(vis, [poly], len(st.pts) == 4, (0, 255, 255), 1)

    # Magnifier: crop anh GOC quanh vi tri con tro -> phong to ZOOM lan.
    cx, cy = st.cursor
    fx, fy = int(round(cx / scale)), int(round(cy / scale))
    h, w = img_full.shape[:2]
    x0, x1 = max(0, fx - ZOOM_R), min(w, fx + ZOOM_R)
    y0, y1 = max(0, fy - ZOOM_R), min(h, fy + ZOOM_R)
    if x1 > x0 and y1 > y0:
        mag = cv2.resize(img_full[y0:y1, x0:x1], None, fx=ZOOM, fy=ZOOM,
                         interpolation=cv2.INTER_NEAREST)
        mh, mw = mag.shape[:2]
        ccx, ccy = int((fx - x0) * ZOOM), int((fy - y0) * ZOOM)
        cv2.line(mag, (ccx - 12, ccy), (ccx + 12, ccy), (0, 0, 255), 1)
        cv2.line(mag, (ccx, ccy - 12), (ccx, ccy + 12), (0, 0, 255), 1)
        cv2.rectangle(mag, (0, 0), (mw - 1, mh - 1), (255, 255, 255), 1)
        vis[10:10 + mh, vis.shape[1] - mw - 10:vis.shape[1] - 10] = mag

    cv2.rectangle(vis, (0, 0), (vis.shape[1], 28), (0, 0, 0), -1)
    cv2.putText(vis, msg, (8, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.55,
                (255, 255, 255), 1)
    return vis


def _preview_topview(img_full: np.ndarray, th: TableHomography) -> np.ndarray:
    """Warp frame về top-view metric để mắt kiểm 4 click có ra hình bàn không."""
    px_per_m = 220
    margin = 30
    w_px = int(TABLE_W_M * px_per_m)
    l_px = int(TABLE_L_M * px_per_m)
    # anh_goc -> ban (m) -> pixel top-view (x sang phai, y XUONG nhu anh)
    S = np.array([[px_per_m, 0, margin], [0, px_per_m, margin], [0, 0, 1]],
                 dtype=np.float64)
    M = S @ th.H
    out = cv2.warpPerspective(img_full, M, (w_px + 2 * margin, l_px + 2 * margin))
    cv2.rectangle(out, (margin, margin), (margin + w_px, margin + l_px),
                  (0, 255, 255), 2)
    cv2.putText(out, "Top-view: mat ban phai vua khit khung vang. Enter=OK r=cham lai",
                (8, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.45, (0, 255, 255), 1)
    return out


def click_one(shot_dir: Path, display_h: int) -> str:
    """Trả 'saved' | 'skipped' | 'quit'."""
    frames = sorted((shot_dir / "frames").glob("*.jpg"))
    if not frames:
        print(f"[WARN] {shot_dir.name}: chua co frames/ — chay extract_shots.py truoc.")
        return "skipped"
    img_full = cv2.imread(str(frames[0]))
    if img_full is None:
        print(f"[WARN] {shot_dir.name}: khong doc duoc {frames[0].name}.")
        return "skipped"

    scale = min(1.0, display_h / img_full.shape[0])
    base_disp = cv2.resize(img_full, None, fx=scale, fy=scale,
                           interpolation=cv2.INTER_AREA) if scale < 1.0 else img_full

    st = ClickState()

    def on_mouse(event, x, y, flags, param):
        st.cursor = (x, y)
        if event == cv2.EVENT_LBUTTONDOWN and len(st.pts) < 4:
            st.pts.append((x / scale, y / scale))

    cv2.namedWindow(WIN)
    cv2.setMouseCallback(WIN, on_mouse)
    mode = "click"  # click | preview
    th = None
    while True:
        if mode == "click":
            msg = f"{shot_dir.name} [{len(st.pts)}/4] {HELP}"
            cv2.imshow(WIN, _draw(base_disp, img_full, st, scale, msg))
        else:
            cv2.imshow(WIN, _preview_topview(img_full, th))
        k = cv2.waitKey(16) & 0xFF
        if k in (ord("q"), 27):
            return "quit"
        if k == ord("s"):
            print(f"[SKIP] {shot_dir.name} — pocket bi che? Hoi Danh chon cu thay.")
            return "skipped"
        if k == ord("u") and mode == "click" and st.pts:
            st.pts.pop()
        if k == ord("r"):
            st.pts.clear()
            mode = "click"
        if k in (13, 32):
            if mode == "click" and len(st.pts) == 4:
                oriented = orient_corners(np.asarray(st.pts, dtype=np.float64))
                try:
                    th = TableHomography(oriented, table_w=TABLE_W_M,
                                         table_l=TABLE_L_M)
                except ValueError as e:
                    print(f"[ERROR] {shot_dir.name}: homography loi: {e}")
                    st.pts.clear()
                    continue
                mode = "preview"
            elif mode == "preview":
                data = {
                    "shot": shot_dir.name,
                    "image": frames[0].name,
                    "clicked_px": [list(p) for p in st.pts],
                    "oriented_px": th.corners_px.tolist(),
                    "table_w_m": TABLE_W_M,
                    "table_l_m": TABLE_L_M,
                    "H": th.H.tolist(),
                    "created_at": datetime.now(timezone.utc).isoformat(
                        timespec="seconds"),
                }
                out = shot_dir / "homography.json"
                out.write_text(json.dumps(data, indent=2), encoding="utf-8")
                print(f"[OK] {shot_dir.name}: luu {out.name}")
                return "saved"


def main() -> None:
    ap = argparse.ArgumentParser(description="Click 4 pocket goc -> homography (P0)")
    ap.add_argument("--root", type=Path, default=PILOT_ROOT)
    ap.add_argument("--only", type=int, nargs="*", default=None)
    ap.add_argument("--redo", action="store_true",
                    help="lam lai ca cu da co homography.json")
    ap.add_argument("--display-height", type=int, default=950,
                    help="chieu cao cua so hien thi (px)")
    args = ap.parse_args()

    shots = sorted(d for d in args.root.glob("shot_*") if d.is_dir())
    if args.only:
        want = {f"shot_{i:02d}" for i in args.only}
        shots = [d for d in shots if d.name in want]
    todo = [d for d in shots
            if args.redo or not (d / "homography.json").exists()]
    if not todo:
        print("[DONE] Moi cu deu da co homography.json (dung --redo de lam lai).")
        return

    print(f"{len(todo)} cu can click. {HELP}")
    n_saved = n_skip = 0
    for d in todo:
        r = click_one(d, args.display_height)
        if r == "quit":
            print("[QUIT] Dung theo yeu cau.")
            break
        n_saved += r == "saved"
        n_skip += r == "skipped"
    cv2.destroyAllWindows()
    print(f"[SUMMARY] saved={n_saved} skipped={n_skip} / todo={len(todo)}")


if __name__ == "__main__":
    main()