scripts/cv/demo_scan.py

"""Demo mini end-to-end F2 (Bước 4, BRIEF 05/08/2026): ảnh → bi → toạ độ bàn.

Pipeline: ảnh val → YOLO detect → tâm bbox → homography 4 góc → toạ độ bàn
(hệ pooltool, m) → in bảng + lưu overlay (bbox, toạ độ, khung bàn) vào thư
mục artifact.

4 góc bàn mỗi ảnh: dataset pix2pockets KHÔNG có nhãn góc bàn (chỉ có class
Dot = nút kim cương trên thành gỗ) → dùng 4 góc ĐO TAY bằng mắt trên từng
ảnh (05/08, zoom 2x + lưới toạ độ; sai số ước ±10-15 px). Quy ước thứ tự
góc như `table_corners()`: 2 góc đầu là một BĂNG NGẮN, đi vòng quanh bàn —
(0,0) → (W,0) → (W,L) → (0,L).

Tâm bbox ≈ tâm bi chiếu xuống mặt bàn: gần đúng — bi có độ cao nên tâm ảnh
lệch nhẹ so với điểm chạm nỉ, chấp nhận cho baseline (pix2pockets cũng vậy).

Class Dot bị loại khỏi phép chiếu (nút nằm NGOÀI mặt bàn, trên thành gỗ);
vẫn vẽ bbox mảnh trên overlay để soi detector.

    python scripts/cv/demo_scan.py               # 3 ảnh demo mặc định
    python scripts/cv/demo_scan.py --images <stem-hoặc-tên-file> ...
"""

from __future__ import annotations

import argparse
import sys
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 TABLE_L, TABLE_W, TableHomography  # noqa: E402

VAL_IMAGES = ROOT / "datasets" / "pix2pockets" / "yolo" / "valid" / "images"
DEFAULT_WEIGHTS = ROOT / "runs" / "cv" / "cv_baseline_20260805" / "weights" / "best.pt"
DEFAULT_OUT = ROOT.parent / "cv_baseline_20260805" / "demo"

# 4 góc mặt bàn (pixel, ảnh gốc 1920x1080) đo tay 05/08/2026 — xem docstring.
DEMO_CORNERS: dict[str, list[tuple[float, float]]] = {
    # Kozoom, nhìn xiên: trái → dưới (băng ngắn đầu bàn), phải, trên
    "103_png.rf.8837c4ba6e91c59a634fc8cd30cd204b.jpg":
        [(340, 650), (690, 785), (1550, 495), (1030, 250)],
    # Predator, nhìn từ sau băng ngắn: dưới-trái → dưới-phải, trên-phải, trên-trái
    "114_png.rf.98b2144c25b9f48816abd7edb00f365c.jpg":
        [(420, 790), (1555, 795), (1420, 252), (612, 255)],
    "130_png.rf.21f92398160a76a64b948884e4c8103c.jpg":
        [(335, 795), (1545, 790), (1310, 258), (645, 255)],
}

BALL_CLASSES = {"Black", "Cue", "Solid", "Striped"}  # Dot = nút thành gỗ, bỏ chiếu


def _resolve_image(token: str) -> Path:
    p = Path(token)
    if p.exists():
        return p
    cands = sorted(VAL_IMAGES.glob(token + "*"))
    if len(cands) == 1:
        return cands[0]
    sys.exit(f"[ERROR] Cannot resolve image '{token}' ({len(cands)} matches in val).")


def scan_one(model, img_path: Path, corners_px, out_dir: Path, conf: float) -> dict:
    im = cv2.imread(str(img_path))
    if im is None:
        sys.exit(f"[ERROR] Cannot read {img_path}")
    th = TableHomography(np.asarray(corners_px, dtype=float))
    res = model.predict(im, conf=conf, imgsz=640, verbose=False)[0]
    names = res.names

    print(f"\n=== {img_path.name} ===")
    print(f"  {'class':>8} {'conf':>5} {'px':>14} {'table_xy_m':>16} in?")
    n_ball = n_in = 0
    vis = im.copy()
    cv2.polylines(vis, [np.asarray(corners_px, int).reshape(-1, 1, 2)], True, (0, 255, 255), 2)
    for i, (x0, y0) in enumerate(np.asarray(corners_px, int)):
        cv2.circle(vis, (x0, y0), 6, (0, 255, 255), -1)
        cv2.putText(vis, f"C{i}", (x0 + 8, y0 - 8), 0, 0.6, (0, 255, 255), 2)

    for b in res.boxes:
        cls = names[int(b.cls)]
        x1, y1, x2, y2 = (float(v) for v in b.xyxy[0])
        cx, cy = (x1 + x2) / 2.0, (y1 + y2) / 2.0
        if cls not in BALL_CLASSES:
            cv2.rectangle(vis, (int(x1), int(y1)), (int(x2), int(y2)), (140, 140, 140), 1)
            continue
        tx, ty = th.px_to_table((cx, cy))
        inside = 0.0 <= tx <= TABLE_W and 0.0 <= ty <= TABLE_L
        n_ball += 1
        n_in += inside
        print(f"  {cls:>8} {float(b.conf):5.2f} ({cx:6.1f},{cy:6.1f}) "
              f"({tx:6.3f},{ty:6.3f})m  {'in' if inside else 'OUT'}")
        color = (0, 200, 0) if inside else (0, 0, 255)
        cv2.rectangle(vis, (int(x1), int(y1)), (int(x2), int(y2)), color, 2)
        cv2.putText(vis, f"{cls} ({tx:.2f},{ty:.2f})", (int(x1), int(y1) - 6),
                    0, 0.55, color, 2)

    out_dir.mkdir(parents=True, exist_ok=True)
    out = out_dir / ("scan_" + img_path.name)
    cv2.imwrite(str(out), vis)
    print(f"  balls in-table: {n_in}/{n_ball}  overlay: {out}")
    return {"image": img_path.name, "balls": n_ball, "in_table": n_in}


def main() -> None:
    ap = argparse.ArgumentParser(description="Mini end-to-end scan demo (detect + homography)")
    ap.add_argument("--weights", type=Path, default=DEFAULT_WEIGHTS)
    ap.add_argument("--images", nargs="+", default=sorted(DEMO_CORNERS),
                    help="image files or val-image name prefixes (need corners in DEMO_CORNERS)")
    ap.add_argument("--out", type=Path, default=DEFAULT_OUT)
    ap.add_argument("--conf", type=float, default=0.25)
    args = ap.parse_args()

    from ultralytics import YOLO

    model = YOLO(str(args.weights))
    stats = []
    for token in args.images:
        img = _resolve_image(token)
        if img.name not in DEMO_CORNERS:
            sys.exit(f"[ERROR] No hand-measured corners for {img.name} - add to DEMO_CORNERS.")
        stats.append(scan_one(model, img, DEMO_CORNERS[img.name], args.out, args.conf))

    total_b = sum(s["balls"] for s in stats)
    total_in = sum(s["in_table"] for s in stats)
    print(f"\n[G4] {len(stats)} images, balls in-table {total_in}/{total_b}")


if __name__ == "__main__":
    main()