"""Bước 3 P0 (BRIEF 11/08/2026): detector zero-shot + track cue ball thô. Chạy trên ``poolcoach-cv-env`` (CUDA). Với mỗi ``shot_XX`` (đã có frames + homography.json): 1. Detector 5-class hiện có (``cv_full_20260805/best.pt``) chạy zero-shot từng frame, imgsz 640 (đúng cấu hình vận hành cv_worker — KHÔNG chỉnh để làm đẹp số; muốn báo thêm ở conf/imgsz khác thì ghi rõ nhãn). Detection lưu từ conf floor 0.10 (chẩn đoán), cột ``at_op`` đánh dấu det đạt conf vận hành 0.2993 — MỌI số gate chỉ dùng det đạt op. 2. Track cue ball thô: nearest-neighbor qua frame, gating vận tốc VMAX = 11 m/s (design §4.4) theo dt TIMESTAMP THẬT + slack nhiễu; mất det thì coast (không tính coverage), tự bắt lại trong bán kính gating nở dần theo thời gian mất dấu. 3. Xuất per-shot: ``detections.csv`` (mọi det), ``track.csv`` (cue ball, toạ độ bàn mm + timestamp), ``overlay.mp4`` + ``contact_sheet.jpg`` (bắt buộc — mắt người kiểm track, số coverage suông không tin được). D:\\Khoa luan\\poolcoach-cv-env\\Scripts\\python.exe \ scripts/broadcast/track_p0.py """ from __future__ import annotations import argparse import csv import json 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 TableHomography # noqa: E402 PILOT_ROOT = ROOT / "datasets" / "bb9_pilot" DEFAULT_WEIGHTS = Path(r"D:\Khoa luan\cv_full_20260805\best.pt") # = cv_worker OP_CONF = 0.2993 # conf vận hành (pick_conf 05/08) — đồng bộ cv_worker CONF_FLOOR = 0.10 # lưu det từ đây cho chẩn đoán; KHÔNG dùng cho số gate IMGSZ = 640 # cùng imgsz train/val/worker BATCH = 16 VMAX = 11.0 # m/s — gating vận tốc (design §4.4) SLACK_M = 0.06 # nhiễu vị trí + tâm bbox lệch tâm bi (~2R) IN_TOL_M = 0.05 # tâm chiếu lệch ra ngoài mặt bàn quá mức này → loại BALL_CLASSES = {"Black", "Cue", "Solid", "Striped"} # Dot = nút thành gỗ TRAIL_N = 90 # số điểm track vẽ đuôi trên overlay def load_frames_meta(shot_dir: Path) -> list[tuple[str, float]]: rows = [] with open(shot_dir / "frames_meta.csv", encoding="utf-8") as f: for r in csv.DictReader(f): rows.append((r["frame_file"], float(r["t_video_s"]))) return rows def frame_diffs(shot_dir: Path, meta: list[tuple[str, float]], th: TableHomography) -> dict[str, float]: """Mean abs diff (gray) vùng bàn giữa frame liên tiếp — lộ frame nhân đôi. Stream YouTube 30fps upconvert từ nguồn 25fps sẽ có ~1/6 frame gần trùng (đo 11/08 trên trận Florida Open); measure_p0 dùng cột này để đề xuất ngưỡng trên chuỗi đã loại frame trùng, có dán nhãn. """ corners = th.corners_px.astype(int) x0, x1 = corners[:, 0].min(), corners[:, 0].max() y0, y1 = corners[:, 1].min(), corners[:, 1].max() diffs: dict[str, float] = {} prev = None for name, _t in meta: im = cv2.imread(str(shot_dir / "frames" / name), cv2.IMREAD_GRAYSCALE) crop = im[max(0, y0):y1, max(0, x0):x1].astype(np.int16) diffs[name] = (float(np.abs(crop - prev).mean()) if prev is not None else -1.0) prev = crop return diffs def detect_shot(model, shot_dir: Path, th: TableHomography, conf_floor: float, op_conf: float, imgsz: int, device) -> list[dict]: """YOLO trên mọi frame của cú → list det dict (đã chiếu toạ độ bàn).""" meta = load_frames_meta(shot_dir) dets: list[dict] = [] for i in range(0, len(meta), BATCH): chunk = meta[i:i + BATCH] imgs = [] for name, _t in chunk: im = cv2.imread(str(shot_dir / "frames" / name)) if im is None: sys.exit(f"[ERROR] Khong doc duoc {shot_dir.name}/frames/{name}") imgs.append(im) results = model.predict(imgs, conf=conf_floor, imgsz=imgsz, verbose=False, device=device) for (name, t), res in zip(chunk, results): names = res.names 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 tx, ty = th.px_to_table((cx, cy)) w, l = th.table_w, th.table_l inside = (-IN_TOL_M <= tx <= w + IN_TOL_M and -IN_TOL_M <= ty <= l + IN_TOL_M) dets.append({ "frame_file": name, "t_s": t, "cls": cls, "conf": float(b.conf), "px_cx": cx, "px_cy": cy, "px_w": x2 - x1, "px_h": y2 - y1, "table_x_m": float(tx), "table_y_m": float(ty), "in_table": int(inside), "at_op": int(float(b.conf) >= op_conf), }) return dets def track_cue(meta: list[tuple[str, float]], dets: list[dict], op_conf: float, diffs: dict[str, float]) -> list[dict]: """NN + gating vận tốc trên det class Cue đạt op-conf, trong mặt bàn.""" by_frame: dict[str, list[dict]] = {} for d in dets: if d["cls"] == "Cue" and d["at_op"] and d["in_table"]: by_frame.setdefault(d["frame_file"], []).append(d) rows = [] last_pos = None # (x, y) bàn, m last_t = None for name, t in meta: cands = by_frame.get(name, []) chosen = None if cands: if last_pos is None: chosen = max(cands, key=lambda d: d["conf"]) else: dt = max(t - last_t, 1e-3) gate = VMAX * dt + SLACK_M best = None for d in cands: dist = float(np.hypot(d["table_x_m"] - last_pos[0], d["table_y_m"] - last_pos[1])) if dist <= gate and (best is None or dist < best[0]): best = (dist, d) chosen = best[1] if best else None if chosen is not None: last_pos = (chosen["table_x_m"], chosen["table_y_m"]) last_t = t rows.append({ "frame_file": name, "t_s": t, "covered": int(chosen is not None), "px_x": chosen["px_cx"] if chosen else "", "px_y": chosen["px_cy"] if chosen else "", "table_x_m": chosen["table_x_m"] if chosen else "", "table_y_m": chosen["table_y_m"] if chosen else "", "conf": chosen["conf"] if chosen else "", "n_cands": len(cands), "gap_s": 0.0 if chosen is not None or last_t is None else round(t - last_t, 4), "img_diff": round(diffs.get(name, -1.0), 4), }) return rows def render_overlay(shot_dir: Path, meta, dets, track_rows, fps: float) -> None: by_frame_dets: dict[str, list[dict]] = {} for d in dets: if d["at_op"]: by_frame_dets.setdefault(d["frame_file"], []).append(d) track_by_frame = {r["frame_file"]: r for r in track_rows} first = cv2.imread(str(shot_dir / "frames" / meta[0][0])) h, w = first.shape[:2] vw = cv2.VideoWriter(str(shot_dir / "overlay.mp4"), cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h)) trail: list[tuple[int, int]] = [] n_cov = 0 sheet_idx = set(np.linspace(0, len(meta) - 1, 12, dtype=int).tolist()) thumbs = [] for i, (name, t) in enumerate(meta): im = cv2.imread(str(shot_dir / "frames" / name)) for d in by_frame_dets.get(name, []): x1 = int(d["px_cx"] - d["px_w"] / 2); y1 = int(d["px_cy"] - d["px_h"] / 2) x2 = int(d["px_cx"] + d["px_w"] / 2); y2 = int(d["px_cy"] + d["px_h"] / 2) col = (180, 180, 180) if d["cls"] != "Cue" else (0, 220, 220) cv2.rectangle(im, (x1, y1), (x2, y2), col, 1) r = track_by_frame[name] if r["covered"]: n_cov += 1 p = (int(r["px_x"]), int(r["px_y"])) trail.append(p) cv2.circle(im, p, 14, (0, 220, 0), 2) if len(trail) >= 2: cv2.polylines(im, [np.asarray(trail[-TRAIL_N:], np.int32)], False, (0, 220, 0), 2) status = "TRACK" if r["covered"] else "NO TRACK" col = (0, 220, 0) if r["covered"] else (0, 0, 255) cv2.rectangle(im, (0, 0), (w, 30), (0, 0, 0), -1) cv2.putText(im, f"{shot_dir.name} t={t:9.3f}s frame {i + 1}/{len(meta)} " f"cov {100.0 * n_cov / (i + 1):5.1f}% [{status}]", (8, 22), cv2.FONT_HERSHEY_SIMPLEX, 0.6, col, 2) vw.write(im) if i in sheet_idx: thumbs.append(cv2.resize(im, (480, 270))) vw.release() rows = [np.hstack(thumbs[r * 4:(r + 1) * 4]) for r in range(3) if len(thumbs[r * 4:(r + 1) * 4]) == 4] if rows: cv2.imwrite(str(shot_dir / "contact_sheet.jpg"), np.vstack(rows), [cv2.IMWRITE_JPEG_QUALITY, 90]) def write_csv(path: Path, rows: list[dict]) -> None: with open(path, "w", newline="", encoding="utf-8") as f: w = csv.DictWriter(f, fieldnames=list(rows[0].keys())) w.writeheader() w.writerows(rows) def main() -> None: ap = argparse.ArgumentParser(description="Zero-shot detect + track cue ball (P0)") ap.add_argument("--root", type=Path, default=PILOT_ROOT) ap.add_argument("--weights", type=Path, default=DEFAULT_WEIGHTS) ap.add_argument("--conf", type=float, default=OP_CONF, help="conf van hanh cho so gate (mac dinh 0.2993 — dung cham)") ap.add_argument("--conf-floor", type=float, default=CONF_FLOOR) ap.add_argument("--imgsz", type=int, default=IMGSZ, help="doi khac 640 thi so KHONG phai so gate — ghi ro nhan") ap.add_argument("--device", default=0) ap.add_argument("--only", type=int, nargs="*", default=None) args = ap.parse_args() if not args.weights.exists(): sys.exit(f"[ERROR] Khong thay weights: {args.weights} — DUNG, hoi Danh.") from ultralytics import YOLO model = YOLO(str(args.weights)) 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] if not shots: sys.exit("[ERROR] Khong co shot_XX nao trong root.") print(f"weights={args.weights} op_conf={args.conf} imgsz={args.imgsz} " f"floor={args.conf_floor}") if args.conf != OP_CONF or args.imgsz != IMGSZ: print("[NOTE] Cau hinh KHAC van hanh — so sinh ra KHONG phai so gate G0.") summary = [] for d in shots: hpath = d / "homography.json" if not hpath.exists(): print(f"[WARN] {d.name}: chua co homography.json — bo qua.") continue hdata = json.loads(hpath.read_text(encoding="utf-8")) th = TableHomography(np.asarray(hdata["oriented_px"], dtype=np.float64), table_w=hdata["table_w_m"], table_l=hdata["table_l_m"]) meta = load_frames_meta(d) fps = json.loads((d / "meta.json").read_text(encoding="utf-8"))[ "fps_nominal"] dets = detect_shot(model, d, th, args.conf_floor, args.conf, args.imgsz, args.device) diffs = frame_diffs(d, meta, th) track_rows = track_cue(meta, dets, args.conf, diffs) write_csv(d / "detections.csv", dets) if dets else None write_csv(d / "track.csv", track_rows) render_overlay(d, meta, dets, track_rows, fps) n = len(meta) n_cue_frames = len({r["frame_file"] for r in track_rows if r["n_cands"] > 0}) n_cov = sum(r["covered"] for r in track_rows) gaps = [r["gap_s"] for r in track_rows if r["gap_s"]] s = {"shot": d.name, "frames": n, "cue_det_rate": round(n_cue_frames / n, 4), "coverage_frames": round(n_cov / n, 4), "max_gap_s": max(gaps) if gaps else 0.0} summary.append(s) print(f"{d.name}: frames={n} cue_det={s['cue_det_rate']:.1%} " f"cov={s['coverage_frames']:.1%} max_gap={s['max_gap_s']:.2f}s") if summary: write_csv(args.root / "track_summary.csv", summary) n80 = sum(1 for s in summary if s["coverage_frames"] >= 0.80) print(f"\n[G0 so bo] {n80}/{len(summary)} cu coverage >= 80% " f"(so chot o measure_p0 — coverage theo THOI LUONG, day la theo frame)") if __name__ == "__main__": main()