Spaces:
Sleeping
Sleeping
| """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()) | |