Spaces:
Running
Running
Download scripts/cv_worker.py from masterdanh/poolcoach-analyzer: direct link, hf CLI and curl.
- Browser
- Download file 54.6 kB
-
https://huggingface.co/spaces/masterdanh/poolcoach-analyzer/resolve/main/scripts/cv_worker.py
- Command line
-
hf download hf://spaces/masterdanh/poolcoach-analyzer/scripts/cv_worker.py
-
curl -L -o cv_worker.py https://huggingface.co/spaces/masterdanh/poolcoach-analyzer/resolve/main/scripts/cv_worker.py
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()) | |