"""CV worker -- the process that runs the detector, ShotNet and the re-simulation for the Analyzer. It is SEPARATE from the API and runs in the CV environment (CUDA), so the app never imports torch or ultralytics. YOLO is loaded ONCE at boot and warmed with one predict; only when everything is loaded does the heartbeat turn on (never report alive while still starting up). Life cycle: BRPOP pc:jobs:segment | pc:jobs:analyze -> handler -> status key + files Both job types run for minutes, so neither uses request-reply: progress and results go to a status key that the API polls, and the durable output (the manifest, the per-shot JSON, the thumbnails and the overlays) is written to DISK as soon as it exists -- job results were genuinely lost once when they only lived in Redis under a 1 hour TTL. The operating confidence 0.2993 is the argmax of the smoothed any-ball F1 on the validation set (recorded in opconf.json: F1 0.9322, P 0.9642, R 0.9023 on the VAL pipeline; that number is not directly comparable with the predict pipeline running here -- two different rulers). Override it with the env var ``POOLCOACH_CV_CONF``. The heartbeat key is the worker's own (pc:worker:heartbeat:cv), and it is beaten from INSIDE a running job through the progress callback, so a long job does not make the upload route fast-fail while the worker is merely busy. The ``shotnet`` block carries ``confidence_raw`` plus ``det_density``, and ``confidence`` (the DISPLAYED rank) passes through the **detection-density valve** -- see the ``SHOTNET_DET_DENSITY_*`` table. The valve only changes the word and adds one ``warnings`` line: it never touches the model's numbers, hides a column, or blocks inference. An analyze job detects in TWO resolutions (design doc 7, section 3.1): every frame at ``ANALYZE_IMGSZ`` (960), plus one extra pass over the first frame at ``ANALYZE_FIRST_IMGSZ`` (1280) which is the sole source of ``balls_init`` and of the boxes its ball numbers are read from. Balls at the far rail are ~22 px wide and 640 simply does not see them. SEGMENTATION still runs at ``IMGSZ`` (640): the shot cut marks are a baseline that must not move. That same pass predicts down at ``DET_CONF_FLOOR`` and reads the answer in two bands (section 3.2): above the operating confidence nothing changed anywhere, below it a detection may only continue a track that is already alive. An analyze job COMPENSATES FOR THE BALL-CENTRE HEIGHT on every coordinate mapped through the homography (``comp_homography_cls`` -- the camera is decomposed from that clip's own homography, ``poolcoach_cv.camera``); the result carries a ``height_comp: {on, h_m, f_px, cx_m, cy_m}`` flag so the two coordinate conventions can be told apart. A failed decomposition or an out-of-bounds camera means no compensation plus an off flag and a warning. Nothing is applied retroactively to older tracks. The 1-9 ball numbers of ``balls_init`` come from the COLOUR of each ball, against the standard palette of ``poolcoach_cv.ballid``. A tournament using another ball set (Dynaspheres Titan: the 5 is purple, not orange) is served by a JSON profile whose path goes in the env var ``POOLCOACH_BALLID_PROFILE``; the result then carries a ``ballid_profile: {on, name, file}`` block so two eras of ball numbers stay distinguishable, the same way ``height_comp`` separates two coordinate conventions. The profile is read ONCE per job and given to BOTH layers that name a ball -- ``balls_init`` and the overlay video -- because a viewer reads the number on the video and the number on the metric table as one statement. With the variable unset nothing changes at all -- no block, no behaviour, the same palette (design doc 7_FarZone_Detection_BallID.md section 3.3). A segment job (a long video) cuts the shot list with ``poolcoach_cv.segmentation``, writes the manifest plus thumbnails and per-shot clips to FILES immediately, then enqueues N analyze jobs unchanged, adding only a ``result_path`` so the per-shot result lands in a file too. Running it by hand (Redis: docker compose -f docker-compose.dev.yml up -d redis): set REDIS_URL=redis://localhost:6379/0 ..\\poolcoach-cv-env\\Scripts\\python.exe scripts\\cv_worker.py """ from __future__ import annotations import math import os import sys import time from pathlib import Path # The repository is not installed as a package: src/ for poolcoach_cv and # ROOT for app.* (the scripts/ convention -- see CLAUDE.md). ROOT = Path(__file__).resolve().parents[1] for _p in (ROOT / "src", ROOT): if str(_p) not in sys.path: sys.path.insert(0, str(_p)) import numpy as np # noqa: E402 from app import analyzer_store as astore # noqa: E402 -- stdlib only from app import jobqueue # noqa: E402 -- pulls redis, not FastAPI or torch from poolcoach_cv.ballid import (identify_balls, # noqa: E402 -- pure numpy load_profile) from poolcoach_cv.broadcast import DET_CONF_FLOOR # noqa: E402 -- pure numpy from poolcoach_cv.camera import HeightCompensation # noqa: E402 -- pure numpy from poolcoach_cv.homography import (TableHomography, # noqa: E402 orient_corners) DEFAULT_WEIGHTS = str(ROOT / "models" / "cv_full_20260805" / "best.pt") # ---- ShotNet: the learned inverse model (cue V0, phi, a, b) that runs # ALONGSIDE the analytic estimate inside an analyze job. DEFAULT_SHOTNET = str(ROOT / "models" / "shotnet_20260812_c4b" / "best.pt") # the c4b checkpoint -- override with the env var # POOLCOACH_SHOTNET_WEIGHTS; a missing file still # leaves analysis running, just without the shotnet # block (with a clear warning). # # Its acceptance gate was still RED under the rule # "meet the absolute bar AND beat c3" (3 of 4 cells # below the absolute bar -- that result stands and is # not re-scored). The app moved to c4b as a DELIBERATE # DECISION, not because the gate went green: c4b beats # c3 in all four cells on the official ruler (the # >=0.3s slice, 669 paired shots: 2.24 deg / 8.56% / # 74.3% / 72.9% versus 3.03 / 9.73 / 69.7 / 65.3, with # the 95% CIs for dphi and V0 excluding 0), and it does # not regress on the full 5k set. SHOTNET_V0_SCALE_REAL = 2.54 / 1.9812 # converts the model's V0 to m/s on the REAL table # (~1.282). Pure geometry, no fitting: the net was # trained on the pooltool table (L=1.9812 m) while the # real input is normalised isotropically by the # tournament table (L=2.54 m, the same 2:1 ratio), so # the world the net sees is a 0.78x scale model of the # real table with the same timestamps. The V0 it reads # out is therefore m/s of that scaled world; # multiplying by L_real/L_synth returns real m/s. # phi, a and b are invariant under a similarity, so # they need no conversion. SHOTNET_CONF_HIGH = 0.8 # p_ident at or above this -> "high" spin confidence SHOTNET_CONF_MED = 0.5 # p_ident at or above this -> "medium"; below -> # "low" (spin only shows through a collision, so a # low p_ident is the net saying it saw no collision # to read spin from) # ---- Displayed-confidence valve based on TARGET-BALL DETECTION DENSITY. # The lesson from the first real clips: p_ident was 1.0 and the confidence # "high" on BOTH badly wrong shots (one was 84.9 degrees off in phi). The # confidence head CANNOT report unfamiliar input, because it learned "does # this shot have a collision", not "does this input look like training". # So a second valve measures something OUTSIDE the model: the target-ball # detection density of the clip itself (non-cue detections per frame -- the # very quantity that separated those two shots: 1.01 versus 3.09). # # The thresholds come from the measured TRAINING distribution (shard 0, 2000 # synthetic shots): mean 5.28 non-cue detections per frame, p10 = 2.00, # p90 = 8.21. Below p10 means the clip is sparser than 90% of the data the # net saw when its validation numbers were measured -- drop ONE rank; # sparser still (< 1.5) and the problem itself has lost information -- drop # TWO. Augmentation widened what the net TOLERATES but did not recreate lost # information, so the valve keeps the thresholds of the original # distribution. SHOTNET_DET_DENSITY_MED = 2.0 # below -> drop the DISPLAYED rank by 1 SHOTNET_DET_DENSITY_LOW = 1.5 # below -> drop it by 2 SHOTNET_CONF_RANKS = ("low", "medium", "high") def density_confidence(rank: str, density: float) -> str: """Downgrade the DISPLAYED confidence rank by detection density. It ONLY changes the word -- it touches no model number, hides no column and blocks no inference.""" if density < SHOTNET_DET_DENSITY_LOW: step = 2 elif density < SHOTNET_DET_DENSITY_MED: step = 1 else: return rank if rank not in SHOTNET_CONF_RANKS: return rank return SHOTNET_CONF_RANKS[max(SHOTNET_CONF_RANKS.index(rank) - step, 0)] # Rank labels -- they appear in a warning the front-end prints verbatim. _CONF_LABEL = {"low": "low", "medium": "medium", "high": "high"} def density_warning(blk: dict) -> str | None: """The warning line to show when the valve downgraded something -- None when it did not. The server composes the sentence (the same idea as ``spin_evidence[].text``) and the front-end prints it as is.""" raw, shown = blk.get("confidence_raw"), blk.get("confidence") if not raw or raw == shown: return None return (f"Target-ball detection density is only " f"{blk['det_density']:.2f}/frame -- sparser than the data the " f"net was trained on (p10 ~ {SHOTNET_DET_DENSITY_MED:g}/frame, " f"mean 5.3). The DISPLAYED confidence rank of the model block " f"was lowered from \"{_CONF_LABEL.get(raw, raw)}\" to " f"\"{_CONF_LABEL.get(shown, shown)}\"; the model numbers are " f"unchanged, read them with care.") OP_CONF = 0.2993 # the operating confidence (see the module docstring); # override with POOLCOACH_CV_CONF IMGSZ = 640 # the same imgsz as training and validation -- what # SEGMENTATION runs at, and the shot cut marks must # stay bit for bit what they were ANALYZE_IMGSZ = 960 # analyze: the main pass. 640 loses the far half of # the table (a ball at the far rail is ~7 px after # the 3x shrink); 960 finds the cue ball at the far # rail on 16 of 16 frames where 640 finds 5, for # 2.10x the detect cost (far-zone probe 28/08/2026) ANALYZE_FIRST_IMGSZ = 1280 # analyze: the extra pass over the FIRST frame only, # the sole source of balls_init. 1280 is the only # configuration that reads the far-pocket ball of # shot_02 above OP_CONF (0.688 against 0.208 at 960), # and paying for it on one frame per clip costs one # frame instead of 3.77x everything # ------------------------------------------------------------ analyze clip def not_shot_flag(result: dict) -> dict: """The "probably not a shot" valve: a DISPLAY flag resting entirely on three signals the shot already reports about itself in ``result``. No new signal, no change to segmentation, and no shot is deleted or filtered: 1. the analytic V0 is outside the trusted range ``bc.V0_SANE_MPS`` [1, 8] m/s; 2. no motion_start could be found; 3. coverage is below ``bc.COVERAGE_WARN`` (0.8). The three phantom shots of the reference rack (indices 7/8/9) report exactly these signals (v0 0.27 / None / 0.99 m/s, coverage 45.5/38.9/60%) while all 11 real shots are clean on all three (the lowest coverage is 92.6%) -- the acceptance gate was 3 of 3 phantoms caught, 0 of 11 real shots flagged unfairly. A V0 of None WITH a motion_start (the rare "not enough frames after the start to measure V0") does NOT fire signal 1: exactly three signals were agreed, and "outside the range" only applies when there is a number. Such a case usually trips signal 3 anyway if it really is junk. The reasons are sentences composed here, in the same style as ``spin_evidence[].text``. """ from poolcoach_cv import broadcast as bc m = result.get("metrics") or {} reasons: list[str] = [] v0 = m.get("v0_mps") if v0 is not None and not (bc.V0_SANE_MPS[0] <= v0 <= bc.V0_SANE_MPS[1]): reasons.append(f"the analytic V0 {v0:.2f} m/s is outside the " f"trusted range [{bc.V0_SANE_MPS[0]:g}, " f"{bc.V0_SANE_MPS[1]:g}] m/s") if m.get("motion_start_s") is None: reasons.append("the moment the cue ball starts moving " "(motion_start) was not found") cov = m.get("coverage") if cov is not None and cov < bc.COVERAGE_WARN: reasons.append(f"the track covers only {cov:.0%} of the duration " f"(< {bc.COVERAGE_WARN:.0%})") return {"flagged": bool(reasons), "reasons": reasons} def load_shotnet(device: str): """Load ShotNet ONCE at boot (the same rule that keeps YOLO warm). Returns {model, cfg, tag, device}, or None when the file is missing or fails to load -- analysis still runs, the shotnet block is simply absent with a per-job warning.""" import torch from poolcoach_cv.shotnet import ShotNet, ShotNetConfig path = Path(os.environ.get("POOLCOACH_SHOTNET_WEIGHTS", DEFAULT_SHOTNET)) if not path.exists(): print(f"[cv_worker] ShotNet not found at {path} -- analysis will run " f"WITHOUT the shotnet block (set POOLCOACH_SHOTNET_WEIGHTS if " f"the model lives elsewhere)", flush=True) return None dev = device if device else ("cuda" if torch.cuda.is_available() else "cpu") if dev.strip().isdigit(): # YOLO convention: device="0" is cuda:0 dev = f"cuda:{dev.strip()}" ck = torch.load(path, map_location=dev) cfg = ShotNetConfig(**ck["model_config"]) net = ShotNet(cfg).to(dev) net.load_state_dict(ck["state_dict"]) net.eval() return {"model": net, "cfg": cfg, "tag": path.parent.name, "device": dev} def load_resim(): """Load and JIT pooltool for the re-simulation ONCE at boot, in the same style as load_shotnet: missing or broken means None, analysis still runs and the resim block is absent with a per-job warning. A cold JIT takes 40-90s, so it has to happen BEFORE the heartbeat turns on.""" try: from poolcoach_cv import resim as rs except Exception as e: # noqa: BLE001 print(f"[cv_worker] could not import poolcoach_cv.resim " f"({type(e).__name__}: {e}) -- re-simulation OFF", flush=True) return None state = rs.warmup() if state is None: return None return {"module": rs, "state": state} def shotnet_infer(sn: dict, rows: list[dict], others: list[dict]) -> dict: """The app's track and detections -> the ``shotnet`` block of the result JSON. The features are built with the EXACT loader code (shot_from_track + featurize_shot -- one source of truth), with the real table passed into featurize through broadcast's (w, l).""" import torch from poolcoach_cv import broadcast as bc from poolcoach_cv.shotnet import featurize_shot, shot_from_track, \ spin_classes t0 = time.perf_counter() shot = shot_from_track(rows, others) feats, t = featurize_shot(shot, w=bc.TABLE_W_M, l=bc.TABLE_L_M, deltas=sn["cfg"].use_deltas) dev = sn["device"] x = torch.from_numpy(feats)[None].to(dev) tt = torch.from_numpy(t)[None].to(dev) mask = torch.ones(1, len(t), dtype=torch.bool, device=dev) p = sn["model"].predict(x, tt, mask) a, b = float(p["a"][0]), float(p["b"][0]) vert, _side = spin_classes(np.array([a]), np.array([b])) p_ident = float(p["p_ident"][0]) conf_raw = ("high" if p_ident >= SHOTNET_CONF_HIGH else "medium" if p_ident >= SHOTNET_CONF_MED else "low") # this clip's own target-ball detection density (non-cue detections per # frame) -- a quantity outside the model, measured on the very input fed # into the net density = len(others) / max(len(rows), 1) conf = density_confidence(conf_raw, density) return { "model": sn["tag"], # CUE V0 (the label/net ruler), converted geometrically to m/s on the # real table; a DIFFERENT ruler from metrics.v0_mps (the BALL speed # measured off the track) "v0_cue_mps": round(float(p["v0"][0]) * SHOTNET_V0_SCALE_REAL, 3), "phi_deg": round(float(p["phi_deg"][0]), 1), "a": round(a, 3), "b": round(b, 3), # the same mapping as the training harness: side is the plain sign of # a (no dead zone), vert is quantised into 3 classes at B_STUN_MAX "spin_vert": str(vert[0]), "spin_side": "side-L" if a > 0 else "side-R", "identifiable_prob": round(p_ident, 3), # `confidence` is the DISPLAYED rank (already through the density # valve); `confidence_raw` is the raw rank from the identifiable head # -- both are kept so no number is lost "confidence": conf, "confidence_raw": conf_raw, "det_density": round(density, 3), "inference_ms": round((time.perf_counter() - t0) * 1000, 1), "n_target_slots": int(shot["n_balls"]) - 1, } def comp_homography_cls(image_wh: tuple[float, float], ball_r: float, sink: dict): """A TableHomography subclass that compensates for the ball-centre height. ``px_to_table`` applies k = R/(h-R) along the camera-to-point direction, with the camera decomposed from THAT CLIP's own homography (never a shared constant). Inside ``analyze_clip`` EVERY mapped coordinate (the cue track, the non-cue detections, the first-frame balls_init) goes through ``px_to_table``, so the compensation covers every mapped position, not just points near a rail. A failed decomposition or an implausible camera makes HeightCompensation fall back to the identity -- better no compensation than one built on a junk camera. ``sink`` receives the ``height_comp`` metadata when the homography is constructed, and the worker writes it into the job JSON. It is patched over ``bc.TableHomography`` around ONE analyze_clip call (a script-level adapter that does not modify broadcast.py); the worker serves jobs sequentially, so there is no race. """ class _CompHomography(TableHomography): def __init__(self, corners_px, **kw) -> None: super().__init__(corners_px, **kw) self._comp = HeightCompensation.from_homography( self.H, image_wh, ball_r) sink.clear() sink.update(self._comp.meta()) def px_to_table(self, pts_px): return self._comp.apply(super().px_to_table(pts_px)) return _CompHomography def clip_frame_wh(clip_path) -> tuple[float, float]: """The clip's (width, height) in pixels -- the nominal principal point for the decomposition sits at its centre. Unreadable gives (0, 0), which turns the compensator off; a genuinely broken clip raises from analyze_clip a moment later with its own message.""" import cv2 cap = cv2.VideoCapture(str(clip_path)) try: return (cap.get(cv2.CAP_PROP_FRAME_WIDTH), cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) finally: cap.release() def match_ballid_to_init(balls_init: list[dict], numbered: list[tuple], tol_m: float) -> int: """Attach (number, number_conf) to the NEAREST ``balls_init`` entry within ``tol_m`` -- pure arithmetic, split out so it can be unit tested in the app environment without cv2. ``numbered``: [(tx, ty, number, number_conf), ...], the table coordinates of the first-frame boxes, already passed through the same homography as balls_init so they match almost point for point. Higher confidence is attached first, and an entry that already has a number is never overwritten (when two nearby boxes compete for one entry, the higher-confidence one wins). Returns how many entries were attached.""" import math n = 0 for tx, ty, num, nc in sorted(numbered, key=lambda q: -q[3]): if num is None: continue best, bd = None, float("inf") for entry in balls_init: if entry.get("type") != "ball" or "number" in entry: continue d = math.hypot(entry["x_m"] - tx, entry["y_m"] - ty) if d < bd: best, bd = entry, d if best is not None and bd <= tol_m: best["number"] = int(num) best["number_conf"] = round(float(nc), 3) n += 1 return n BALLID_PROFILE_ENV = "POOLCOACH_BALLID_PROFILE" def ballid_profile() -> tuple[dict | None, dict, list[str]]: """The ball-set colour profile of THIS job -> ``(profile, meta, warnings)``. The path comes from the env var ``POOLCOACH_BALLID_PROFILE`` and is read PER JOB, not once at boot: a profile is a property of the footage being analysed, not of the worker, and re-reading a few hundred bytes costs nothing next to a clip. ``meta`` is what the job JSON carries, in the same spirit as ``height_comp``: which ball set these numbers belong to, so two eras of ball numbers never get compared by accident. It is EMPTY when the env var is unset -- and an empty meta is written nowhere, which is what keeps an unconfigured worker's JSON byte for byte what it was before profiles existed. When the var IS set but the file is unusable, the meta says so (``on: False`` plus the reason) rather than staying silent: a profile that was asked for and did not load has to be visible. """ raw = (os.environ.get(BALLID_PROFILE_ENV) or "").strip() if not raw: return None, {}, [] profile, notes = load_profile(raw) if profile is None: return None, {"on": False, "file": raw, "reason": notes[-1] if notes else "unusable"}, notes return profile, {"on": True, "name": profile["name"], "file": raw}, notes def ballid_balls_init(clip_path, corners_px, result: dict, first_dets: list, profile: dict | None = None ) -> str | None: """Attach the 1-9 ball numbers to ``balls_init``, from ``identify_balls`` (the two-stage colour BallID) on the FIRST FRAME with exactly the boxes the numeric layer used -- which since the hybrid detect means the boxes of the first-frame 1280 pass (``FirstFramePass.boxes``), the very ones ``balls_init`` was built from, not the main pass the overlay draws -- with the box centres mapped through the job's SAME height-compensated homography (``comp_homography_cls`` is rebuilt -- the decomposition is deterministic, so the same input gives the same map). An unrecognised colour, or an older result, simply leaves the entry without a ``number`` field and the front-end draws a grey ball rather than inventing one. ``profile``: the ball-set table from `ballid_profile`, or None for the standard palette. Returns None on success, or a warning message when something went wrong -- the NUMERIC result is never touched either way. """ import cv2 from poolcoach_cv import broadcast as bc balls_init = result.get("balls_init") or [] balls = [(cls, conf, box) for cls, conf, box in first_dets if cls in bc.BALL_CLASSES] others = [b for b in balls if b[0] != "Cue"] if not balls_init or not others: return None cap = cv2.VideoCapture(str(clip_path)) try: ok, frame = cap.read() finally: cap.release() if not ok: return ("The first frame could not be reread to attach ball " "numbers -- balls on the metric table show grey.") cue = next((b for b in balls if b[0] == "Cue"), None) assigns, _wb = identify_balls(frame, [b[2] for b in others], cue[2] if cue else None, profile=profile) th_cls = comp_homography_cls(clip_frame_wh(clip_path), bc.BALL_R_M, {}) th = th_cls(orient_corners(np.asarray(corners_px, dtype=np.float64)), table_w=bc.TABLE_W_M, table_l=bc.TABLE_L_M) numbered = [] for (_cls, _conf, (x1, y1, x2, y2)), (num, nc) in zip(others, assigns): tx, ty = th.px_to_table(((x1 + x2) / 2.0, (y1 + y2) / 2.0)) numbered.append((float(tx), float(ty), num, float(nc or 0.0))) match_ballid_to_init(balls_init, numbered, bc.BALL_DEDUP_M) return None def analyze_once(payload: dict, model, conf: float, predict_kw: dict, status, shotnet: dict | None = None, resim: dict | None = None) -> dict: """ONE analyze job: read the clip from its path -> poolcoach_cv.broadcast -> then ShotNet over the same track and detections, so the result gains a ``shotnet`` block (a pure addition; ``shotnet=None`` leaves the analytic numbers plus a warning). Raises ValueError for anything that is the input's fault (a missing, broken or over-long clip, degenerate corners). `status(dict)` reports progress, and the progress callback also BEATS the heartbeat, because the job outlasts the 15s TTL many times over -- without the beat the upload route would fast-fail while the worker is simply busy. """ clip_path = payload.get("clip_path") if not clip_path or not Path(clip_path).exists(): raise ValueError("The clip is no longer in the temporary " "directory -- upload again (temporary files live " "for a session, they are not kept).") from poolcoach_cv import broadcast as bc def progress(stage: str, frac: float) -> None: jobqueue.beat(key=jobqueue.CV_HEARTBEAT_KEY) status({"status": "running", "progress": round(min(frac, 0.99), 3), "stage": stage}) # ShotNet needs exactly analyze_track's input (per-frame rows plus the # non-cue detections), which analyze_clip does not return, so it is # captured with a module-level wrapper around ONE call (a worker-side # adapter that does not modify broadcast.py; the worker serves jobs # sequentially, so there is no race). TableHomography is patched the same # way: px_to_table compensates the ball-centre height for EVERY mapped # position (comp_homography_cls), which means ShotNet also eats a # compensated track. captured: dict = {} comp_meta: dict = {} orig_analyze_track = bc.analyze_track orig_homography = bc.TableHomography def _capture(rows, others=None): captured["rows"], captured["others"] = rows, others return orig_analyze_track(rows, others=others) # Per-shot overlay: wrap the model to RECORD exactly the boxes the # numeric layer used, so rendering does not detect again (the same # adapter-around-one-call pattern). Only the multi-shot flow sends an # overlay_path. overlay_path = payload.get("overlay_path") rec = None if overlay_path: from poolcoach_cv.overlay import RecordingModel # min_conf = the operating point: the overlay draws and numbers the # HIGH band only, the recovery lane is for the trackers alone rec = RecordingModel(model, min_conf=conf) # The hybrid detect of design 7 section 3.1: every frame at # ANALYZE_IMGSZ, the first frame ALSO at ANALYZE_FIRST_IMGSZ, and it is # that reading balls_init comes from. The pass gets the RAW model on # purpose -- routing it through `rec` would append a frame the clip does # not have and shift every overlay drawing by one. first_pass = bc.FirstFramePass(model, ANALYZE_FIRST_IMGSZ) t0 = time.perf_counter() bc.analyze_track = _capture bc.TableHomography = comp_homography_cls(clip_frame_wh(clip_path), bc.BALL_R_M, comp_meta) try: result = bc.analyze_clip(clip_path, payload.get("corners"), rec if rec is not None else model, op_conf=conf, imgsz=ANALYZE_IMGSZ, first_pass=first_pass, progress=progress, device=predict_kw.get("device")) finally: bc.analyze_track = orig_analyze_track bc.TableHomography = orig_homography result["metrics"]["elapsed_s"] = round(time.perf_counter() - t0, 2) # the coordinate-convention flag: was this track height-compensated or # not? Reading the JSON later tells you which era it belongs to; older # results have no flag, which means uncompensated. result["height_comp"] = dict(comp_meta) if comp_meta else {"on": False} if not result["height_comp"]["on"]: result["warnings"].append( "Ball-centre height was not compensated (" + result["height_comp"].get("reason", "reason unknown") + ") -- coordinates follow the old convention, a ball near a " "rail may sit 1-3cm too far from it.") # the shotnet block -- TWO methods shown side by side (the analytic # numbers are untouched); a broken or missing net must NOT kill the # analytic result, it is declared through warnings if shotnet is None: result["warnings"].append( "ShotNet did not load (a missing checkpoint?) -- only the " "analytic numbers are available.") elif captured.get("rows"): progress("shotnet", 0.99) try: result["shotnet"] = shotnet_infer(shotnet, captured["rows"], captured.get("others") or []) # the density valve adds ONE warning line and nothing else: no # model number changes and nothing is blocked note = density_warning(result["shotnet"]) if note: result["warnings"].append(note) except Exception as e: # noqa: BLE001 -- a broken net is not a lost analytic result["warnings"].append( f"ShotNet failed during inference ({type(e).__name__}: " f"{e}) -- only the analytic numbers are available.") # Re-simulation RMSE: feed the inferred parameters into pooltool, replay # the shot, and produce the ``resim`` block (two sets, shotnet and # analytic, each with its series and rmse). A step ADDED AFTER analysis # that touches no existing number; a broken resim or a missing pooltool # is declared through warnings rather than killing the result. if resim is None: result["warnings"].append( "pooltool is not available in the CV venv -- no re-simulation " "RMSE.") else: progress("resim", 0.99) try: result["resim"] = resim["module"].resim_result(result, resim["state"]) except Exception as e: # noqa: BLE001 -- a broken resim is not a lost analysis result["warnings"].append( f"Re-simulation failed ({type(e).__name__}: {e}) -- no " f"RMSE.") # the "probably not a shot" valve: a DISPLAY flag from the self-reported # signals, always present in a new result (flagged False means checked # and clean); an older result has no key, i.e. it predates the valve result["suspect_not_shot"] = not_shot_flag(result) # the ball-set colour profile: read ONCE per job and handed to BOTH # layers that name a ball -- balls_init below and the overlay further # down. Reading it twice would let a file swapped mid-job number the # metric table one way and the video another, which is the very # divergence this block exists to close. profile: dict | None = None if rec is not None: profile, prof_meta, prof_notes = ballid_profile() if prof_meta: # only when a profile was ASKED FOR: an unconfigured worker # writes no key at all, so the JSON schema is unchanged result["ballid_profile"] = prof_meta result["warnings"].extend(prof_notes) # first-frame ball numbers: attach `number` to balls_init from the boxes # balls_init itself came from (the first-frame 1280 pass) -- only on the # flow that has a RecordingModel (the per-shot jobs of a multi-shot # video). A failure is only a warning and touches no measurement. if rec is not None and first_pass.boxes: try: note = ballid_balls_init(clip_path, payload.get("corners"), result, first_pass.boxes, profile=profile) if note: result["warnings"].append(note) except Exception as e: # noqa: BLE001 -- bad colour is not a lost result result["warnings"].append( f"Ball numbers could not be attached from the first frame " f"({type(e).__name__}: {e}) -- balls on the metric table " f"show grey.") # the per-shot overlay.mp4: rendered RIGHT HERE in the job, while the # clip is STILL on disk (handle_analyze deletes it once analyze_once # returns, and an overlay cannot be produced retroactively -- the lesson # of the first full run). Flagged shots are rendered too (with a red # banner on the frame, which helps a human check them); a failed render # is only a warning and never kills the numbers. if rec is not None: from poolcoach_cv.overlay import render_shot_overlay progress("overlay", 0.99) try: note = render_shot_overlay( clip_path, payload.get("corners"), result, rec.frames, overlay_path, label=Path(str(overlay_path)).stem, beat=lambda: jobqueue.beat(key=jobqueue.CV_HEARTBEAT_KEY), profile=profile) result["overlay"] = {"file": Path(str(overlay_path)).name} if note: result["warnings"].append(note) except Exception as e: # noqa: BLE001 -- a broken overlay is not a lost measurement result["warnings"].append( f"The overlay could not be rendered " f"({type(e).__name__}: {e}) -- the measurements are " f"unaffected.") return result def handle_analyze(payload: dict, model, conf: float, predict_kw: dict, shotnet: dict | None = None, resim: dict | None = None) -> dict: """The transport reply protocol, though the real reader is the status key the API polls -- the reply itself lands on an orphaned result list and expires. Status writes are best-effort: one dropped Redis beat mid-job must not kill the worker (the job runs to the end, and a GET sees the stale status until the next beat overwrites it). The temporary clip is DELETED on success and on an input error (copyright), but KEPT on an internal error so the cause can still be traced. ``result_path`` is optional in the payload (the per-shot jobs of the multi-shot flow): the result or the error is ALSO written to a FILE right here, because job results were genuinely lost once when they only lived in Redis under a 1 hour TTL. """ analyze_id = str(payload.get("analyze_id", "")) result_path = payload.get("result_path") def status(data: dict) -> None: try: jobqueue.set_analyze_status(analyze_id, data) except jobqueue.QueueDown: pass def persist(data: dict) -> None: if not result_path: return try: astore.write_json(Path(result_path), data) except OSError as e: print(f"[cv_worker] WARNING: could not write {result_path}: " f"{type(e).__name__}: {e}", flush=True) def drop_clip() -> None: p = payload.get("clip_path") if p: Path(p).unlink(missing_ok=True) status({"status": "running", "progress": 0.0, "stage": "decode"}) try: result = analyze_once(payload, model, conf, predict_kw, status, shotnet=shotnet, resim=resim) except ValueError as e: drop_clip() persist({"error": "validation", "message": str(e)}) status({"status": "error", "message": str(e)}) return {"ok": False, "error": "validation", "message": str(e)} except Exception as e: # noqa: BLE001 -- a broken job is not a dead worker persist({"error": "internal", "message": f"{type(e).__name__}: {e}"}) status({"status": "error", "message": f"{type(e).__name__}: {e}"}) return {"ok": False, "error": "internal", "message": f"{type(e).__name__}: {e}"} persist(result) # the file BEFORE Redis -- the disk is the truth drop_clip() status({"status": "done", "progress": 1.0, "result": result}) return {"ok": True, "result": {"analyze_id": analyze_id}} # --------------------------------------------- segment a multi-shot video def cut_clip(video_path: Path, t0: float, dur: float, out: Path) -> str | None: """Cut ONE per-shot clip with the system ffmpeg (re-encoding to h264, so the seek is frame accurate -- a stream copy would slip back to the previous keyframe by as much as a second). Returns None on success or an error message on failure. ``-an`` drops the audio.""" import shutil import subprocess exe = shutil.which("ffmpeg") if exe is None: return ("the machine has no ffmpeg -- per-shot clips cannot be " "cut") out.parent.mkdir(parents=True, exist_ok=True) try: r = subprocess.run( [exe, "-y", "-v", "error", "-ss", f"{max(t0, 0.0):.3f}", "-i", str(video_path), "-t", f"{dur:.3f}", "-an", "-c:v", "libx264", "-preset", "veryfast", "-crf", "18", str(out)], capture_output=True, text=True, timeout=300) except subprocess.TimeoutExpired: return "ffmpeg took over 300s to cut the clip -- is it stuck?" if r.returncode != 0 or not out.exists(): err = (r.stderr or "unknown ffmpeg error").strip().splitlines()[-1] return f"ffmpeg could not cut the clip ({err[:200]})" return None def segment_once(payload: dict, model, conf: float, predict_kw: dict, status) -> dict: """ONE segment job: a long video -> the shot list -> the manifest written to a FILE immediately -> thumbnails and per-shot clips -> N analyze jobs enqueued unchanged. Raises ValueError for anything that is the input's fault (a missing, broken or over-long video, degenerate corners). It is the longest job in the system (a 5-20 minute video), so the progress callback beats the heartbeat exactly as the analyze job does. """ video_path = payload.get("video_path") if not video_path or not Path(video_path).exists(): raise ValueError("The video is no longer on disk -- upload " "again.") video_path = Path(video_path) out_dir = Path(payload["out_dir"]) video_id = str(payload.get("video_id", "")) from poolcoach_cv.segmentation import segment_video def progress(stage: str, frac: float) -> None: jobqueue.beat(key=jobqueue.CV_HEARTBEAT_KEY) status({"status": "running", "progress": round(min(frac, 0.99), 3), "stage": stage}) t0 = time.perf_counter() seg = segment_video(video_path, payload.get("corners"), model, op_conf=conf, device=predict_kw.get("device"), max_dur_s=astore.VIDEO_MAX_DUR_S, progress=progress) # the per-shot manifest -- initial status queued (a healthy shot) or # error (a shot segmentation already convicted: scene lost, boundary # never closed, ...) shots = [] for s in seg["shots"]: idx = int(s["idx"]) dur = float(s["t_end_s"]) - float(s["t_start_s"]) entry = dict(s) entry["analyze_id"] = astore.shot_analyze_id(video_id, idx) entry["thumb"] = astore.thumb_path(out_dir, idx).name if s["status"] == "ok": if dur > bc_max_dur(): entry["status"] = "error" entry["reason"] = (f"the shot is {dur:.1f}s long, over the " f"{bc_max_dur():g}s analysis clip " f"ceiling") else: entry["status"] = "queued" entry["clip"] = f"clips/{astore.clip_path(out_dir, idx).name}" shots.append(entry) # thumbnail of each shot's first frame (seeking by absolute PTS, the same # scale as t_start) import cv2 cap = cv2.VideoCapture(str(video_path)) try: for entry in shots: jobqueue.beat(key=jobqueue.CV_HEARTBEAT_KEY) cap.set(cv2.CAP_PROP_POS_MSEC, float(entry["t_start_s"]) * 1000.0) ok, frame = cap.read() if not ok: continue h, w = frame.shape[:2] scale = 320.0 / max(w, 1) thumb = cv2.resize(frame, (320, max(1, round(h * scale)))) cv2.imwrite(str(astore.thumb_path(out_dir, int(entry["idx"]))), thumb, [cv2.IMWRITE_JPEG_QUALITY, 85]) finally: cap.release() # cut the clip of every healthy shot -- a failure turns THAT shot into an # error with a reason while the others carry on (one bad shot does not # kill the list) n_cut = 0 for entry in shots: if entry["status"] != "queued": continue jobqueue.beat(key=jobqueue.CV_HEARTBEAT_KEY) status({"status": "running", "stage": "cut", "progress": round(0.92 + 0.07 * n_cut / max(len(shots), 1), 3)}) err = cut_clip(video_path, float(entry["t_start_s"]), float(entry["t_end_s"]) - float(entry["t_start_s"]), astore.clip_path(out_dir, int(entry["idx"]))) if err is not None: entry["status"] = "error" entry["reason"] = err entry.pop("clip", None) n_cut += 1 manifest = { "status": "done", "video_id": video_id, "filename": payload.get("filename"), "fps_nominal": seg.get("fps_nominal"), "t_first_s": seg.get("t_first_s"), "t_last_s": seg.get("t_last_s"), "duration_s": seg.get("duration_s"), "n_frames": seg.get("n_frames"), "n_dup_frames": seg.get("n_dup_frames"), "elapsed_scan_s": round(time.perf_counter() - t0, 2), "warnings": seg.get("warnings", []), "shots": shots, } # the manifest reaches DISK before any shot is analyzed -- if the worker # dies mid-way, the shot list and its reasons are still there astore.write_json(astore.manifest_path(out_dir), manifest) # enqueue the per-shot analyze jobs unchanged, adding only result_path; # the "queued" status is written BEFORE the enqueue (that race was caught # by a test once) for entry in shots: if entry["status"] != "queued": continue idx = int(entry["idx"]) aid = entry["analyze_id"] try: jobqueue.set_analyze_status(aid, {"status": "queued", "progress": 0.0}) jobqueue.enqueue( {"analyze_id": aid, "clip_path": str(astore.clip_path(out_dir, idx)), "corners": payload.get("corners"), "result_path": str(astore.shot_json_path(out_dir, idx)), # the overlay is rendered inside the analyze job, before the # per-shot clip is deleted "overlay_path": str(astore.overlay_path(out_dir, idx))}, jobs_key=jobqueue.ANALYZE_JOBS_KEY, job_id=aid) except jobqueue.QueueDown: # Redis dropped this beat: the shot sits "queued" in the manifest # with no job, so a GET reports unknown; the recovery path is to # upload again. print(f"[cv_worker] WARNING: could not enqueue {aid} " f"(Redis down)", flush=True) # the source video has done its job (manifest, thumbnails and clips all # exist) -- delete it so copyrighted footage does not linger; the # per-shot results live on as JSON files video_path.unlink(missing_ok=True) n_queued = sum(1 for s in shots if s["status"] == "queued") return {"n_shots": len(shots), "n_queued": n_queued, "elapsed_scan_s": manifest["elapsed_scan_s"]} def bc_max_dur() -> float: """The analyze clip duration ceiling (30s), read from broadcast so this module carries no copy that could drift. A function rather than a constant because broadcast can be imported lazily in any environment.""" from poolcoach_cv import broadcast as bc return bc.ANALYZE_MAX_DUR_S def handle_segment(payload: dict, model, conf: float, predict_kw: dict) -> dict: """Same reply protocol as handle_analyze -- the real readers are the per-video status key and the manifest file. An error is ALSO written as an error manifest to DISK (file first), so that once the status key expires after an hour the user can still see why this video produced no shot list.""" video_id = str(payload.get("video_id", "")) out_dir = Path(payload.get("out_dir", "")) def status(data: dict) -> None: try: jobqueue.set_video_status(video_id, data) except jobqueue.QueueDown: pass def persist_error(message: str) -> None: if not str(out_dir): return try: astore.write_json(astore.manifest_path(out_dir), {"status": "error", "video_id": video_id, "message": message}) except OSError as e: print(f"[cv_worker] WARNING: could not write the error " f"manifest: {type(e).__name__}: {e}", flush=True) status({"status": "running", "progress": 0.0, "stage": "decode"}) try: result = segment_once(payload, model, conf, predict_kw, status) except ValueError as e: p = payload.get("video_path") if p: Path(p).unlink(missing_ok=True) # broken video -- no leftovers persist_error(str(e)) status({"status": "error", "message": str(e)}) return {"ok": False, "error": "validation", "message": str(e)} except Exception as e: # noqa: BLE001 -- a broken job is not a dead worker # KEEP the video on an internal error, so the cause stays traceable persist_error(f"{type(e).__name__}: {e}") status({"status": "error", "message": f"{type(e).__name__}: {e}"}) return {"ok": False, "error": "internal", "message": f"{type(e).__name__}: {e}"} status({"status": "done", "progress": 1.0, **result}) return {"ok": True, "result": {"video_id": video_id, **result}} def resolve_weights() -> Path: """Path to the detector weights -- the env var ``POOLCOACH_CV_WEIGHTS`` first, then the copy in the repository at ``models/cv_full_20260805/best.pt`` (in LFS, the very file the Dockerfile copies by name).""" env = os.environ.get("POOLCOACH_CV_WEIGHTS") return Path(env) if env else Path(DEFAULT_WEIGHTS) def load_all(conf: float, device: str, predict_kw: dict): """Load YOLO, ShotNet and pooltool (for the re-simulation) ONCE -- the worker's whole "starting up" phase, split out so the embedded mode (app/localcv.py) SHARES this code and there is no second copy of the load and warmup order. Raises FileNotFoundError when the detector weights are missing; a missing ShotNet or pooltool returns None instead (analysis still runs, with a per-job warning).""" weights = resolve_weights() if not weights.exists(): raise FileNotFoundError( f"weights not found at {weights} -- set POOLCOACH_CV_WEIGHTS if " f"the model lives elsewhere") t0 = time.time() from ultralytics import YOLO model = YOLO(str(weights)) # warm up with one predict on a black image: this builds the CUDA context # and fuses the model now, so the first real job does not pay for it (the # same reason the physics layer is warmed at boot). Once per imgsz the # worker actually runs -- segmentation, the analyze main pass and the # analyze first-frame pass each size their own buffers. for sz in (IMGSZ, ANALYZE_IMGSZ, ANALYZE_FIRST_IMGSZ): model.predict(np.zeros((sz, sz, 3), dtype=np.uint8), conf=conf, imgsz=sz, verbose=False, **predict_kw) # ShotNet: loaded ONCE at boot like YOLO -- 10.7 MB, about a second; a # missing file gives None and analysis still runs (with a per-job warning) shotnet = load_shotnet(device) # pooltool for the re-simulation: a cold JIT takes 40-90s, so it happens # BEFORE the heartbeat; missing pooltool gives None and analysis survives resim = load_resim() print(f"[cv_worker] model {weights.name} ready after " f"{time.time() - t0:.1f}s (conf={conf:g}" f", floor={DET_CONF_FLOOR:g}" f", imgsz={IMGSZ}/{ANALYZE_IMGSZ}/{ANALYZE_FIRST_IMGSZ}" f"{', device=' + device if device else ''}" f", shotnet={'OK ' + shotnet['tag'] if shotnet else 'NONE'}" f", resim={'OK' if resim else 'NONE'})" f" -- heartbeat on, waiting for jobs on " f"'{jobqueue.SEGMENT_JOBS_KEY}' + '{jobqueue.ANALYZE_JOBS_KEY}'", flush=True) return model, shotnet, resim def serve_forever(model, conf: float, predict_kw: dict, shotnet, resim, should_stop=None) -> int: """The serving loop over both queues -- the body of the worker, shared with the embedded mode. ``should_stop()`` (optional) lets the embedded mode exit cleanly when the app shuts down; the worker process does not pass it and exits on Ctrl+C. """ import redis served = 0 try: while True: if should_stop is not None and should_stop(): print(f"[cv_worker] stopping on request after {served} jobs", flush=True) break try: jobqueue.beat(key=jobqueue.CV_HEARTBEAT_KEY) t_job = time.perf_counter() # The 2 queues take turns with short waits (2 x 2s, about the # old 5s cadence) so the heartbeat keeps beating. segment # goes LAST: one segment job spawns N analyze jobs, and those # have to finish before a new video is accepted, so the shot # list of the video being watched fills in at the right pace. if jobqueue.serve_one( lambda p: handle_analyze(p, model, conf, predict_kw, shotnet=shotnet, resim=resim), timeout_s=2, jobs_key=jobqueue.ANALYZE_JOBS_KEY): served += 1 print(f"[cv_worker] analyze job #{served} done in " f"{time.perf_counter() - t_job:.2f}s", flush=True) elif jobqueue.serve_one( lambda p: handle_segment(p, model, conf, predict_kw), timeout_s=2, jobs_key=jobqueue.SEGMENT_JOBS_KEY): served += 1 print(f"[cv_worker] segment job #{served} done in " f"{time.perf_counter() - t_job:.2f}s", flush=True) except redis.RedisError as e: # a dead or restarting Redis must NOT kill a worker that has # already loaded its models print(f"[cv_worker] lost the Redis connection " f"({type(e).__name__}: {e}) -- retrying in 2s", flush=True) time.sleep(2) except KeyboardInterrupt: print(f"\n[cv_worker] Ctrl+C -- clean exit after {served} jobs", flush=True) finally: jobqueue.clear_heartbeat(key=jobqueue.CV_HEARTBEAT_KEY) if should_stop is None: # only the worker PROCESS may close the transport; the embedded # mode SHARES the client with the app, so tearing it down here # would kill the app too jobqueue.teardown() return 0 def main() -> int: if not os.environ.get("REDIS_URL"): print("cv_worker: the REDIS_URL env var is missing -- a standalone " "worker only makes sense with a queue (for example " "redis://localhost:6379/0). To run it in the SAME process as " "the app, set POOLCOACH_LOCAL_CV=1 when starting the app.", flush=True) return 2 conf = float(os.environ.get("POOLCOACH_CV_CONF", OP_CONF)) device = os.environ.get("POOLCOACH_CV_DEVICE", "") predict_kw = {"device": device} if device else {} jobqueue.setup() try: model, shotnet, resim = load_all(conf, device, predict_kw) except FileNotFoundError as e: print(f"cv_worker: {e}", flush=True) return 2 return serve_forever(model, conf, predict_kw, shotnet, resim) if __name__ == "__main__": sys.exit(main())