"""Real-frame gate check for hand-off 4, steps 1-3 -- design doc 7_FarZone_Detection_BallID.md section 6. Called by ``check_overlay_polish.bat``, which owns the environment. The unit tests pin the three pure functions; this pins what they do to the actual footage, which is the half a unit test cannot reach: step 1 ``_on_table_dets`` -- how many recorded boxes per clip project OFF the cloth (the boxes Danh saw drawn on the wall), where they land in table coordinates, and that none of them survives into the drawing; step 2 ``_static_balls`` -- how many frames of shot_02 carry a dashed box for ball 5, the ball whose 0.208 confidence keeps it out of ``frames_dets`` on every single frame while ``balls_init`` reads it at 0.688; step 3 ``CUE_EXCL_M`` -- the object-ball tracks WITH and WITHOUT the cue exclusion, from ONE detection pass (``track_others`` is captured on the way through and re-run by hand), plus the end point, heading, speed and pocket distance of ball 5's track that BRIEF 28/08 asks to be reported for Danh's eye. It also renders each clip's overlay TWICE -- once as the code now stands and once with both new drawing rules switched off -- and writes the same frame out of both as a PNG pair. Neither the numbers nor the pictures are tuned against: a bad number goes in the HANDOFF, it does not move a constant. Inputs are the clips the far-zone probe already cut from the archive VOD (``out_probe_farzone_20260828\\clips``), so the archive is not touched at all and no frame leaves this machine. """ from __future__ import annotations import importlib.util import json import math import os import sys import time from pathlib import Path import numpy as np APP_ROOT = Path(os.environ.get("POOLCOACH_APP_ROOT", r"D:\poolcoach\poolcoach-app")) CLIPS = Path(os.environ["POOLCOACH_POLISH_CLIPS"]) OUT = Path(os.environ["POOLCOACH_POLISH_OUT"]) BUNDLE = Path(os.environ.get( "POOLCOACH_POLISH_BUNDLE", APP_ROOT / "app" / "demo_analyzer" / "ad534bb871ba4193aa1a7bbebc1d484b")) DEVICE = os.environ.get("POOLCOACH_CV_DEVICE", "cpu") # The 4 corner pockets of run A2b (14/08) and every regeneration since. The # probe clips are cuts in TIME only, so the same pixels apply. CORNERS = [[527.5, 869.7], [1380.1, 864.0], [1232.3, 278.5], [673.0, 272.8]] SHOTS = (2, 5) # 02 carries ball 5 and the steal; 05 is the # control: a clip with no faint ball at all BALL5_XY_M = (1.1447, 2.4676) # the probe's MEASURED pixel centroid of BALL5_TOL_M = 0.06 # ball 5, mapped -- not a detection for _p in (APP_ROOT / "src", APP_ROOT): if str(_p) not in sys.path: sys.path.insert(0, str(_p)) def load_worker(): """``scripts/cv_worker.py`` as a module, loaded BY PATH the way ``app/localcv.py`` loads it -- importing rather than copying is what makes this measure the SHIPPED chain.""" name = "poolcoach_cv_worker" if name in sys.modules: return sys.modules[name] path = APP_ROOT / "scripts" / "cv_worker.py" spec = importlib.util.spec_from_file_location(name, path) if spec is None or spec.loader is None: raise ImportError(f"could not load {path}") mod = importlib.util.module_from_spec(spec) sys.modules[name] = mod spec.loader.exec_module(mod) return mod def analyse(cw, clip: Path, model): """One clip through the SAME chain an analyze job runs: the recording wrapper at the operating point, the 1280 first-frame pass, the height-compensated homography, then the first-frame ball numbers. Returns (result, rec, first_pass, captured) with ``captured`` holding the exact arguments ``track_others`` was called with.""" from poolcoach_cv import broadcast as bc from poolcoach_cv.overlay import RecordingModel rec = RecordingModel(model, min_conf=cw.OP_CONF) first_pass = bc.FirstFramePass(model, cw.ANALYZE_FIRST_IMGSZ) comp_meta: dict = {} captured: dict = {} orig_others = bc.track_others orig_homography = bc.TableHomography def capture(meta, others_by_frame, diffs, balls_init, low_by_frame=None, cue_rows=None): captured.update(meta=meta, others_by_frame=others_by_frame, diffs=diffs, balls_init=balls_init, low_by_frame=low_by_frame, cue_rows=cue_rows) return orig_others(meta, others_by_frame, diffs, balls_init, low_by_frame=low_by_frame, cue_rows=cue_rows) bc.track_others = capture bc.TableHomography = cw.comp_homography_cls(cw.clip_frame_wh(clip), bc.BALL_R_M, comp_meta) try: result = bc.analyze_clip(clip, CORNERS, rec, op_conf=cw.OP_CONF, imgsz=cw.ANALYZE_IMGSZ, first_pass=first_pass, device=DEVICE) finally: bc.track_others = orig_others bc.TableHomography = orig_homography result["height_comp"] = dict(comp_meta) if comp_meta else {"on": False} if first_pass.boxes: cw.ballid_balls_init(clip, CORNERS, result, first_pass.boxes, profile=None) return result, rec, first_pass, captured def plain_homography(): from poolcoach_cv.broadcast import TABLE_L_M, TABLE_W_M from poolcoach_cv.homography import TableHomography, orient_corners return TableHomography(orient_corners(np.asarray(CORNERS, dtype=np.float64)), table_w=TABLE_W_M, table_l=TABLE_L_M) def step1(rec, th, hc) -> dict: """Every recorded BALL box that projects off the cloth, and proof that the drawing keeps none of them.""" from poolcoach_cv import overlay as ov from poolcoach_cv.broadcast import BALL_CLASSES off: list[dict] = [] n_ball = 0 worst = (0, 0) # (count, frame index) for i, dets in enumerate(rec.frames): balls = [d for d in dets if d[0] in BALL_CLASSES] kept = ov._on_table_dets(th, dets, hc) n_ball += len(balls) n_off = len(balls) - len(kept) if n_off > worst[0]: worst = (n_off, i) if not n_off: continue kept_boxes = {tuple(b[2]) for b, _xy in kept} for cls, conf, box in balls: if tuple(box) in kept_boxes: continue xy = ov._comp_xy(np.asarray(th.px_to_table( ((box[0] + box[2]) / 2.0, (box[1] + box[3]) / 2.0)), dtype=np.float64).reshape(2), hc) off.append({"frame": i, "cls": cls, "conf": round(conf, 4), "x_m": round(float(xy[0]), 4), "y_m": round(float(xy[1]), 4)}) return {"n_frames": len(rec.frames), "n_ball_boxes": n_ball, "n_off_table": len(off), "worst_frame": worst[1], "worst_count": worst[0], "off_table": off[:40], "y_range": [min((d["y_m"] for d in off), default=None), max((d["y_m"] for d in off), default=None)]} def step2(result, rec, th, hc) -> dict: """How many frames each ``balls_init`` entry draws itself on, and which entry is ball 5.""" from poolcoach_cv import overlay as ov statics = ov._static_balls(result, th) drawn = [0] * len(statics) for dets in rec.frames: kept = ov._on_table_dets(th, dets, hc) seen = [xy for b, xy in kept if b[0] != "Cue"] for s in ov._undetected(statics, seen): drawn[statics.index(s)] += 1 rows = [] for k, ((x, y), box, num) in enumerate(statics): rows.append({"i": k, "x_m": round(x, 4), "y_m": round(y, 4), "number": num, "frames_drawn": drawn[k], "box_px": [round(v, 1) for v in box], "side_px": round(box[2] - box[0], 1), "is_ball5": math.hypot(x - BALL5_XY_M[0], y - BALL5_XY_M[1]) <= BALL5_TOL_M}) return {"n_frames": len(rec.frames), "entries": rows} def track_report(tracks: list[dict], balls_init: list[dict]) -> list[dict]: """Each object track as the handful of numbers the BRIEF asks for.""" from poolcoach_cv import broadcast as bc out = [] for tr in tracks: pts = tr["track"] last = pts[-1] chord = bc.end_chord(pts) pockets = bc.pocket_xy() near = min(pockets, key=lambda p: math.hypot(last["x_m"] - p[0], last["y_m"] - p[1])) idx = tr["init_idx"] row = {"init_idx": idx, "number": (balls_init[idx].get("number") if 0 <= idx < len(balls_init) else None), "n_points": len(pts), "t_last_s": last["t_s"], "last_xy": [last["x_m"], last["y_m"]], "pocketed_guess": tr.get("pocketed_guess"), "pocket_xy": tr.get("pocket_xy"), "nearest_pocket": [round(v, 4) for v in near], "d_pocket_m": round(math.hypot(last["x_m"] - near[0], last["y_m"] - near[1]), 4)} if chord is not None: speed, dx, dy = chord row["end_speed_mps"] = round(speed, 3) row["end_heading_deg"] = round(math.degrees(math.atan2(dy, dx)), 1) out.append(row) return out def step3(captured, result) -> dict: """The tracks with and without the exclusion, off ONE detection pass, plus how many candidates the exclusion actually struck off.""" from poolcoach_cv import broadcast as bc meta = captured["meta"] others = captured["others_by_frame"] low = captured["low_by_frame"] or {} rows = captured["cue_rows"] init = captured["balls_init"] without = bc.track_others(meta, others, captured["diffs"], init, low_by_frame=captured["low_by_frame"], cue_rows=None) cue_pts = bc.cue_positions(rows) n_hi = n_lo = 0 for name, t in meta[1:]: cue = bc.cue_at(cue_pts, float(t)) if cue is None: continue for band, bucket in (("hi", others.get(name) or []), ("lo", low.get(name) or [])): for d in bucket: if math.hypot(float(d["table_x_m"]) - cue[0], float(d["table_y_m"]) - cue[1]) \ < bc.CUE_EXCL_M: if band == "hi": n_hi += 1 else: n_lo += 1 return {"cue_excl_m": bc.CUE_EXCL_M, "n_excluded_high": n_hi, "n_excluded_low": n_lo, "with": track_report(result.get("ball_tracks") or [], init), "without": track_report(without, init)} def render_pair(clip: Path, result: dict, rec, stem: str, frames: list[int]) -> dict: """The overlay as the code now stands and as it was before steps 1-2, plus the same frames written out of both as PNGs.""" import cv2 from poolcoach_cv import overlay as ov out = {} for label in ("after", "before"): path = OUT / f"{stem}_{label}.mp4" keep_on_table = ov._on_table_dets keep_static = ov._static_balls if label == "before": # the two rules of steps 1 and 2 switched off, and NOTHING else: # every recorded ball box drawn, no box from the layout ov._on_table_dets = lambda th, dets, hc: [ (d, (0.0, 0.0)) for d in dets if d[0] in ov.BALL_CLASSES] ov._static_balls = lambda result, th: [] try: note = ov.render_shot_overlay(clip, CORNERS, result, rec.frames, path, label=f"{stem} {label}") finally: ov._on_table_dets = keep_on_table ov._static_balls = keep_static out[label] = {"file": path.name, "warning": note} cap = cv2.VideoCapture(str(path)) try: for f in frames: cap.set(cv2.CAP_PROP_POS_FRAMES, f) ok, img = cap.read() if ok: cv2.imwrite(str(OUT / f"{stem}_{label}_f{f:04d}.png"), img) finally: cap.release() return out def main() -> int: OUT.mkdir(parents=True, exist_ok=True) cw = load_worker() from poolcoach_cv import broadcast as bc print(f"[polish] app root : {APP_ROOT}") print(f"[polish] clips : {CLIPS}") print(f"[polish] baseline : {BUNDLE}") print(f"[polish] detect : analyze {cw.ANALYZE_IMGSZ}, first frame " f"{cw.ANALYZE_FIRST_IMGSZ}, op_conf {cw.OP_CONF}, floor " f"{bc.DET_CONF_FLOOR}") print(f"[polish] cue_excl : {bc.CUE_EXCL_M} m", flush=True) weights = cw.resolve_weights() print(f"[polish] loading {weights} ...", flush=True) from ultralytics import YOLO model = YOLO(str(weights)) for sz in (cw.ANALYZE_IMGSZ, cw.ANALYZE_FIRST_IMGSZ): model.predict(np.zeros((sz, sz, 3), dtype=np.uint8), conf=bc.DET_CONF_FLOOR, imgsz=sz, verbose=False, device=DEVICE) print("[polish] model ready", flush=True) th = plain_homography() report: dict = {"corners": CORNERS, "cue_excl_m": bc.CUE_EXCL_M, "shots": {}} for idx in SHOTS: stem = f"shot_{idx:02d}" clip = CLIPS / f"{stem}.mp4" if not clip.exists(): print(f"[polish] FAILED: {clip} is missing") return 2 t0 = time.time() print(f"\n[polish] === {stem} ===", flush=True) result, rec, _fp, captured = analyse(cw, clip, model) hc = result.get("height_comp") print(f"[polish] analysed in {time.time() - t0:.1f}s -- " f"{len(rec.frames)} frames, " f"{len(result['balls_init'])} balls_init, " f"{len(result.get('ball_tracks') or [])} ball tracks", flush=True) s1 = step1(rec, th, hc) print(f"[polish] STEP1 {s1['n_ball_boxes']} ball boxes recorded, " f"{s1['n_off_table']} project OFF the cloth " f"(y {s1['y_range'][0]} .. {s1['y_range'][1]}), worst frame " f"{s1['worst_frame']} with {s1['worst_count']}") for d in s1["off_table"][:6]: print(f"[polish] off-cloth f{d['frame']:4d} {d['cls']:<7} " f"conf {d['conf']:.3f} -> ({d['x_m']}, {d['y_m']})") s2 = step2(result, rec, th, hc) print(f"[polish] STEP2 static boxes over {s2['n_frames']} frames:") for r in s2["entries"]: print(f"[polish] init[{r['i']}] ({r['x_m']}, {r['y_m']}) " f"number {r['number']} drawn on {r['frames_drawn']} " f"frames side {r['side_px']} px" f"{' <-- ball 5' if r['is_ball5'] else ''}") s3 = step3(captured, result) print(f"[polish] STEP3 struck off {s3['n_excluded_high']} high and " f"{s3['n_excluded_low']} low candidates") for label in ("without", "with"): for r in s3[label]: print(f"[polish] {label:<7} init[{r['init_idx']}] " f"n={r['n_points']:3d} t_last={r['t_last_s']:.3f} " f"last={r['last_xy']} pocket=" f"{r['pocketed_guess']} {r.get('pocket_xy')} " f"d_pocket={r['d_pocket_m']} " f"speed={r.get('end_speed_mps')} " f"heading={r.get('end_heading_deg')}") frames = sorted({0, s1["worst_frame"], min(len(rec.frames) - 1, 33)}) vids = render_pair(clip, result, rec, stem, frames) print(f"[polish] rendered {vids} -- PNG frames {frames}", flush=True) report["shots"][stem] = {"step1": s1, "step2": s2, "step3": s3, "overlays": vids, "png_frames": frames, "balls_init": result["balls_init"], "ball_tracks": result.get("ball_tracks")} (OUT / "overlay_polish_check.json").write_text( json.dumps(report, indent=1), encoding="utf-8") print(f"\n[polish] report -> {OUT / 'overlay_polish_check.json'}") return 0 if __name__ == "__main__": raise SystemExit(main())