poolcoach-analyzer / scripts /cv_worker.py
masterdanh's picture
fix(analyzer): overlay BallID honours ball-set profile
5b4ac58
Raw History Blame Contribute Delete
54.6 kB
"""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())