"""
claude_anim_v2.py — broadcast-ready NextJump overlay video from the ONNX export.
Pipeline
1. Pass 1 (analysis): decode the video once, building
- 256x256 model frames for the ONNX counting model, sampled at 60 fps
(the training rate - feeding e.g. 120 fps footage makes it over-count), and
- MoveNet-Thunder keypoints (with crop-region tracking) for framing.
2. Sliding-window ONNX inference -> period length, periodicity, marks, phase.
3. Keypoints -> a smoothed "virtual camera" crop path that keeps the jumper
(and the rope arc) well framed at the target aspect ratio.
4. Pass 2 (render): decode again, crop/zoom each frame along the camera path,
composite broadcast graphics and pipe straight into ffmpeg (audio copied).
Example
python claude_anim_v2.py --video ali.mp4 --format vertical
python claude_anim_v2.py --video 30s_wr_male_new_b_60fps.mp4 \
--format horizontal --title "30s SPEED | WORLD RECORD"
"""
import os
import sys
import math
import time
import shutil
import tempfile
import argparse
import subprocess
import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont
# scipy is imported lazily: the render worker processes never need it (each import costs memory)
def gaussian_filter1d(*args, **kwargs):
from scipy.ndimage import gaussian_filter1d as f
return f(*args, **kwargs)
def median_filter(*args, **kwargs):
from scipy.ndimage import median_filter as f
return f(*args, **kwargs)
def cpu_count():
"""CPUs this process may actually use (container / affinity limits, e.g. a 2-vCPU Space on a big host)."""
try:
return max(1, len(os.sched_getaffinity(0)))
except AttributeError:
return max(1, os.cpu_count() or 1)
def _running_median(x, size):
"""1-D median filter with edge replication (numpy only; used by the overlay in render workers)."""
if len(x) == 0:
return x
h = size // 2
xp = np.pad(np.asarray(x, float), h, mode='edge')
return np.median(np.lib.stride_tricks.sliding_window_view(xp, size), axis=1)
ROOT = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.join(ROOT, 'movenet-pytorch'))
# Output order of checkpoint/onnx/nextjump_all_128.onnx. The graph's output
# *names* are shifted by one (masks were prepended to RepNet.forward's return
# tuple), so always index by position, never by name.
OUT_MASKS, OUT_PERIOD, OUT_PERIODICITY, OUT_MARKS, OUT_EVENT, OUT_FOOT, OUT_PHASE = range(7)
# COCO-17 keypoint indices (MoveNet order)
NOSE, L_SHOULDER, R_SHOULDER, L_HIP, R_HIP, L_ANKLE, R_ANKLE = 0, 5, 6, 11, 12, 15, 16
SKELETON = [(0, 1), (0, 2), (1, 3), (2, 4), (5, 6), (5, 7), (7, 9), (6, 8), (8, 10),
(5, 11), (6, 12), (11, 12), (11, 13), (13, 15), (12, 14), (14, 16)]
# ── Palette (matches claude_anim.py) ──
C_BG = (13, 17, 23)
C_ACCENT = (88, 166, 255)
C_FEET = (255, 157, 0)
C_ROPE = (78, 205, 196)
C_GOOD = (63, 185, 80)
C_BAD = (248, 81, 73)
C_TEXT = (230, 237, 243)
C_MUTED = (139, 148, 158)
MODEL_FPS = 60.0 # the counting model was trained on 60 fps clips (see eval.py data dirs)
FORMATS = {'vertical': (1080, 1920), 'horizontal': (1920, 1080)}
def sigmoid(x):
return 1 / (1 + np.exp(-x))
def hex_rgb(h):
h = h.lstrip('#')
return tuple(int(h[i:i + 2], 16) for i in (0, 2, 4))
def find_ffmpeg():
# NEXTJUMP_FFMPEG lets you point at a newer build (e.g. one whose NVENC works with your driver)
override = os.environ.get('NEXTJUMP_FFMPEG')
if override and os.path.isfile(override):
return override
exe = shutil.which('ffmpeg')
if exe:
return exe
for cand in (os.path.join(sys.prefix, 'Library', 'bin', 'ffmpeg.exe'),
os.path.join(sys.prefix, 'bin', 'ffmpeg')):
if os.path.isfile(cand):
return cand
raise RuntimeError('ffmpeg not found on PATH or in the active Python environment')
def iter_frames(path, max_frames, mirror):
"""Yield RGB frames from *path* (optionally mirrored), at most *max_frames*."""
cap = cv2.VideoCapture(path)
n = 0
while n < max_frames:
ret, frame = cap.read()
if not ret:
break
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
if mirror:
frame = np.ascontiguousarray(frame[:, ::-1])
yield frame
n += 1
cap.release()
def iter_selected(path, max_frames, mirror, need):
"""Yield (i, BGR frame) only for source frames where need(i) is true. need() is called for every
frame in order. Unneeded frames are only grabbed (decoded but never colour-converted / copied),
~6x cheaper than reading them -- a 120 fps source analysed at 60 fps skips half the work."""
cap = cv2.VideoCapture(path)
try:
for i in range(max_frames):
if not cap.grab():
break
if need(i):
ok, frame = cap.retrieve()
if not ok:
break
if mirror:
frame = np.ascontiguousarray(frame[:, ::-1])
yield i, frame
finally:
cap.release()
def model_input(bgr, size=256):
"""Counting-model input (RGB, size x size). Two-step resize: bilinear to an integer multiple of the
target, then exact-ratio INTER_AREA -- within ~0.5 grey levels of a direct INTER_AREA, ~35% faster."""
h, w = bgr.shape[:2]
kx, ky = max(1, w // size), max(1, h // size)
if kx > 1 or ky > 1:
bgr = cv2.resize(bgr, (size * kx, size * ky), interpolation=cv2.INTER_LINEAR)
return cv2.cvtColor(cv2.resize(bgr, (size, size), interpolation=cv2.INTER_AREA), cv2.COLOR_BGR2RGB)
def schedule(n_frames, src_fps, fps):
"""Source-frame index of every frame of a *fps* timeline over n_frames source frames."""
rs = Resampler(src_fps, fps)
return np.array([i for i in range(n_frames) for _ in range(rs.take(i))], dtype=np.int64)
def _total_frames(path, max_frames):
cap = cv2.VideoCapture(path)
total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
cap.release()
return min(total, max_frames) if total > 0 else max_frames
# ════════════════════════════ pose tracking ════════════════════════════
class PoseTracker:
"""MoveNet-Thunder single-person pose with the standard crop-region tracker:
run on the whole (padded) frame until the jumper is found, then re-run on a
square crop around the previous pose for much higher keypoint precision."""
SIZE = 256
def __init__(self, device, conf=0.25):
from movenet.models.model_factory import load_model
self.model = load_model('movenet_thunder', ft_size=64,
model_dir=os.path.join(ROOT, 'movenet-pytorch', '_models'))
self.model = self.model.to(device).eval()
self.device = device
self.conf = conf
self.region = None # (cx, cy, size) in source pixels
def _full_region(self, w, h):
return (w / 2, h / 2, float(max(w, h)))
def crop(self, frame, region):
cx, cy, size = region
s = self.SIZE / size
M = np.float32([[s, 0, self.SIZE / 2 - cx * s], [0, s, self.SIZE / 2 - cy * s]])
return cv2.warpAffine(frame, M, (self.SIZE, self.SIZE), flags=cv2.INTER_AREA if s < 1 else cv2.INTER_LINEAR,
borderMode=cv2.BORDER_CONSTANT, borderValue=0)
@staticmethod
def _to_frame(kp, region):
cx, cy, size = region
out = np.empty_like(kp)
out[:, 0] = cx - size / 2 + kp[:, 1] * size # x (px)
out[:, 1] = cy - size / 2 + kp[:, 0] * size # y (px)
out[:, 2] = kp[:, 2]
return out
def _infer(self, frame, region):
cx, cy, size = region
crop = self.crop(frame, region)
import torch
with torch.no_grad():
x = torch.from_numpy(crop).float().unsqueeze(0).to(self.device)
kp = self.model(x)[0, 0].cpu().numpy() # (17, 3) = y, x, conf in [0, 1]
out = np.empty_like(kp)
out[:, 0] = cx - size / 2 + kp[:, 1] * size # x (px)
out[:, 1] = cy - size / 2 + kp[:, 0] * size # y (px)
out[:, 2] = kp[:, 2]
return out
def _next_region(self, kp, w, h):
good = kp[:, 2] > self.conf
torso = [L_SHOULDER, R_SHOULDER, L_HIP, R_HIP]
if good.sum() < 5 or not good[torso].any():
return None
hips = [i for i in (L_HIP, R_HIP) if good[i]] or [i for i in torso if good[i]]
cx, cy = kp[hips, 0].mean(), kp[hips, 1].mean()
pts = kp[good, :2]
body_half = max(np.abs(pts[:, 0] - cx).max(), np.abs(pts[:, 1] - cy).max())
torso_pts = kp[[i for i in torso if good[i]], :2]
torso_half = max(np.abs(torso_pts[:, 0] - cx).max(), np.abs(torso_pts[:, 1] - cy).max())
half = max(torso_half * 1.9, body_half * 1.25, 0.08 * max(w, h))
half = min(half, max(w, h) / 2)
return (cx, cy, 2 * half)
def __call__(self, frame, region=None):
if region is not None: # seeded from the selected jumper's detector box: never drifts to someone else
return self._infer(frame, region)
h, w = frame.shape[:2]
region = self.region or self._full_region(w, h)
kp = self._infer(frame, region)
self.region = self._next_region(kp, w, h)
if self.region is None and region != self._full_region(w, h):
# Lost the jumper inside the crop — retry on the full frame.
kp = self._infer(frame, self._full_region(w, h))
self.region = self._next_region(kp, w, h)
return kp
# ════════════════════════════ jumper selection ════════════════════════════
#
# Pose-only tracking follows whoever MoveNet happens to lock on to, which can be a judge. "smart"
# framing instead (1) detects every person a few times per second and links the boxes into tracks,
# (2) probes each candidate by running the counting model on a short clip cropped around them and
# scoring how periodic (= actively jumping) they look, and (3) frames the N most active people.
def square_crop(frame, cx, cy, side, size):
"""Square crop centred on (cx, cy), zero-padded outside the frame, resized to size x size."""
h, w = frame.shape[:2]
half = max(2, int(round(side / 2)))
x0, y0 = int(round(cx)) - half, int(round(cy)) - half
x1, y1 = x0 + 2 * half, y0 + 2 * half
crop = frame[max(0, y0):min(h, y1), max(0, x0):min(w, x1)]
if crop.size == 0:
return np.zeros((size, size, 3), np.uint8)
crop = cv2.copyMakeBorder(crop, max(0, -y0), max(0, y1 - h), max(0, -x0), max(0, x1 - w), cv2.BORDER_CONSTANT, value=0)
return cv2.resize(crop, (size, size), interpolation=cv2.INTER_AREA if 2 * half > size else cv2.INTER_LINEAR)
def load_person_detector(spec, device):
"""'coco' -> torchvision Faster R-CNN MobileNetV3 (COCO person class; weights auto-download once).
Otherwise a path to a torch-pickled torchvision detector whose class 1 is a person / jumper
(e.g. NextJump-Tournament's fasterrcnn_model_finetuned.pth)."""
if spec in (None, '', 'coco'):
from torchvision.models.detection import (fasterrcnn_mobilenet_v3_large_fpn,
FasterRCNN_MobileNet_V3_Large_FPN_Weights)
model = fasterrcnn_mobilenet_v3_large_fpn(weights=FasterRCNN_MobileNet_V3_Large_FPN_Weights.COCO_V1,
progress=False)
else:
import torch
model = torch.load(spec, map_location='cpu', weights_only=False)
return model.to(device).eval()
def detect_people(model, frame, device, score_thr=0.6, min_height=0.12, max_side=960):
"""Person boxes (x0, y0, x1, y1) in source pixels; tiny (far-away) people are dropped."""
h, w = frame.shape[:2]
sc = min(1.0, max_side / max(h, w))
img = cv2.resize(frame, (int(w * sc), int(h * sc)), interpolation=cv2.INTER_AREA) if sc < 1 else frame
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # frames arrive as BGR from iter_selected
import torch
with torch.no_grad():
x = torch.from_numpy(np.ascontiguousarray(img)).to(device).permute(2, 0, 1).float() / 255
pred = model([x])[0]
keep = (pred['labels'] == 1) & (pred['scores'] >= score_thr)
boxes = pred['boxes'][keep].cpu().numpy() / sc
scores = pred['scores'][keep].cpu().numpy()
tall = (boxes[:, 3] - boxes[:, 1]) >= min_height * h
return boxes[tall], scores[tall]
def _iou(a, b):
x0 = np.maximum(a[:, None, 0], b[None, :, 0])
y0 = np.maximum(a[:, None, 1], b[None, :, 1])
x1 = np.minimum(a[:, None, 2], b[None, :, 2])
y1 = np.minimum(a[:, None, 3], b[None, :, 3])
inter = np.clip(x1 - x0, 0, None) * np.clip(y1 - y0, 0, None)
area = lambda r: (r[:, 2] - r[:, 0]) * (r[:, 3] - r[:, 1])
return inter / (area(a)[:, None] + area(b)[None, :] - inter + 1e-9)
def build_tracks(dets, max_gap=1.0, min_iou=0.2, min_coverage=0.15):
"""Greedy IoU linking of per-frame detections into person tracks."""
tracks = []
for t, boxes, scores in dets:
used = set()
live = [k for k, tr in enumerate(tracks) if t - tr['t'][-1] <= max_gap]
if live and len(boxes):
iou = _iou(np.array([tracks[k]['b'][-1] for k in live]), boxes)
for a, b in sorted(np.argwhere(iou > min_iou), key=lambda ab: -iou[ab[0], ab[1]]):
if live[a] in used or ('_b', b) in used:
continue
used.update((live[a], ('_b', b)))
tr = tracks[live[a]]
tr['t'].append(t), tr['b'].append(boxes[b]), tr['s'].append(scores[b])
for b in range(len(boxes)):
if ('_b', b) not in used:
tracks.append(dict(t=[t], b=[boxes[b]], s=[scores[b]]))
n_det = max(1, len(dets))
out = []
for tr in tracks:
if len(tr['t']) >= max(3, min_coverage * n_det):
out.append(dict(t=np.array(tr['t']), b=np.array(tr['b']), s=np.array(tr['s'])))
return out
def prefilter_tracks(tracks, duration, keep):
"""Cheap ranking before the (expensive) model probe: people who are on screen for long, are large
(near the camera) and whose box bounces vertically (jumping) come first."""
def cheap(tr):
h = tr['b'][:, 3] - tr['b'][:, 1]
coverage = min(1.0, (tr['t'][-1] - tr['t'][0]) / max(duration, 1e-6)) * len(tr['t']) / max(1, len(tr['t']))
bounce = np.std(np.diff(tr['b'][:, 3])) / max(1.0, np.median(h)) if len(h) > 2 else 0.0
return coverage * np.median(h) * (1 + 10 * bounce)
order = sorted(range(len(tracks)), key=lambda k: -cheap(tracks[k]))
return [tracks[k] for k in order[:keep]]
def track_boxes(tr, times):
"""Track box linearly interpolated (held constant outside its span) at the given times -> (T, 4)."""
times = np.atleast_1d(times)
return np.stack([np.interp(times, tr['t'], tr['b'][:, k]) for k in range(4)], -1)
def probe_windows(duration, win_s, event=None, n=5):
lo, hi = (event[0] + 0.5, event[1]) if event is not None else (0.0, duration)
if hi - lo <= win_s:
return [(max(0.0, lo), max(0.0, lo) + win_s)]
centres = [lo + (hi - lo) * f for f in np.linspace(0.1, 0.9, n)]
return [(c - win_s / 2, c + win_s / 2) for c in centres]
def scan_people(path, max_frames, mirror, src_fps, model_fps, det_fps, detector, device, windows,
probe_side=640, det_side=960, model_sink=None, img_size=256, progress=None):
"""Decode pass: person detections at det_fps, downscaled frames (at model_fps) in the probe windows
and -- with collect_model -- the full-frame counting-model input, so no second pass is needed for it.
Returns dict(dets, probes, scale, model_frames, n_frames, frame_size)."""
total = _total_frames(path, max_frames)
det_rs, pr_rs, mo_rs = Resampler(src_fps, det_fps), Resampler(src_fps, model_fps), Resampler(src_fps, model_fps)
dets, probes, model_frames, scale = [], [[] for _ in windows], model_sink, None
collect_model = model_sink is not None
cur, last = [0, 0, 0], [-1]
def need(i):
t = i / src_fps
nd = det_rs.take(i)
npr = pr_rs.take(i)
nm = mo_rs.take(i) if collect_model else 0
cur[:] = nd, (npr if npr and any(a <= t < b for a, b in windows) else 0), nm
last[0] = i
if i % 50 == 0:
print(f'\rFinding people {i + 1}/{total} ', end='', flush=True)
if progress:
progress(min(1.0, (i + 1) / max(1, total)))
return any(cur)
from concurrent.futures import ThreadPoolExecutor
size, pending = None, []
with ThreadPoolExecutor(1) as det_pool: # detection (GPU, releases the GIL) overlaps decoding
for i, frame in iter_selected(path, max_frames, mirror, need):
nd, npr, nm = cur
size = (frame.shape[1], frame.shape[0])
t = i / src_fps
if nd:
pending.append((t, det_pool.submit(detect_people, detector, frame, device, max_side=det_side)))
if npr:
if scale is None:
scale = min(1.0, probe_side / max(frame.shape[:2]))
small = cv2.resize(frame, None, fx=scale, fy=scale, interpolation=cv2.INTER_AREA) if scale < 1 else frame
small = cv2.cvtColor(small, cv2.COLOR_BGR2RGB)
for w, (a, b) in enumerate(windows):
if a <= t < b:
probes[w] += [(t, small)] * npr
if nm:
model_frames += [model_input(frame, img_size)] * nm
dets = [(t, *fu.result()) for t, fu in pending]
print(f'\rFound {sum(len(d[1]) for d in dets)} person detections in {len(dets)} frames ')
return dict(dets=dets, probes=probes, scale=scale or 1.0, n_frames=last[0] + 1, frame_size=size,
model_frames=model_frames)
def score_tracks(sess, tracks, probes, scale, batch_size=1, progress=None, gpu=None):
"""Activity per track: counting-model periodicity on a crop around the person in each probe window,
averaged over that person's best half of the windows (so idle time before / after the event, or a
short pause, doesn't hide a jumper; bystanders never look periodic in any window)."""
seq_len = sess.get_inputs()[0].shape[1]
per_window = np.zeros((len(tracks), len(probes)))
jobs = []
for w, frames in enumerate(probes):
if len(frames) < seq_len // 2:
continue
frames = (frames + [frames[-1]] * seq_len)[:seq_len]
tm = np.array([t for t, _ in frames])
for k, tr in enumerate(tracks):
if not (tr['t'][0] - 0.3 <= tm[len(tm) // 2] <= tr['t'][-1] + 0.3):
continue # not on screen during this window -> scores 0 here
B = track_boxes(tr, tm) * scale
cx = gaussian_filter1d((B[:, 0] + B[:, 2]) / 2, 8, mode='nearest')
cy = gaussian_filter1d((B[:, 1] + B[:, 3]) / 2, 8, mode='nearest')
side = 1.4 * np.max(np.maximum(B[:, 2] - B[:, 0], B[:, 3] - B[:, 1]))
jobs.append((k, w, frames, cx, cy, side)) # crops are built lazily (memory)
for b in range(0, len(jobs), batch_size):
chunk = jobs[b:b + batch_size]
X = np.stack([np.stack([square_crop(f, x, y, side, 256) for (_, f), x, y in zip(frames, cx, cy)])
for _, _, frames, cx, cy, side in chunk]).astype(np.float32).transpose(0, 1, 4, 2, 3)
out = gpu.submit(sess.run, None, {'video': X}).result() if gpu else sess.run(None, {'video': X})
for j, (k, w, *_) in enumerate(chunk):
per_window[k, w] = float(np.mean(sigmoid(out[OUT_PERIODICITY][j].reshape(seq_len, -1)[:, 0])))
if progress:
progress(min(1.0, (b + batch_size) / len(jobs)))
used = [w for w, fr in enumerate(probes) if len(fr) >= seq_len // 2]
if not used:
return np.zeros(len(tracks))
top = max(1, int(math.ceil(len(used) / 2)))
return np.sort(per_window[:, used], 1)[:, ::-1][:, :top].mean(1)
def select_jumpers(tracks, scores, n, frame_w=None, spread_penalty=0.6, duration=None):
"""The most active person, then (for groups) the most active people near them: each addition is
ranked by activity minus a penalty for distance from the group so far (keeps one team together).
Activity is weighted by the fraction of the video each person is on screen, which breaks ties between
equally rhythmic people in favour of the athletes (warm-up jumpers / passers-by come and go)."""
if duration:
cover = np.array([min(1.0, (tr['t'][-1] - tr['t'][0]) / duration) for tr in tracks])
scores = np.asarray(scores) * (0.5 + 0.5 * cover)
centre = lambda k: ((tracks[k]['b'][:, 0] + tracks[k]['b'][:, 2]) / 2).mean()
fw = frame_w or max(1.0, max(tr['b'][:, 2].max() for tr in tracks))
sel = [int(np.argmax(scores))] if len(scores) else []
while len(sel) < min(n, len(tracks)):
gx = np.mean([centre(k) for k in sel])
rest = [k for k in range(len(tracks)) if k not in sel]
sel.append(max(rest, key=lambda k: scores[k] - spread_penalty * abs(centre(k) - gx) / fw))
print('Jumper candidates (activity = periodicity of the counting model on a crop around each person):')
for k in np.argsort(-scores)[:8]:
tr = tracks[k]
b = tr['b'].mean(0)
print(f' {"*" if k in sel else " "} #{k:<3d} activity {scores[k]:.2f} on screen {tr["t"][0]:5.1f}-{tr["t"][-1]:5.1f} s '
f'box ~({b[0]:.0f},{b[1]:.0f})-({b[2]:.0f},{b[3]:.0f})')
return sel
def dets_frame_w(dets, path):
cap = cv2.VideoCapture(path)
w = cap.get(cv2.CAP_PROP_FRAME_WIDTH)
cap.release()
return w or None
def group_box(tracks, sel, times):
"""Box around the selected jumpers, centred on their centre of mass and wide enough to hold all of them."""
B = np.stack([track_boxes(tracks[k], times) for k in sel]) # (N, T, 4)
cx = ((B[..., 0] + B[..., 2]) / 2).mean(0)
cy = ((B[..., 1] + B[..., 3]) / 2).mean(0)
hw = np.maximum(cx - B[..., 0].min(0), B[..., 2].max(0) - cx)
hh = np.maximum(cy - B[..., 1].min(0), B[..., 3].max(0) - cy)
return np.stack([cx - hw, cy - hh, cx + hw, cy + hh], -1)
# ════════════════════════════ beep detection ════════════════════════════
def _stream_start(ffprobe, path, kind):
try:
out = subprocess.run([ffprobe, '-v', 'error', '-select_streams', f'{kind}:0', '-show_entries',
'stream=start_time', '-of', 'csv=p=0', path], capture_output=True, text=True, check=True)
return float(out.stdout.strip().splitlines()[0])
except (subprocess.CalledProcessError, ValueError, IndexError, OSError):
return 0.0
def detect_beeps(video_path, ffmpeg, event_length, beep_file, beep_height=0.8, tolerance=0.5):
"""Find the official start / end beeps (same strategy as NextJump-Tournament/app.py:
cross-correlate the audio with the reference beep, then take the pair of peaks whose spacing
is closest to the event length, lowering the peak threshold until a good pair is found).
Differences from the app: times are the *onset* of the beep tone (the app's mode='same'
correlation peak sits ~0.18 s after the onset, i.e. about one jump late), and the
audio/video stream start offset is compensated. Returns (start_s, end_s) in video time, or None."""
from scipy.io import wavfile
from scipy.signal import correlate, find_peaks
fs, ref = wavfile.read(beep_file)
ref = ref.astype(np.float64)
if ref.ndim > 1:
ref = ref.mean(1)
env = np.abs(ref)
ref_onset = np.argmax(env > 0.1 * env.max()) / fs # tone starts ~20 ms into beep.WAV
proc = subprocess.run([ffmpeg, '-v', 'error', '-i', video_path, '-vn', '-ac', '1', '-ar', str(fs),
'-f', 's16le', '-'], capture_output=True)
audio = np.frombuffer(proc.stdout, np.int16).astype(np.float64)
if proc.returncode != 0 or len(audio) < len(ref) * 2:
print('Beep detection: no usable audio track')
return None
ffprobe = os.path.join(os.path.dirname(ffmpeg), 'ffprobe' + os.path.splitext(ffmpeg)[1])
offset = _stream_start(ffprobe, video_path, 'a') - _stream_start(ffprobe, video_path, 'v')
corr = correlate(audio, ref, mode='valid') # index = sample where the reference starts
corr = 2 * (corr - corr.min()) / (corr.max() - corr.min() + 1e-12) - 1
best, best_diff = None, float('inf')
h = beep_height
while h >= 0.3 - 1e-9:
peaks, _ = find_peaks(corr, height=h, distance=fs // 2)
for a in range(len(peaks)):
for b in range(a + 1, len(peaks)):
diff = abs((peaks[b] - peaks[a]) / fs - event_length)
if diff < best_diff:
best_diff, best = diff, (peaks[a], peaks[b])
if best_diff < 0.25:
break
h -= 0.05
if best is None or best_diff > tolerance:
print(f'Beep detection: no beep pair ~{event_length:g} s apart found'
+ (f' (closest was off by {best_diff:.2f} s)' if best else ''))
return None
t0, t1 = best[0] / fs + ref_onset + offset, best[1] / fs + ref_onset + offset
print(f'Beeps: start {t0:.3f} s, end {t1:.3f} s (event {t1 - t0:.3f} s)')
return t0, t1
# ════════════════════════════ analysis pass ════════════════════════════
class Resampler:
"""Nearest-frame resampling from src_fps to dst_fps, fed one source frame at a time.
take(i) returns how many destination samples use source frame i (0, 1, or >1 when upsampling)."""
def __init__(self, src_fps, dst_fps):
self.r, self.k = src_fps / dst_fps, 0
def take(self, i):
n = 0
while int(math.floor(self.k * self.r + 0.5)) == i:
n += 1
self.k += 1
return n
def pick_rate(src_fps, target_fps):
"""Rate closest to *target_fps* reachable from *src_fps* by keeping every n-th frame
(integer decimation => even frame spacing, no judder). Never upsamples."""
if target_fps is None or src_fps <= target_fps * 1.2:
return src_fps
return src_fps / max(1, round(src_fps / target_fps))
def analyse_video(path, max_frames, mirror, img_size, device, src_fps, model_fps, out_fps, run_pose=True,
pose_fps=None, collect_model=True, progress=None, pose_region=None, model_crop=None,
model_sink=None):
"""Decode pass building
- model frames at *model_fps* (the rate the counting model was trained on) if collect_model,
- keypoints at *pose_fps* (default: every output frame), interpolated onto the output frames,
- the source-frame index of every output frame (at *out_fps*).
pose_region(t) -> (cx, cy, size) seeds MoveNet on the selected jumper; model_crop(t) -> (cx, cy, side)
feeds the counting model a square crop around the selected jumper(s) instead of the full frame.
Returns (model_frames | None, keypoints | None, frame_size, out_src)."""
total = _total_frames(path, max_frames)
tracker = PoseTracker(device) if run_pose else None
model_rs, out_rs = Resampler(src_fps, model_fps), Resampler(src_fps, out_fps)
pose_rs = Resampler(src_fps, pose_fps or out_fps)
model_frames = model_sink if model_sink is not None else []
pose_t, pose_kp, out_src = [], [], []
cur, last = [0, 0], [-1]
def need(i):
nm = model_rs.take(i) if collect_model else 0
out_src.extend([i] * out_rs.take(i))
npose = pose_rs.take(i) if tracker is not None else 0
cur[:] = nm, npose
last[0] = i
if i % 50 == 0:
print(f'\rAnalysing frame {i + 1}/{total} ', end='', flush=True)
if progress:
progress(min(1.0, (i + 1) / max(1, total)))
return any(cur)
size = None
for i, frame in iter_selected(path, max_frames, mirror, need):
nm, npose = cur
size = (frame.shape[1], frame.shape[0])
t = i / src_fps
if nm:
if model_crop is not None:
model_frames += [cv2.cvtColor(square_crop(frame, *model_crop(t), img_size), cv2.COLOR_BGR2RGB)] * nm
else:
# Same preprocessing as training / graphics/main.py: plain resize to a square.
model_frames += [model_input(frame, img_size)] * nm
if npose:
pose_t.append(t)
pose_kp.append(tracker(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB), pose_region(t) if pose_region else None))
out_src = np.array(out_src, dtype=np.int64)
kps = None
if pose_kp:
P, tt = np.stack(pose_kp), out_src / src_fps
kps = np.stack([np.stack([np.interp(tt, pose_t, P[:, j, c]) for c in range(3)], -1) for j in range(17)], 1)
print(f'\rAnalysed {last[0] + 1} frames ({size[0]}x{size[1]} @ {src_fps:.2f} fps) -> '
+ (f'model {len(model_frames)} @ {model_fps:.2f} fps, ' if collect_model else '')
+ (f'pose {len(pose_kp)} @ {pose_fps or out_fps:.1f} fps, ' if pose_kp else '')
+ f'output {len(out_src)} @ {out_fps:.2f} fps')
if not collect_model:
model_frames = None
elif model_sink is None:
model_frames = np.stack(model_frames)
return model_frames, kps, size, out_src
def make_session(onnx_path, use_cuda):
import torch # imported before onnxruntime so its CUDA DLLs are loaded
import onnxruntime as ort
opts = ort.SessionOptions()
opts.log_severity_level = 4 # silence the noisy CUDA-provider fallback errors
providers = ['CPUExecutionProvider']
if use_cuda and 'CUDAExecutionProvider' in ort.get_available_providers():
# The default arena grows by powers of two until VRAM is full; on Windows the driver then
# spills into shared system memory and every window gets ~10x slower. Allocate exactly
# what is requested and cap the arena so torch (MoveNet / detector) keeps headroom.
total = torch.cuda.get_device_properties(0).total_memory if torch.cuda.is_available() else 8 << 30
providers.insert(0, ('CUDAExecutionProvider', {'arena_extend_strategy': 'kSameAsRequested',
'gpu_mem_limit': str(int(total * 0.6)),
'cudnn_conv_algo_search': 'HEURISTIC'}))
sess = ort.InferenceSession(onnx_path, sess_options=opts, providers=providers)
if use_cuda and 'CUDAExecutionProvider' not in sess.get_providers():
print('Note: onnxruntime CUDA provider unavailable (CUDA 12 / cuDNN 9 libs not found) - running ONNX on CPU')
print(f'ONNX providers: {sess.get_providers()}')
return sess
def run_model(sess, frames, stride, batch_size, progress=None):
"""Sliding-window inference with overlap averaging. *frames*: (N, H, W, 3) uint8."""
seq_len = sess.get_inputs()[0].shape[1]
n = len(frames)
if n < seq_len:
frames = np.concatenate([frames, np.repeat(frames[-1:], seq_len - n, axis=0)])
starts = list(range(0, len(frames) - seq_len + 1, stride))
if starts[-1] != len(frames) - seq_len:
starts.append(len(frames) - seq_len)
m = len(frames)
acc = {k: np.zeros((m, d)) for k, d in (('period', 2), ('periodicity', 1), ('marks', 1), ('phase', 2))}
overlaps = np.zeros((m, 1))
for b in range(0, len(starts), batch_size):
chunk = starts[b:b + batch_size]
X = np.stack([frames[s:s + seq_len] for s in chunk]).astype(np.float32).transpose(0, 1, 4, 2, 3)
try:
out = sess.run(None, {'video': X})
except Exception as e:
if len(chunk) == 1 or 'allocate' not in str(e):
raise
# out of (GPU) memory for this batch size: run the windows one at a time instead
outs = [sess.run(None, {'video': X[j:j + 1]}) for j in range(len(chunk))]
out = [np.concatenate([o[k] for o in outs]) for k in range(len(outs[0]))]
for j, s in enumerate(chunk):
acc['period'][s:s + seq_len] += out[OUT_PERIOD][j].reshape(seq_len, -1)[:, :2]
acc['periodicity'][s:s + seq_len] += out[OUT_PERIODICITY][j].reshape(seq_len, -1)[:, :1]
acc['marks'][s:s + seq_len] += out[OUT_MARKS][j].reshape(seq_len, -1)[:, :1]
acc['phase'][s:s + seq_len] += out[OUT_PHASE][j].reshape(seq_len, -1)[:, :2]
overlaps[s:s + seq_len] += 1
print(f'\rInference window {min(b + batch_size, len(starts))}/{len(starts)} ', end='', flush=True)
if progress:
progress(min(b + batch_size, len(starts)) / len(starts))
print()
res = {k: (v / np.maximum(overlaps, 1))[:n] for k, v in acc.items()}
return res
def _run_window(sess, frames):
L = len(frames)
X = np.stack(frames)[None].astype(np.float32).transpose(0, 1, 4, 2, 3)
out = sess.run(None, {'video': X})
return (out[OUT_PERIOD][0].reshape(L, -1)[:, :2], out[OUT_PERIODICITY][0].reshape(L, -1)[:, :1],
out[OUT_MARKS][0].reshape(L, -1)[:, :1], out[OUT_PHASE][0].reshape(L, -1)[:, :2])
class StreamingCounter(list):
"""The list of counting-model frames, which runs each sliding window on the (single-thread) GPU queue
as soon as its frames have arrived -- inference overlaps decoding instead of waiting for it.
Same windows and overlap averaging as run_model()."""
def __init__(self, sess, stride, gpu):
super().__init__()
self.sess, self.stride, self.gpu = sess, stride, gpu
self.L = sess.get_inputs()[0].shape[1]
self.next, self.jobs = 0, {}
def __iadd__(self, items):
self.extend(items)
while self.next + self.L <= len(self):
self.jobs[self.next] = self.gpu.submit(_run_window, self.sess, self[self.next:self.next + self.L])
self.next += self.stride
return self
def finish(self, progress=None):
n = len(self)
frames = list(self) + [self[-1]] * max(0, self.L - n)
last = len(frames) - self.L
if last not in self.jobs:
self.jobs[last] = self.gpu.submit(_run_window, self.sess, frames[last:last + self.L])
m = len(frames)
acc = [np.zeros((m, d)) for d in (2, 1, 1, 2)]
overlaps = np.zeros((m, 1))
for k, (st, fu) in enumerate(sorted(self.jobs.items())):
for a, o in zip(acc, fu.result()):
a[st:st + self.L] += o
overlaps[st:st + self.L] += 1
print(f'\rInference window {k + 1}/{len(self.jobs)} ', end='', flush=True)
if progress:
progress((k + 1) / len(self.jobs))
print()
return {key: (a / np.maximum(overlaps, 1))[:n] for key, a in zip(('period', 'periodicity', 'marks', 'phase'), acc)}
def compute_analytics(res, fps, periodicity_thresh, mark_thresh, event=None, skip_after_miss=False,
min_miss_s=0.25):
"""*event*: optional (start_s, end_s) from the beeps. When given, counting runs only between
the beeps: nothing before the start beep, frozen at the end beep.
*skip_after_miss*: rulebook judging — the first jump after each miss is not counted."""
from scipy.signal import medfilt, find_peaks
period_legs = res['period'][:, 0]
period_arms = res['period'][:, 1]
periodicity = sigmoid(medfilt(res['periodicity'][:, 0], 5))
marks = sigmoid(res['marks'][:, 0])
phase = np.arctan2(-res['phase'][:, 1], res['phase'][:, 0]) # negated sin: bottom = start of jump
active = (periodicity > periodicity_thresh).astype(float)
jps_legs = np.where(period_legs > 2, fps / np.maximum(period_legs, 1e-6), 0)
jps_arms = np.where(period_arms > 2, fps / np.maximum(period_arms, 1e-6), 0)
jps_legs = gaussian_filter1d(jps_legs, 2) * active
jps_arms = gaussian_filter1d(jps_arms, 2) * active
n = len(marks)
if event is not None:
# first model frame at/after the start beep .. last model frame before the end beep
ev_s = int(min(n - 1, math.ceil(event[0] * fps - 1e-6)))
ev_e = int(max(ev_s, min(n - 1, math.ceil(event[1] * fps - 1e-6) - 1)))
in_ev = np.zeros(n, bool)
in_ev[ev_s:ev_e + 1] = True
active = active * in_ev
jps_legs, jps_arms = jps_legs * in_ev, jps_arms * in_ev
peaks, _ = find_peaks(marks, distance=3, height=mark_thresh)
is_peak = np.zeros(len(marks), bool)
is_peak[peaks] = True
# Count accumulation — identical logic to claude_anim.py / graphics/main.py
count = np.zeros(len(marks))
c = 0.0
for k in range(len(marks)):
if event is not None and k > ev_e:
count[k] = count[k - 1] # frozen at the end beep
continue
if period_legs[k] < 2 or not active[k]:
c = float(int(c))
elif is_peak[k]:
c = float(int(c)) if math.modf(c)[0] < 0.1 else float(int(c) + 1.01)
else:
c += 1.0 / (period_legs[k] + 1e-8)
count[k] = c
count_int = count.astype(int)
if event is not None:
jump_start, jump_end = ev_s, ev_e
else:
started = np.where(count_int > 0)[0]
jump_start = int(started[0]) if len(started) else 0
act_idx = np.where(active > 0)[0]
jump_end = int(act_idx[-1]) if len(act_idx) else len(count) - 1
jump_end = max(jump_end, jump_start)
# Start pace = mean speed over the first 3 s of actual jumping (from the first counted jump,
# so the reaction time after the start beep doesn't drag it down); pace line starts there too.
counted = np.where(count_int[jump_start:jump_end + 1] > 0)[0]
first_jump = jump_start + int(counted[0]) if len(counted) else jump_start
win = min(int(3 * fps), jump_end + 1 - first_jump)
baseline = float(np.mean(jps_legs[first_jump:first_jump + win])) if win > 0 else 0.0
pace = np.zeros(len(count))
t = (np.arange(len(count)) - first_jump) / fps
pace[first_jump:] = baseline * np.clip(t[first_jump:], 0, (jump_end - first_jump) / fps)
# Misses: the model's periodicity drops out for >= min_miss_s after jumping has started.
# Rulebook option: the first counted jump after each miss is not counted (judges' realtime rule).
misses, skipped = [], []
inact = np.r_[False, active[first_jump + 1:jump_end + 1] == 0, False]
edges = np.flatnonzero(np.diff(inact.astype(int)))
for a, b in zip(edges[::2] + first_jump, edges[1::2] + first_jump): # [a, b) inactive model frames
if b - a < min_miss_s * fps:
continue
misses.append((a, b))
if skip_after_miss:
inc = np.flatnonzero(np.diff(count_int[b - 1:jump_end + 1]) > 0)
if len(inc):
f = b + int(inc[0]) # frame where the first jump after the miss is counted
skipped.append(f)
count_int[f:] -= 1
count_int = np.maximum(count_int, 0)
if misses:
print(f'Misses: {len(misses)} at ' + ', '.join(f'{a / fps:.1f} s' for a, _ in misses)
+ (f' -> {len(skipped)} jump(s) after a miss not counted' if skip_after_miss else ''))
final = int(count_int[jump_end])
# Per-jump timestamps from the mark peaks, refined to sub-frame precision (parabolic fit) so
# rhythm stats aren't quantised to whole frames.
bt = []
for k in peaks:
if jump_start <= k <= jump_end and 0 < k < len(marks) - 1:
a, b, c = marks[k - 1], marks[k], marks[k + 1]
den = a - 2 * b + c
bt.append(k + (0.5 * (a - c) / den if den < 0 else 0.0))
beat_times = np.array(bt) / fps
print(f'Jump window: frame {jump_start} -> {jump_end} ({(jump_end - jump_start) / fps:.1f} s), '
f'{final} jumps, start pace {baseline:.2f} JPS')
return dict(jps=jps_legs, jps_arms=jps_arms, periodicity=periodicity, active=active, marks=marks,
phase=phase, count=count_int, pace=pace, jump_start=jump_start, jump_end=jump_end,
final=final, baseline=baseline, beat_times=beat_times, event=event, first_jump=first_jump,
miss_times=[(a / fps, b / fps) for a, b in misses], skip_times=[f / fps for f in skipped],
skip_after_miss=skip_after_miss)
def to_output_timeline(an, model_fps, times):
"""Resample per-model-frame analytics onto output frames at the given times (seconds)."""
n = len(an['count'])
j = np.clip(np.asarray(times) * model_fps, 0, n - 1)
grid = np.arange(n)
out = dict(an)
for k in ('jps', 'jps_arms', 'periodicity', 'marks', 'pace'):
out[k] = np.interp(j, grid, an[k])
out['active'] = an['active'][np.round(j).astype(int)]
out['phase'] = np.arctan2(np.interp(j, grid, np.sin(an['phase'])), np.interp(j, grid, np.cos(an['phase'])))
out['count'] = an['count'][np.floor(j + 1e-6).astype(int)]
times = np.asarray(times)
out['times'] = times
out['first_jump'] = int(min(np.searchsorted(j, an['first_jump']), len(j) - 1))
if an.get('event') is not None:
# exact beep times: first output frame at/after the start beep, last frame before the end beep
out['jump_start'] = int(min(np.searchsorted(times, an['event'][0] - 1e-6), len(j) - 1))
out['jump_end'] = int(max(out['jump_start'], np.searchsorted(times, an['event'][1] - 1e-6) - 1))
out['count'][:out['jump_start']] = 0
else:
out['jump_start'] = int(min(np.searchsorted(j, an['jump_start']), len(j) - 1))
out['jump_end'] = int(max(out['jump_start'], np.searchsorted(j, an['jump_end'], side='right') - 1))
if an.get('event') is not None:
out['count'][out['jump_end']:] = an['final']
else:
out['count'][out['jump_end']:] = np.maximum(out['count'][out['jump_end']:], an['final'])
return out
# ════════════════════════════ virtual camera ════════════════════════════
def _interp_invalid(x, valid):
idx = np.arange(len(x))
if valid.sum() == 0:
return None
return np.interp(idx, idx[valid], x[valid])
def boxes_from_keypoints(kps, conf=0.3):
"""Per-frame body box from keypoints, padded from eyes/ankles to head-top/soles; NaN where unreliable."""
good = kps[:, :, 2] > conf
valid = good.sum(1) >= 6
x = np.where(good, kps[:, :, 0], np.nan)
y = np.where(good, kps[:, :, 1], np.nan)
with np.errstate(all='ignore'):
x0, x1 = np.nanmin(x, 1), np.nanmax(x, 1)
y0, y1 = np.nanmin(y, 1), np.nanmax(y, 1)
span = y1 - y0
boxes = np.stack([x0, y0 - 0.12 * span, x1, y1 + 0.06 * span], 1)
boxes[~valid] = np.nan
return boxes
def compute_camera_path(boxes, frame_size, out_size, fps, fill, anchor_y, max_zoom, smooth_sec):
"""Turn per-frame subject boxes (x0, y0, x1, y1; NaN rows = unknown) into a smooth crop path
(cx, cy, crop_w, crop_h).
The crop height is chosen so the jumper's body occupies *fill* of the frame
height (leaving room for the rope arc), the body centre sits at *anchor_y*
of the crop height, and the path is heavily low-passed so the camera glides
rather than bouncing with every jump."""
W, H = frame_size
aspect = out_size[0] / out_size[1]
n = len(boxes) if boxes is not None else 0
# Largest crop of the output aspect that fits inside the source.
max_h = min(H, W / aspect)
if n == 0:
return np.tile([W / 2, H / 2, max_h * aspect, max_h], (1, 1)), False
valid = np.isfinite(boxes).all(1)
if valid.sum() < max(5, 0.1 * n):
print('Warning: jumper rarely detected — falling back to a centred crop')
return np.tile([W / 2, H / 2, max_h * aspect, max_h], (n, 1)), False
x0, y0, x1, y1 = boxes.T
cx = _interp_invalid((x0 + x1) / 2, valid)
cy = _interp_invalid((y0 + y1) / 2, valid)
bh = _interp_invalid(y1 - y0, valid)
bw = _interp_invalid(x1 - x0, valid)
# Robust smoothing: median kills pose glitches, gaussian removes jump bounce.
med = max(3, int(fps * 0.5) | 1)
sig = max(1.0, smooth_sec * fps)
cx = gaussian_filter1d(median_filter(cx, med, mode='nearest'), sig, mode='nearest')
cy = gaussian_filter1d(median_filter(cy, med, mode='nearest'), sig, mode='nearest')
# Size: rolling max (arms overhead / rope) then slower smoothing to avoid "breathing" zoom.
bh = gaussian_filter1d(median_filter(bh, int(fps) | 1, mode='nearest'), sig * 2, mode='nearest')
bw = gaussian_filter1d(median_filter(bw, int(fps) | 1, mode='nearest'), sig * 2, mode='nearest')
crop_h = np.maximum(bh / fill, (bw * 1.35) / aspect) # keep arms + rope width in frame
crop_h = np.clip(crop_h, max_h / max_zoom, max_h)
crop_w = crop_h * aspect
# Place body centre at anchor_y of the crop, then keep the crop inside the frame.
ccx = np.clip(cx, crop_w / 2, W - crop_w / 2)
ccy = cy + (0.5 - anchor_y) * crop_h
ccy = np.clip(ccy, crop_h / 2, H - crop_h / 2)
print(f'Camera: crop {crop_w.min():.0f}-{crop_w.max():.0f} x {crop_h.min():.0f}-{crop_h.max():.0f} px '
f'from {W}x{H} (zoom {max_h / crop_h.max():.2f}-{max_h / crop_h.min():.2f}x)')
return np.stack([ccx, ccy, crop_w, crop_h], 1), True
def crop_frame(frame, cam, out_size, interp=cv2.INTER_CUBIC):
"""Sub-pixel crop + resize (smooth camera motion without 1-px jitter). Returns (view, M) where M maps
source pixels to output pixels.
The crop is taken at source resolution with getRectSubPix (bilinear, sub-pixel centre, replicated
border) and then resized once -- several times cheaper than warping straight to the output size."""
cx, cy, cw, ch = cam
W, H = out_size
pw, ph = max(2, int(round(cw))), max(2, int(round(ch)))
patch = cv2.getRectSubPix(frame, (pw, ph), (float(cx), float(cy)))
sx, sy = W / pw, H / ph
view = cv2.resize(patch, (W, H), interpolation=cv2.INTER_AREA if sx < 1 else interp)
M = np.float32([[sx, 0, W / 2 - cx * sx], [0, sy, H / 2 - cy * sy]])
return view, M
# ════════════════════════════ graphics ════════════════════════════
def load_font(size, bold=True):
cands = [os.path.join(ROOT, 'graphics', 'art', 'SalmaproMedium.otf')]
win = os.path.join(os.environ.get('WINDIR', 'C:/Windows'), 'Fonts')
cands += [os.path.join(win, 'segoeuib.ttf' if bold else 'segoeui.ttf'),
os.path.join(win, 'arialbd.ttf' if bold else 'arial.ttf'), 'DejaVuSans-Bold.ttf']
for c in cands:
try:
return ImageFont.truetype(c, int(size))
except OSError:
continue
return ImageFont.load_default()
def load_numeric_font(size):
# Bold sans with even digit widths so the counter doesn't wobble.
win = os.path.join(os.environ.get('WINDIR', 'C:/Windows'), 'Fonts')
for c in (os.path.join(win, 'seguisb.ttf'), os.path.join(win, 'segoeuib.ttf'),
os.path.join(win, 'arialbd.ttf'), 'DejaVuSans-Bold.ttf'):
try:
return ImageFont.truetype(c, int(size))
except OSError:
continue
return load_font(size)
C_GOLD = (240, 192, 64)
GAUGE_STOPS = [(0.0, C_ACCENT), (0.45, C_ROPE), (0.75, C_FEET), (1.0, C_BAD)]
def lerp_colors(stops, t):
pos = [s[0] for s in stops]
return np.stack([np.interp(t, pos, [s[1][c] for s in stops]) for c in range(3)], -1)
def over(dst, src):
"""Porter-Duff 'over' for float RGBA arrays in [0, 1] (straight alpha)."""
sa, da = src[..., 3:4], dst[..., 3:4]
oa = sa + da * (1 - sa)
rgb = (src[..., :3] * sa + dst[..., :3] * da * (1 - sa)) / np.maximum(oa, 1e-6)
return np.concatenate([rgb, oa], -1)
class SpeedGauge:
"""Procedural 270-degree arc gauge: gradient value arc with glow, end-cap and peak marker.
Rendered analytically in numpy so it is anti-aliased at 1x without supersampling."""
START, SWEEP = 135.0, 270.0 # degrees, clockwise from 3 o'clock (image coords)
def __init__(self, size, max_jps, u):
self.S = S = int(size)
self.max = max_jps
yy, xx = np.mgrid[0:S, 0:S].astype(np.float32) + 0.5
dx, dy = xx - S / 2, yy - S / 2
r = np.hypot(dx, dy)
P = ((np.degrees(np.arctan2(dy, dx)) - self.START) % 360) / self.SWEEP
gap_mid = (1 + 360 / self.SWEEP) / 2
self.P = np.where(P > gap_mid, P - 360 / self.SWEEP, P) # progress along arc, <0 / >1 in the gap
self.R = R = S * 0.40
self.w = w = S * 0.085
self.px = np.radians(self.SWEEP) * R # pixels per unit progress
ring = np.clip(w / 2 - np.abs(r - R) + 0.5, 0, 1)
in_sweep = np.clip(np.minimum(self.P, 1 - self.P) * self.px + 0.5, 0, 1)
self.ring = (ring * in_sweep)[..., None]
self.rgb8 = lerp_colors(GAUGE_STOPS, np.clip(self.P, 0, 1)).astype(np.uint8)
glow = cv2.GaussianBlur(self.ring[..., 0], (0, 0), w * 0.55)
glow = glow / glow.max() * 0.55
ring2 = self.ring[..., 0]
self.val_a = (ring2 + glow * (1 - ring2)).astype(np.float32) # arc + glow, combined alpha
self.P = self.P.astype(np.float32)
# anti-aliased white end-cap patch
cr = w * 0.30
k = int(np.ceil(cr + 2)) * 2 + 1
yk, xk = np.mgrid[0:k, 0:k] + 0.5 - k / 2
cap = np.zeros((k, k, 4), np.uint8)
cap[..., :3] = 255
cap[..., 3] = np.uint8(np.clip(cr - np.hypot(xk, yk) + 0.5, 0, 1) * 255)
self.cap = Image.fromarray(cap)
# static base: dark face + faint track + ticks + numerals (drawn once)
face = np.clip((R - w / 2 - 2) - r + 0.5, 0, 1)[..., None] * 0.55
base = over(np.concatenate([np.broadcast_to(np.array(C_BG) / 255, (S, S, 3)), face], -1),
np.concatenate([np.ones((S, S, 3)), self.ring * 0.13], -1))
ss = 2
ticks = Image.new('RGBA', (S * ss, S * ss), (0, 0, 0, 0))
d = ImageDraw.Draw(ticks)
f_num = load_font(22 * u * ss)
c = S * ss / 2
for k in range(int(max_jps * 2) + 1):
v = k / 2
a = np.radians(self.START + self.SWEEP * v / max_jps)
major = k % 2 == 0
r0 = (R + w / 2 + 5 * u) * ss
r1 = r0 + (14 if major else 7) * u * ss
d.line([(c + r0 * np.cos(a), c + r0 * np.sin(a)), (c + r1 * np.cos(a), c + r1 * np.sin(a))],
fill=(255, 255, 255, 190 if major else 90), width=int(max(1, (3 if major else 2) * u * ss)))
if major and int(v) % 2 == 0:
rt = (R - w / 2 - 22 * u) * ss
d.text((c + rt * np.cos(a), c + rt * np.sin(a)), f'{int(v)}', font=f_num,
fill=C_MUTED + (220,), anchor='mm')
ticks = np.asarray(ticks.resize((S, S), Image.LANCZOS)).astype(np.float32) / 255
self.base = Image.fromarray(np.uint8(over(base, ticks) * 255))
def _pos(self, v):
a = np.radians(self.START + self.SWEEP * np.clip(v / self.max, 0, 1))
return self.S / 2 + self.R * np.cos(a), self.S / 2 + self.R * np.sin(a), a
def render(self, v, peak=None):
t = float(np.clip(v / self.max, 0, 1))
out = np.empty((self.S, self.S, 4), np.uint8)
out[..., :3] = self.rgb8
if t > 0:
m = np.clip((t - self.P) * self.px + 0.5, 0, 1)
out[..., 3] = (m * self.val_a * 255).astype(np.uint8)
else:
out[..., 3] = 0
img = Image.fromarray(out)
if t > 0: # bright end cap
hx, hy, _ = self._pos(v)
k = self.cap.size[0]
img.alpha_composite(self.cap, (int(round(hx - k / 2)), int(round(hy - k / 2))))
if peak is not None and peak > 0: # peak marker: short gold bar just outside the ring
_, _, a = self._pos(peak)
r0, r1 = self.R + self.w / 2 + 2, self.R + self.w * 1.4
c = self.S / 2
ImageDraw.Draw(img).line([(c + r0 * np.cos(a), c + r0 * np.sin(a)), (c + r1 * np.cos(a), c + r1 * np.sin(a))],
fill=C_GOLD + (255,), width=max(2, int(self.w * 0.22)))
return img
def vertical_gradient(w, h, color, a0, a1, power=1.6):
t = np.linspace(0, 1, h)[:, None] ** power
alpha = (a0 + (a1 - a0) * t) * 255
arr = np.zeros((h, w, 4), np.uint8)
arr[..., :3] = color
arr[..., 3] = np.repeat(alpha, w, 1).astype(np.uint8)
return Image.fromarray(arr)
class Layout:
"""All positions in output pixels, derived from the output format. Boxes are (x0, y0, x1, y1)."""
def __init__(self, fmt, W, H):
self.fmt, self.W, self.H = fmt, W, H
u = self.u = (W if fmt == 'vertical' else H) / 1080
m = self.margin = int(56 * u)
s = lambda v: int(v * u)
if fmt == 'vertical':
self.top_grad_h = int(H * 0.16)
self.bot_grad_y = int(H * 0.50)
self.logo_w = int(W * 0.36)
self.logo_xy = ('center', s(60))
self.title_y = s(60) + int(self.logo_w * 183 / 1058) + s(24)
self.phase_box = (W - m - s(170), s(230), W - m, s(400))
self.gauge_box = (W - m - s(360), H - s(800), W - m, H - s(440))
self.count_xy = (m, H - s(540))
self.count_size = s(240)
self.trace_box = (m, H - s(410), m + s(600), H - s(240))
self.splits_box = (m + s(630), H - s(410), W - m, H - s(240))
self.beat_box = (m, H - s(205), W - m, H - s(135))
self.chip_y = H - s(80)
else:
self.top_grad_h = int(H * 0.20)
self.bot_grad_y = int(H * 0.55)
self.logo_w = s(380)
self.logo_xy = (m, s(44))
self.title_y = None
self.phase_box = (W - m - s(170), s(40), W - m, s(210))
# left column: speed trace + count; right column: splits + gauge; centre kept clear
self.trace_box = (m, H - s(560), m + s(560), H - s(400))
self.count_xy = (m, H - s(165))
self.count_size = s(210)
self.splits_box = (W - m - s(400), H - s(610), W - m, H - s(450))
self.gauge_box = (W - m - s(390), H - s(430), W - m - s(10), H - s(50))
self.beat_box = (W // 2 - s(330), H - s(100), W // 2 + s(330), H - s(40))
self.chip_y = H - s(58)
def split_seconds(duration):
for s in (5, 10, 15, 20, 30, 60):
if duration / s <= 8:
return s
return int(math.ceil(duration / 8))
class Overlay:
def __init__(self, fmt, out_size, an, fps, title, dial_max, show_pace):
W, H = out_size
self.W, self.H, self.fps, self.an, self.show_pace = W, H, fps, an, show_pace
L = self.L = Layout(fmt, W, H)
u = L.u
self.f_count = load_numeric_font(L.count_size)
self.f_label = load_font(38 * u)
self.f_small = load_font(28 * u)
self.f_tiny = load_font(22 * u)
self.f_gauge = load_numeric_font(92 * u)
self.f_title = load_font(42 * u)
js, je = an['jump_start'], an['jump_end']
self.js, self.je = js, max(je, js + 1)
n = len(an['count'])
# ── derived stats ──
# detected jump times (model mark peaks) -> fractional output-frame positions
bt = an.get('beat_times')
if bt is None or len(bt) < 2:
bt = np.where(np.diff(an['count'], prepend=0) > 0)[0] / fps
self.jumps = np.asarray(bt) * fps # model time == source time; output frame m sits at m / fps
iv = np.diff(bt)
med = _running_median(iv, 9)
self.jump_dev = np.r_[0, np.abs(iv - med) / np.maximum(med, 1e-6)] # timing deviation per jump
self.peak = np.maximum.accumulate(an['jps'])
dur = (self.je - js) / fps
self.split_s = split_seconds(dur)
nsp = max(1, int(math.ceil(dur / self.split_s - 1e-6)))
self.split_b = [min(js + int(round(k * self.split_s * fps)), self.je + 1) for k in range(nsp + 1)]
c_at = lambda f: int(an['count'][min(f, n - 1)]) if f >= 0 else 0
self.c_at = c_at
self.split_final = [c_at(self.split_b[k + 1] - 1) - c_at(self.split_b[k] - 1) for k in range(nsp)]
self.split_max = max(1, max(self.split_final)) * 1.15
act = an['jps'][js:self.je + 1]
self.avg_jps = float(act[act > 0].mean()) if (act > 0).any() else 0.0
# ── static layer: gradients, logo, title, cards, gauge base ──
static = Image.new('RGBA', (W, H), (0, 0, 0, 0))
static.alpha_composite(vertical_gradient(W, L.top_grad_h, C_BG, 0.75, 0.0, power=1.0), (0, 0))
static.alpha_composite(vertical_gradient(W, H - L.bot_grad_y, C_BG, 0.0, 0.92), (0, L.bot_grad_y))
d = ImageDraw.Draw(static)
logo_path = os.path.join(ROOT, 'nextjump_logo.png')
logo_h = 0
if os.path.isfile(logo_path):
logo = Image.open(logo_path).convert('RGBA')
logo_h = int(L.logo_w * logo.height / logo.width)
logo = logo.resize((L.logo_w, logo_h), Image.LANCZOS)
lx = (W - L.logo_w) // 2 if L.logo_xy[0] == 'center' else L.logo_xy[0]
static.alpha_composite(logo, (lx, L.logo_xy[1]))
else:
d.text((W // 2 if fmt == 'vertical' else L.margin, L.logo_xy[1]), 'NEXTJUMP',
font=self.f_title, fill=C_ACCENT + (255,), anchor='mt' if fmt == 'vertical' else 'lt')
if title:
if fmt == 'vertical':
d.text((W // 2, L.title_y), title.upper(), font=self.f_title, fill=C_TEXT + (235,),
anchor='mt', stroke_width=int(2 * u), stroke_fill=C_BG + (160,))
else:
tx = L.logo_xy[0] + L.logo_w + int(36 * u)
ty = L.logo_xy[1] + logo_h // 2
d.line([(tx, ty - int(26 * u)), (tx, ty + int(26 * u))], fill=C_ACCENT + (255,), width=max(2, int(4 * u)))
d.text((tx + int(30 * u), ty), title.upper(), font=self.f_title, fill=C_TEXT + (235,),
anchor='lm', stroke_width=int(2 * u), stroke_fill=C_BG + (160,))
pad = int(16 * u)
for box, label in ((L.trace_box, 'SPEED (JUMPS / SEC)'), (L.splits_box, f'{self.split_s} S SPLITS'),
(L.beat_box, None)):
x0, y0, x1, y1 = box
d.rounded_rectangle([x0 - pad, y0 - pad, x1 + pad, y1 + pad], radius=int(18 * u), fill=C_BG + (150,))
if label:
d.text((x0, y0), label, font=self.f_tiny, fill=C_MUTED + (255,), anchor='lt')
gx0, gy0, gx1, gy1 = L.gauge_box
self.gauge = SpeedGauge(min(gx1 - gx0, gy1 - gy0), dial_max, u)
static.alpha_composite(self.gauge.base, (gx0, gy0))
self.static = static
alpha_rows = np.flatnonzero(np.asarray(static)[:, :, 3].max(1) > 0)
bands, start = [], None
for r0, r1 in zip(np.r_[alpha_rows[:1], alpha_rows[1:][np.diff(alpha_rows) > 1]],
np.r_[alpha_rows[:-1][np.diff(alpha_rows) > 1], alpha_rows[-1:]]):
bands.append((int(r0), static.crop((0, int(r0), W, int(r1) + 1))))
self.static_bands = bands
# ── speed trace axis (fixed: the full jump window is known offline) ──
seg = an['jps'][js:self.je + 1]
self.ymax = max(1.0, float(np.percentile(seg, 99)) * 1.2 if len(seg) else 1.0, an['baseline'] * 1.2)
# frames since the count last went up (drives the count "pop"), precomputed so render() is stateless
ar = np.arange(n)
inc = np.where(np.diff(an['count'], prepend=an['count'][0]) > 0, ar, -10 ** 6)
self.pop_age = ar - np.maximum.accumulate(inc)
# beep-timed event (None when beep detection is off / failed)
self.event = an.get('event')
self.times = an.get('times')
if self.times is None:
self.times = np.arange(n) / fps
self.f_banner = load_font(88 * u)
# misses and (rulebook) skipped jumps as output-frame positions
fr = lambda t: int(np.searchsorted(self.times, t))
self.misses = [(fr(a), fr(b)) for a, b in an.get('miss_times', [])]
self.skips = [fr(t) for t in an.get('skip_times', [])]
self.skip_rule = an.get('skip_after_miss', False)
self.skipped_beats = {int(np.argmin(np.abs(self.jumps - f))) for f in self.skips} if len(self.jumps) else set()
# ── helpers ──
def _chip(self, d, xy, text, font, fg, bg, anchor='lm', pad=(18, 8)):
u = self.L.u
bbox = d.textbbox(xy, text, font=font, anchor=anchor)
px, py = int(pad[0] * u), int(pad[1] * u)
d.rounded_rectangle([bbox[0] - px, bbox[1] - py, bbox[2] + px, bbox[3] + py],
radius=int(14 * u), fill=bg)
d.text(xy, text, font=font, fill=fg, anchor=anchor)
return bbox[2] + px
def _trace(self, idx):
"""Anti-aliased JPS trace over the whole run (2x supersampled), start pace dashed."""
x0, y0, x1, y1 = self.L.trace_box
w, h = x1 - x0, y1 - y0
ss = 2
img = Image.new('RGBA', (w * ss, h * ss), (0, 0, 0, 0))
d = ImageDraw.Draw(img)
u = self.L.u * ss
an = self.an
top = int(38 * u)
H_ = h * ss
X = lambda k: (k - self.js) / (self.je - self.js) * (w * ss - 1)
Y = lambda v: H_ - 2 - (v / self.ymax) * (H_ - top - 4)
d.line([(0, H_ - 2), (w * ss, H_ - 2)], fill=(255, 255, 255, 70), width=max(1, int(2 * u)))
if self.event is not None:
lw = max(2, int(3 * u))
d.line([(lw / 2, top - 6 * u), (lw / 2, H_ - 2)], fill=C_GOOD + (230,), width=lw)
d.line([(w * ss - lw / 2, top - 6 * u), (w * ss - lw / 2, H_ - 2)], fill=C_BAD + (230,), width=lw)
for a, b in self.misses: # misses so far, shaded red
if a <= idx:
d.rectangle([X(a), top, max(X(a) + 2 * u, X(min(b, idx))), H_ - 2], fill=C_BAD + (70,))
if self.show_pace and an['baseline'] > 0:
yb = Y(an['baseline'])
for xs in range(0, w * ss, int(28 * u)):
d.line([(xs, yb), (xs + int(14 * u), yb)], fill=C_ACCENT + (170,), width=max(1, int(2 * u)))
end = min(idx, self.je)
if end > self.js:
ks = np.unique(np.linspace(self.js, end, min(end - self.js + 1, w)).astype(int))
pts = [(X(k), Y(v)) for k, v in zip(ks, an['jps'][ks])]
d.polygon([(pts[0][0], H_ - 2)] + pts + [(pts[-1][0], H_ - 2)], fill=C_FEET + (55,))
d.line(pts, fill=C_FEET + (255,), width=max(2, int(4 * u)))
kp = self.js + int(np.argmax(an['jps'][self.js:end + 1])) # peak so far
px, py = X(kp), Y(an['jps'][kp])
r = int(6 * u)
d.ellipse([px - r, py - r, px + r, py + r], fill=C_GOLD + (255,))
r = int(9 * u)
px, py = pts[-1]
d.ellipse([px - r, py - r, px + r, py + r], fill=(255, 255, 255, 255), outline=C_FEET + (255,),
width=max(1, int(3 * u)))
return img.reduce(ss), (x0, y0)
def _splits(self, d, idx):
x0, y0, x1, y1 = self.L.splits_box
u = self.L.u
top = y0 + int(34 * u)
base = y1 - int(4 * u)
n = len(self.split_final)
gap = int(10 * u)
bw = (x1 - x0 - gap * (n - 1)) / n
cur = int(np.searchsorted(self.split_b, idx, side='right')) - 1
done = [k for k in range(n) if k < cur]
best = max(done, key=lambda k: self.split_final[k]) if len(done) > 1 else None
for k in range(n):
bx0 = x0 + k * (bw + gap)
d.rounded_rectangle([bx0, top, bx0 + bw, base], radius=int(6 * u), fill=(255, 255, 255, 18))
if k > cur or idx < self.js:
continue
val = self.split_final[k] if k < cur else self.c_at(idx) - self.c_at(self.split_b[k] - 1)
if val <= 0:
continue
by = base - (base - top - int(26 * u)) * val / self.split_max
col = C_FEET if k == cur else (C_GOLD if k == best else C_ACCENT)
d.rounded_rectangle([bx0, by, bx0 + bw, base], radius=int(6 * u), fill=col + (235,))
d.text((bx0 + bw / 2, by - int(4 * u)), f'{val}', font=self.f_tiny, fill=(255, 255, 255, 255), anchor='mb')
def _beats(self, d, idx):
"""Scrolling strip of the last few seconds of jumps, coloured by timing consistency."""
x0, y0, x1, y1 = self.L.beat_box
u = self.L.u
label_w = int(200 * u)
hist = self.jumps[self.jumps <= idx]
# rhythm score: 100 * (1 - coefficient of variation) of the last 12 intervals
iv = np.diff(hist[-13:]) / self.fps
score = f'{max(0, 100 * (1 - iv.std() / iv.mean())):.0f}%' if len(iv) >= 4 else '--'
cy = (y0 + y1) / 2
d.text((x0, cy - int(2 * u)), 'RHYTHM', font=self.f_tiny, fill=C_MUTED + (255,), anchor='ls')
d.text((x0, cy + int(4 * u)), score, font=self.f_label, fill=(255, 255, 255, 255), anchor='lt')
tx0 = x0 + label_w
win = 3.0 * self.fps
d.line([(tx0, cy), (x1, cy)], fill=(255, 255, 255, 40), width=max(1, int(2 * u)))
th = (y1 - y0) * 0.36
if self.event is not None:
for f, col, lab in ((self.js, C_GOOD, 'START'), (self.je + 1, C_BAD, 'END')):
if idx - win < f <= idx:
x = x1 - (idx - f) / win * (x1 - tx0)
d.line([(x, y0 - 4 * u), (x, y1 + 4 * u)], fill=col + (235,), width=max(2, int(3 * u)))
d.text((x + 6 * u, y0 - 2 * u), lab, font=self.f_tiny, fill=col + (255,), anchor='lt')
for a, b in self.misses:
if a <= idx and b > idx - win:
xa = x1 - (idx - max(a, idx - win)) / win * (x1 - tx0)
xb = x1 - (idx - min(b, idx)) / win * (x1 - tx0)
d.rectangle([xa, y0, max(xa + 2 * u, xb), y1], fill=C_BAD + (60,))
for j in np.where((self.jumps <= idx) & (self.jumps > idx - win))[0]:
age = (idx - self.jumps[j]) / win
x = x1 - age * (x1 - tx0)
dev = self.jump_dev[j]
col = C_ROPE if dev < 0.08 else (C_FEET if dev < 0.2 else C_BAD)
fresh = (idx - self.jumps[j]) < 0.08 * self.fps
hh = th * (1.35 if fresh else 1.0)
a = int(255 * (0.25 + 0.75 * (1 - age)))
if j in self.skipped_beats: # rulebook: first jump after a miss is not counted
d.rounded_rectangle([x - 4 * u, cy - hh, x + 4 * u, cy + hh], radius=int(3 * u),
outline=(255, 255, 255, a), width=max(1, int(2 * u)))
continue
d.rounded_rectangle([x - 3 * u, cy - hh, x + 3 * u, cy + hh], radius=int(3 * u),
fill=((255, 255, 255) if fresh else col) + (a,))
def _phase(self, idx):
"""Rope-phase orbit: where the rope is in its rotation (bottom = passing the feet)."""
x0, y0, x1, y1 = self.L.phase_box
S = min(x1 - x0, y1 - y0)
ss = 2
img = Image.new('RGBA', (S * ss, S * ss), (0, 0, 0, 0))
d = ImageDraw.Draw(img)
u = self.L.u * ss
c = S * ss / 2
R = c * 0.72
d.ellipse([c - c * 0.98, c - c * 0.98, c + c * 0.98, c + c * 0.98], fill=C_BG + (150,))
d.ellipse([c - R, c - R, c + R, c + R], outline=(255, 255, 255, 60), width=max(1, int(3 * u)))
d.line([(c, c + R - 10 * u), (c, c + R + 10 * u)], fill=C_FEET + (220,), width=max(1, int(3 * u)))
active = self.an['active'][idx] and self.js <= idx <= self.je
if active:
n_tr = max(2, int(0.12 * self.fps))
ph = self.an['phase'][max(0, idx - n_tr):idx + 1]
for k, a in enumerate(ph):
t = (k + 1) / len(ph)
r = (3 + 6 * t) * u
px, py = c + R * np.cos(a), c - R * np.sin(a)
d.ellipse([px - r, py - r, px + r, py + r], fill=C_ROPE + (int(40 + 150 * t),))
a = ph[-1]
px, py = c + R * np.cos(a), c - R * np.sin(a)
d.line([(c, c), (px, py)], fill=C_ROPE + (110,), width=max(1, int(3 * u)))
r = 11 * u
d.ellipse([px - r, py - r, px + r, py + r], fill=(255, 255, 255, 255), outline=C_ROPE + (255,),
width=max(1, int(3 * u)))
d.text((c, c), 'ROPE', font=self.f_tiny.font_variant(size=int(self.f_tiny.size * ss)),
fill=C_MUTED + (255,), anchor='mm')
return img.reduce(ss), (x0 + (x1 - x0 - S) // 2, y0)
def render(self, frame, idx):
an, L, u = self.an, self.L, self.L.u
# RGB base + 'RGBA' draw mode => translucent shapes blend instead of overwriting
img = Image.fromarray(frame)
for y, band in self.static_bands:
img.paste(band, (0, y), band)
d = ImageDraw.Draw(img, 'RGBA')
cnt = int(an['count'][idx])
done = idx > self.je if self.event is not None else idx > self.je + int(0.3 * self.fps)
jumping = self.js <= idx <= self.je
pop = int(max(0, 6 - self.pop_age[idx]))
# ── big count ──
cx, cy = L.count_xy
f = self.f_count if pop == 0 else self.f_count.font_variant(size=int(L.count_size * (1 + 0.06 * pop / 6)))
d.text((cx, cy), f'{an["final"] if done else cnt}', font=f, fill=(255, 255, 255, 255), anchor='ls',
stroke_width=int(3 * u), stroke_fill=C_BG + (120,))
sub_y = cy + int(50 * u)
if done:
label = 'FINAL SCORE' + (f' · {len(self.misses)} MISS{"ES" if len(self.misses) != 1 else ""}'
if self.misses else '')
d.text((cx + int(6 * u), sub_y), label, font=self.f_label, fill=C_ACCENT + (255,), anchor='ls')
else:
x = d.textbbox((cx + int(6 * u), sub_y), 'JUMPS', font=self.f_label, anchor='ls')[2]
d.text((cx + int(6 * u), sub_y), 'JUMPS', font=self.f_label, fill=C_TEXT + (255,), anchor='ls',
stroke_width=int(2 * u), stroke_fill=C_BG + (140,))
if jumping and idx > an.get('first_jump', self.js) + int(3 * self.fps):
# projection: current count + current speed for the rest of the run
proj = int(round(cnt + an['jps'][idx] * (self.je - idx) / self.fps))
d.text((x + int(22 * u), sub_y), f'ON PACE FOR {proj}', font=self.f_small,
fill=C_TEXT + (220,), anchor='ls')
# ── gauge ──
gx0, gy0, gx1, gy1 = L.gauge_box
v = self.avg_jps if done else float(an['jps'][idx])
peak = float(self.peak[min(idx, self.je)])
g = self.gauge.render(v, peak if idx >= self.js else None)
img.paste(g, (gx0, gy0), g)
gc = (gx0 + self.gauge.S / 2, gy0 + self.gauge.S / 2)
d.text((gc[0], gc[1] + int(8 * u)), f'{v:.1f}', font=self.f_gauge, fill=(255, 255, 255, 255), anchor='ms')
d.text((gc[0], gc[1] + int(46 * u)), 'AVG JUMPS / SEC' if done else 'JUMPS / SEC',
font=self.f_tiny, fill=C_MUTED + (255,), anchor='ms')
if idx >= self.js and peak > 0:
d.text((gc[0], gy0 + self.gauge.S * 0.90), f'PEAK {peak:.1f}', font=self.f_small,
fill=C_GOLD + (255,), anchor='ms')
# ── panels ──
tr, pos = self._trace(idx)
img.paste(tr, pos, tr)
self._splits(d, idx)
self._beats(d, idx)
ph, pos = self._phase(idx)
img.paste(ph, pos, ph)
# ── chips: timer | pace delta | status ──
x = L.margin
if self.event is not None:
t = float(np.clip(self.times[idx] - self.event[0], 0, self.event[1] - self.event[0]))
else:
t = max(0.0, (min(idx, self.je) - self.js) / self.fps) if idx >= self.js else 0.0
x = self._chip(d, (x, L.chip_y), f'{int(t // 60)}:{t % 60:04.1f}', self.f_label,
(255, 255, 255, 255), (255, 255, 255, 40)) + int(18 * u)
if self.show_pace and idx > an.get('first_jump', self.js) + int(3 * self.fps):
k = min(idx, self.je)
delta = int(round(an['count'][k] - an['pace'][k]))
col = C_GOOD if delta >= 0 else C_BAD
x = self._chip(d, (x, L.chip_y), f'{delta:+d} VS START PACE',
self.f_small, (255, 255, 255, 255), col + (200,)) + int(18 * u)
in_miss = any(a <= idx < b for a, b in self.misses)
after_skip = any(f <= idx < f + int(1.0 * self.fps) for f in self.skips)
waiting_skip = self.skip_rule and any(b <= idx < f for (a, b), f in zip(self.misses, self.skips))
if jumping and (in_miss or (idx >= an.get('first_jump', self.js) and not an['active'][idx])):
self._chip(d, (x, L.chip_y), 'MISS', self.f_small, (255, 255, 255, 255), C_BAD + (230,))
elif jumping and (after_skip or waiting_skip):
self._chip(d, (x, L.chip_y), 'FIRST JUMP AFTER MISS NOT COUNTED', self.f_small,
(255, 255, 255, 255), C_BAD + (170,))
elif self.event is not None and idx < self.js:
self._chip(d, (x, L.chip_y), 'WAITING FOR BEEP', self.f_small, (255, 255, 255, 255), (255, 255, 255, 40))
if self.event is not None:
self._event_bar(d, idx)
self._beep_flash(d, idx)
return np.asarray(img)
def _event_bar(self, d, idx):
"""Full-width event clock along the bottom edge: fills from the start beep to the end beep."""
u, W, H = self.L.u, self.W, self.H
h = max(4, int(8 * u))
y0 = H - h
dur = self.event[1] - self.event[0]
t = float(np.clip((self.times[idx] - self.event[0]) / dur, 0, 1))
d.rectangle([0, y0, W, H], fill=(255, 255, 255, 45))
if t > 0:
d.rectangle([0, y0, int(W * t), H], fill=(C_BAD if t >= 1 else C_ACCENT) + (235,))
for k in range(1, len(self.split_b) - 1): # split boundaries
xs = int(W * (self.split_b[k] - self.js) / max(1, self.je + 1 - self.js))
d.line([(xs, y0), (xs, H)], fill=C_BG + (200,), width=max(1, int(2 * u)))
def _beep_flash(self, d, idx):
"""Border flash + banner for ~1.2 s after each beep (green START, red TIME)."""
u, W, H = self.L.u, self.W, self.H
for t_ev, col, text in ((self.event[0], C_GOOD, 'START'), (self.event[1], C_BAD, 'TIME')):
age = self.times[idx] - t_ev
if not 0 <= age < 1.2:
continue
k = 1 - age / 1.2
bw = int(16 * u)
d.rectangle([bw // 2, bw // 2, W - bw // 2, H - bw // 2], outline=col + (int(230 * k),), width=bw)
scale = 1 + 0.25 * max(0.0, 1 - age / 0.15) # quick pop-in
f = self.f_banner.font_variant(size=int(self.f_banner.size * scale))
cx, cy = W / 2, H * 0.30
bb = d.textbbox((cx, cy), text, font=f, anchor='mm')
px, py = int(40 * u), int(16 * u)
a = int(255 * min(1.0, 3 * k))
d.rounded_rectangle([bb[0] - px, bb[1] - py, bb[2] + px, bb[3] + py], radius=int(22 * u),
fill=col + (int(0.9 * a),))
d.text((cx, cy), text, font=f, fill=(255, 255, 255, a), anchor='mm')
# ════════════════════════════ parallel render ════════════════════════════
#
# Each worker process renders + encodes a contiguous segment of the output on its own (decode -> crop ->
# overlay -> its own ffmpeg), then the segments are joined losslessly (concat demuxer, -c copy) and the
# audio is muxed in. Frames are fully independent (the overlay is stateless), so this scales with cores.
_RENDER = {}
def _nvenc_args(ffmpeg, cq):
"""h264_nvenc arguments if this ffmpeg + driver can actually use NVENC, else None (cached)."""
key = ('nvenc', ffmpeg)
if key not in _RENDER:
_RENDER[key] = None
for preset in (['-preset', 'p5', '-tune', 'hq'], ['-preset', 'slow']):
args = ['-c:v', 'h264_nvenc'] + preset + ['-rc', 'vbr', '-cq', str(cq), '-b:v', '0', '-pix_fmt', 'yuv420p']
r = subprocess.run([ffmpeg, '-v', 'error', '-f', 'lavfi', '-i', 'color=black:s=256x256:d=0.1', '-frames:v', '2']
+ args + ['-f', 'null', '-'], capture_output=True)
if r.returncode == 0:
_RENDER[key] = args
break
return _RENDER[key]
def encoder_args(ffmpeg, encoder, preset, crf, threads):
if encoder in ('auto', 'nvenc'):
nv = _nvenc_args(ffmpeg, crf + 2)
if nv:
return nv, 'h264_nvenc'
if encoder == 'nvenc':
print('Note: NVENC not usable with this ffmpeg / driver - using libx264')
return (['-c:v', 'libx264', '-preset', preset, '-crf', str(crf), '-pix_fmt', 'yuv420p', '-profile:v', 'high',
'-threads', str(threads)], f'libx264 {preset}')
def _render_worker_boot(counter, cv_threads=1):
"""Pool initializer: runs when the (pre-spawned) worker starts, long before the render config exists."""
cv2.setNumThreads(cv_threads) # many workers x OpenCV's own thread pool would oversubscribe the CPU
if cv_threads: # ... and so would 8 workers x a 16-thread H.264 decoder each
os.environ.setdefault('OPENCV_FFMPEG_CAPTURE_OPTIONS', 'threads;2')
_RENDER['counter'] = counter
def _render_worker_init(cfg_path, counter=None, cv_threads=1):
import pickle
if counter is not None:
_render_worker_boot(counter, cv_threads)
if _RENDER.get('cfg_path') == cfg_path:
return
with open(cfg_path, 'rb') as f:
cfg = pickle.load(f)
cfg['ov'] = Overlay(*cfg['overlay'])
cfg['cfg_path'] = cfg_path
_RENDER.update(cfg)
def _render_segment(seg, cfg_path=None):
if cfg_path:
_render_worker_init(cfg_path)
m0, m1, seg_path = seg
R = _RENDER
t_start = time.perf_counter()
W, H = R['size']
cmd = [R['ffmpeg'], '-y', '-loglevel', 'error', '-f', 'rawvideo', '-pix_fmt', 'yuv420p', '-s', f'{W}x{H}',
'-r', f'{R["fps"]:.6f}', '-i', '-'] + R['enc'] + ['-an', seg_path]
proc = subprocess.Popen(cmd, stdin=subprocess.PIPE)
cap = cv2.VideoCapture(R['video'])
i, frame, done = -1, None, 0
prof = [0.0, 0.0, 0.0]
first = int(R['out_src'][m0])
if first > 0:
# Seek straight to the segment (OpenCV decodes forward from the previous keyframe, frame-exact);
# otherwise every worker would decode the whole video up to its segment.
cap.set(cv2.CAP_PROP_POS_FRAMES, first)
if int(round(cap.get(cv2.CAP_PROP_POS_FRAMES))) == first:
i = first - 1
else: # container without reliable seeking: start over
cap.release()
cap = cv2.VideoCapture(R['video'])
try:
for m in range(m0, m1):
src = int(R['out_src'][m])
while i < src:
if not cap.grab():
raise RuntimeError(f'video ended early at frame {i}')
i, frame = i + 1, None
if frame is None:
ok, bgr = cap.retrieve()
frame = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
if R['mirror']:
frame = np.ascontiguousarray(frame[:, ::-1])
t1 = time.perf_counter()
view, M = crop_frame(frame, R['cam'][m], (W, H), R['interp'])
if R['kps'] is not None:
draw_skeleton(view, R['kps'][m], M)
t2 = time.perf_counter()
yuv = cv2.cvtColor(R['ov'].render(view, m), cv2.COLOR_RGB2YUV_I420)
t3 = time.perf_counter()
proc.stdin.write(yuv.tobytes())
t4 = time.perf_counter()
prof[0] += t2 - t1; prof[1] += t3 - t2; prof[2] += t4 - t3
done += 1
if done % 10 == 0:
with R['counter'].get_lock():
R['counter'].value += 10
finally:
proc.stdin.close()
proc.wait()
cap.release()
if proc.returncode != 0:
raise RuntimeError(f'ffmpeg (segment) exited with code {proc.returncode}')
with R['counter'].get_lock():
R['counter'].value += done % 10
if os.environ.get('NEXTJUMP_PROFILE'):
print(f'[seg {m0}] crop {prof[0]:.1f}s overlay {prof[1]:.1f}s write {prof[2]:.1f}s', flush=True)
return seg_path, time.perf_counter() - t_start, m1 - m0
class RenderPool:
"""Spawn the render worker processes early (they only import this module and wait), so their start-up
cost overlaps the GPU analysis instead of delaying the render."""
def __init__(self, workers):
import threading
import multiprocessing as mp
self.workers = workers
self.ctx = mp.get_context('spawn')
self.counter = self.ctx.Value('i', 0)
self._ex = None
self._starter = threading.Thread(target=self._start, daemon=True) if workers > 1 else None
if self._starter:
self._starter.start()
def _start(self):
from concurrent.futures import ProcessPoolExecutor
while self.workers > 1:
ex = ProcessPoolExecutor(self.workers, mp_context=self.ctx, initializer=_render_worker_boot,
initargs=(self.counter,))
try:
for fu in [ex.submit(time.sleep, 0) for _ in range(self.workers)]: # make every process start
fu.result()
self._ex = ex
return
except Exception as e: # e.g. a worker could not start because memory is short
ex.shutdown(wait=False, cancel_futures=True)
print(f'\nNote: render pool with {self.workers} workers failed ({type(e).__name__}); retrying with fewer')
self.workers //= 2
@property
def ex(self):
if self._starter:
self._starter.join()
return self._ex
def close(self):
if self.ex:
self.ex.shutdown(wait=True, cancel_futures=True)
def render_video(cfg, n, workers, output, audio_src, progress=None, pool=None):
import pickle
tmp = tempfile.mkdtemp(prefix='nextjump_render_')
own_pool = pool is None
pool = pool or RenderPool(max(1, min(workers, n // 30)))
try:
cfg_path = os.path.join(tmp, 'cfg.pkl')
with open(cfg_path, 'wb') as f:
pickle.dump(cfg, f, protocol=pickle.HIGHEST_PROTOCOL)
k = max(1, min(pool.workers, n // 30))
bounds = np.linspace(0, n, k + 1).astype(int)
segs = [(int(bounds[j]), int(bounds[j + 1]), os.path.join(tmp, f'seg{j:02d}.mp4')) for j in range(k)]
counter = pool.counter
counter.value = 0
print(f'Rendering {n} frames in {k} parallel segment(s)')
def tick():
print(f'\rRendering {min(n, counter.value)}/{n} ', end='', flush=True)
if progress:
progress(min(1.0, counter.value / n))
if k == 1 or pool.ex is None:
segs = [(0, n, os.path.join(tmp, 'seg00.mp4'))]
_render_worker_init(cfg_path, counter, cv_threads=0)
stats = [_render_segment(segs[0])]
else:
futs = [pool.ex.submit(_render_segment, sg, cfg_path) for sg in segs]
while not all(fu.done() for fu in futs):
time.sleep(0.25)
tick()
stats = [fu.result() for fu in futs] # re-raises worker errors
tick()
print(f'\nSegments: ' + ', '.join(f'{nf} fr / {dt:.1f}s' for _, dt, nf in stats))
lst = os.path.join(tmp, 'segments.txt')
with open(lst, 'w') as f:
f.writelines(f"file '{sg[2].replace(os.sep, '/')}'\n" for sg in segs)
cmd = [cfg['ffmpeg'], '-y', '-loglevel', 'error', '-f', 'concat', '-safe', '0', '-i', lst]
if audio_src:
cmd += ['-i', audio_src, '-map', '0:v', '-map', '1:a?', '-c:a', 'aac', '-b:a', '192k', '-shortest']
cmd += ['-c:v', 'copy', '-movflags', '+faststart', output]
subprocess.run(cmd, check=True)
finally:
if own_pool:
pool.close()
shutil.rmtree(tmp, ignore_errors=True)
def draw_skeleton(frame, kp, M, conf=0.3):
pts = np.c_[kp[:, :2], np.ones(17)] @ M.T
t = max(2, int(frame.shape[1] / 400))
for a, b in SKELETON:
if kp[a, 2] > conf and kp[b, 2] > conf:
cv2.line(frame, tuple(np.int32(pts[a])), tuple(np.int32(pts[b])), C_ACCENT, t, cv2.LINE_AA)
for i in range(17):
if kp[i, 2] > conf:
cv2.circle(frame, tuple(np.int32(pts[i])), t + 2, (255, 255, 255), -1, cv2.LINE_AA)
# ════════════════════════════ main ════════════════════════════
class StageTimer:
"""Wall-clock per pipeline stage; printed at the end of every run."""
def __init__(self):
self.t0 = self.t = time.perf_counter()
self.stages = []
def mark(self, name):
now = time.perf_counter()
self.stages.append((name, now - self.t))
self.t = now
def summary(self):
total = time.perf_counter() - self.t0
return 'Timing: ' + ', '.join(f'{k} {v:.1f}s' for k, v in self.stages if v >= 0.05) + f' | total {total:.1f}s'
def main(argv=None, progress=None):
"""CLI entry point; also callable from Python (e.g. the Gradio Space):
main(['--video', 'in.mp4', '--beep'], progress=lambda frac, desc: ...)
Returns a summary dict (output path, final count, beep times, ...)."""
def report(lo, hi, desc):
return (lambda f: progress(lo + (hi - lo) * f, desc)) if progress else None
p = argparse.ArgumentParser(description='Broadcast-ready NextJump visual from the ONNX export')
p.add_argument('--video', required=True)
p.add_argument('--output', default=None, help='Output mp4 (default: <video>_<format>.mp4)')
p.add_argument('--format', choices=list(FORMATS), default='vertical',
help='vertical = 1080x1920 (9:16), horizontal = 1920x1080 (16:9)')
p.add_argument('--resolution_scale', type=float, default=1.0, help='e.g. 2.0 for 4K / 2160x3840')
p.add_argument('--onnx', default=os.path.join(ROOT, 'checkpoint', 'onnx', 'nextjump_all_128.onnx'))
p.add_argument('--max_frames', type=int, default=100000)
p.add_argument('--turbo', action='store_true',
help='Fastest settings for instant replay: 30 fps analysis + output, lighter person detection '
'and probing, no pose model (frames on detector boxes), fast encoder. Any option given '
'explicitly still wins.')
p.add_argument('--stride', type=int, default=None, help='Sliding-window stride (frames; default 32, turbo 64)')
p.add_argument('--batch_size', type=int, default=1, help='Windows per ONNX call')
p.add_argument('--model_fps', type=float, default=None,
help='Frame rate fed to the counting model (default 60 = training rate; turbo 30). '
'Higher-rate sources are decimated to this')
p.add_argument('--out_fps', type=float, default=None,
help='Max output frame rate (default 60, turbo 30; sources above it are decimated)')
p.add_argument('--mirror', action='store_true')
p.add_argument('--cpu', action='store_true')
p.add_argument('--title', default='', help='Event / athlete title shown in the header')
# framing
p.add_argument('--fill', type=float, default=None,
help='Fraction of frame height the jumper occupies (default 0.50 vertical, 0.60 horizontal)')
p.add_argument('--anchor_y', type=float, default=None,
help='Vertical position of body centre in the frame, 0=top (default 0.40 vertical, 0.47 horizontal)')
p.add_argument('--max_zoom', type=float, default=2.5, help='Max zoom relative to the largest possible crop')
p.add_argument('--smooth', type=float, default=0.8, help='Camera smoothing (seconds)')
p.add_argument('--framing', choices=['smart', 'pose', 'off'], default='smart',
help='smart = detect people and frame the most active jumper(s); pose = follow whoever '
'MoveNet locks on to; off = static centred crop')
p.add_argument('--no_crop', action='store_true', help='Same as --framing off')
p.add_argument('--jumpers', type=int, default=1,
help='smart framing: number of people to frame (e.g. 2 for pairs, 3 for double dutch); '
'the camera follows the centre of mass of the group')
p.add_argument('--detector', default='coco',
help="smart framing: 'coco' (torchvision Faster R-CNN MobileNetV3) or a path to a "
'torch-pickled detector such as fasterrcnn_model_finetuned.pth')
p.add_argument('--det_fps', type=float, default=None,
help='smart framing: person detections per second (default 5, turbo 2.5)')
p.add_argument('--count_on_crop', action='store_true',
help='smart framing: feed the counting model a crop around the selected jumper(s) '
'instead of the full frame (helps when others in frame are moving)')
p.add_argument('--skip_after_miss', action='store_true',
help='Rulebook judging: the first jump after each miss is not counted')
# graphics
p.add_argument('--skeleton', action='store_true', help='Draw MoveNet skeleton on the jumper')
p.add_argument('--no_pace', action='store_true', help='Hide start-pace comparison')
p.add_argument('--dial_max', type=float, default=0,
help='Top of the speed gauge in jumps/sec (0 = auto: 10, or more for faster events)')
p.add_argument('--periodicity_threshold', type=float, default=0.2)
p.add_argument('--mark_threshold', type=float, default=0.15)
# event timing
p.add_argument('--beep', action='store_true',
help='Detect the official start/end beeps and count only between them')
p.add_argument('--event_length', type=float, default=30.0, help='Event length in seconds (beep spacing)')
p.add_argument('--beep_file', default=os.path.join(ROOT, 'beep.WAV'), help='Reference beep recording')
p.add_argument('--beep_threshold', type=float, default=0.8, help='Initial correlation peak threshold')
p.add_argument('--crf', type=int, default=None, help='Quality (lower = better; default 17, turbo 19)')
p.add_argument('--x264_preset', default=None, help='x264 speed/size trade-off (default veryfast)')
p.add_argument('--encoder', choices=['auto', 'x264', 'nvenc'], default='auto',
help='auto = NVIDIA NVENC when this ffmpeg + driver support it, else libx264')
p.add_argument('--render_workers', type=int, default=0,
help='Parallel render/encode processes (0 = auto: half the CPU threads, max 8)')
p.add_argument('--no_audio', action='store_true')
args = p.parse_args(argv)
turbo = args.turbo
for name, normal, fast in (('model_fps', MODEL_FPS, 30.0), ('out_fps', 60.0, 30.0), ('stride', 32, 64),
('det_fps', 5.0, 2.5), ('crf', 17, 19), ('x264_preset', 'veryfast', 'veryfast')):
if getattr(args, name) is None:
setattr(args, name, fast if turbo else normal)
W, H = FORMATS[args.format]
W, H = int(round(W * args.resolution_scale / 2) * 2), int(round(H * args.resolution_scale / 2) * 2)
fill = args.fill or (0.50 if args.format == 'vertical' else 0.60)
anchor_y = args.anchor_y if args.anchor_y is not None else (0.40 if args.format == 'vertical' else 0.47)
output = args.output or f'{os.path.splitext(os.path.basename(args.video))[0]}_{args.format}.mp4'
ffmpeg = find_ffmpeg()
T = StageTimer()
workers = args.render_workers or max(1, min(8, cpu_count() // 2))
pool = RenderPool(workers)
try:
return _run(args, T, pool, workers, W, H, fill, anchor_y, output, ffmpeg, report)
finally:
pool.close()
def _run(args, T, pool, workers, W, H, fill, anchor_y, output, ffmpeg, report):
turbo = args.turbo
import torch
device = torch.device('cuda' if torch.cuda.is_available() and not args.cpu else 'cpu')
print(f'Device: {device} | output {W}x{H} ({args.format}) -> {output}')
framing = 'off' if args.no_crop else args.framing
cap = cv2.VideoCapture(args.video)
src_fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
n_src = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
cap.release()
duration = min(n_src if n_src > 0 else args.max_frames, args.max_frames) / src_fps
model_fps, out_fps = pick_rate(src_fps, args.model_fps), pick_rate(src_fps, args.out_fps)
if abs(model_fps - args.model_fps) > 0.1 * args.model_fps:
print(f'Warning: source is {src_fps:.2f} fps; the counting model expects ~{args.model_fps:.0f} fps, '
f'so counts may be less reliable')
sess = make_session(args.onnx, device.type == 'cuda')
from concurrent.futures import ThreadPoolExecutor
gpu = ThreadPoolExecutor(1) # single ONNX queue: streaming counter + jumper probe
sink = StreamingCounter(sess, args.stride, gpu)
T.mark('setup')
# ── 0. event beeps (audio only; also tells the jumper probe where the event is) ──
event = None
if args.beep:
event = detect_beeps(args.video, ffmpeg, args.event_length, args.beep_file, args.beep_threshold)
if event is None:
print('Warning: beeps not found - counting over the whole video instead')
T.mark('beeps')
# ── 0b. smart framing: find people, probe each one's activity, keep the N most active ──
tracks, sel = [], []
if framing == 'smart':
detector = load_person_detector(args.detector, device)
seq_len = sess.get_inputs()[0].shape[1]
windows = probe_windows(duration, seq_len / model_fps, event, n=3 if turbo else 5)
scan = scan_people(args.video, args.max_frames, args.mirror, src_fps, model_fps, args.det_fps, detector,
device, windows, det_side=640 if turbo else 960,
model_sink=None if args.count_on_crop else sink,
progress=report(0.0, 0.25, 'Finding the jumpers'))
dets, probes, pscale = scan['dets'], scan['probes'], scan['scale']
del detector
T.mark('people scan')
if device.type == 'cuda':
torch.cuda.empty_cache()
tracks = build_tracks(dets)
n_all = len(tracks)
tracks = prefilter_tracks(tracks, duration, keep=max(4 if turbo else 6, 2 * args.jumpers + (1 if turbo else 2)))
print(f'{n_all} people tracked; probing the {len(tracks)} most likely jumpers')
if tracks:
scores = score_tracks(sess, tracks, probes, pscale, 1,
progress=report(0.25, 0.32, 'Picking the most active jumper'), gpu=gpu)
sel = select_jumpers(tracks, scores, max(1, args.jumpers), frame_w=dets_frame_w(dets, args.video),
duration=duration)
del probes
T.mark('jumper probe')
if not sel:
print('Warning: no people detected - falling back to pose framing')
framing = 'pose'
base = 0.32 if sel else 0.0
pose_region = model_crop = None
if sel and len(sel) == 1:
tr0 = tracks[sel[0]]
def pose_region(t):
x0, y0, x1, y1 = track_boxes(tr0, t)[0]
return ((x0 + x1) / 2, (y0 + y1) / 2, 1.25 * max(x1 - x0, y1 - y0))
if sel and args.count_on_crop:
gt = np.arange(0, duration + 0.1, 0.1)
G = group_box(tracks, sel, gt)
gcx = gaussian_filter1d((G[:, 0] + G[:, 2]) / 2, 10, mode='nearest')
gcy = gaussian_filter1d((G[:, 1] + G[:, 3]) / 2, 10, mode='nearest')
gside = gaussian_filter1d(median_filter(1.5 * np.maximum(G[:, 2] - G[:, 0], G[:, 3] - G[:, 1]), 21,
mode='nearest'), 10, mode='nearest')
def model_crop(t):
return np.interp(t, gt, gcx), np.interp(t, gt, gcy), np.interp(t, gt, gside)
# ── 1. analysis pass (skipped entirely when the people scan already produced the model input and no
# pose is needed: turbo single-jumper framing runs on the detector boxes) ──
model_frames = scan['model_frames'] if framing == 'smart' or sel else None # full-frame input is selection-independent
need_pose = framing == 'pose' or (len(sel) == 1 and (args.skeleton or not turbo))
pose_fps = out_fps if args.skeleton else min(out_fps, 10.0 if turbo else 15.0) # camera is smoothed anyway
if model_frames is None or need_pose:
mf, kps, frame_size, out_src = analyse_video(
args.video, args.max_frames, args.mirror, 256, device, src_fps, model_fps, out_fps, run_pose=need_pose,
pose_fps=pose_fps, collect_model=model_frames is None, model_sink=sink if model_frames is None else None,
progress=report(base, base + (0.4 - base * 0.6), 'Tracking the jumper'),
pose_region=pose_region, model_crop=model_crop)
model_frames = mf if mf is not None else model_frames
else:
kps, frame_size = None, scan['frame_size']
out_src = schedule(scan['n_frames'], src_fps, out_fps)
n = len(out_src)
out_times = out_src / src_fps
T.mark('analysis pass')
# ── 2. counting model (on the model timeline), then resample onto output frames ──
mid = base + (0.4 - base * 0.6)
if isinstance(model_frames, StreamingCounter):
res = model_frames.finish(progress=report(mid, 0.68, 'Counting jumps'))
else:
res = run_model(sess, model_frames, args.stride, args.batch_size, progress=report(mid, 0.68, 'Counting jumps'))
gpu.shutdown()
del model_frames, sess, sink
an = compute_analytics(res, model_fps, args.periodicity_threshold, args.mark_threshold, event,
skip_after_miss=args.skip_after_miss)
an = to_output_timeline(an, model_fps, out_times)
T.mark('counting model')
# ── 3. camera path ──
boxes = None
if framing == 'pose' and kps is not None:
boxes = boxes_from_keypoints(kps)
elif sel and len(sel) == 1:
# keypoints of the selected jumper; fall back to their detector box where the pose is unreliable
boxes = boxes_from_keypoints(kps) if kps is not None else np.full((n, 4), np.nan)
db = track_boxes(tracks[sel[0]], out_times)
kc = (boxes[:, :2] + boxes[:, 2:]) / 2
off = ~((kc[:, 0] >= db[:, 0]) & (kc[:, 0] <= db[:, 2]) & (kc[:, 1] >= db[:, 1]) & (kc[:, 1] <= db[:, 3]))
boxes[off | ~np.isfinite(boxes).all(1)] = db[off | ~np.isfinite(boxes).all(1)]
elif sel:
boxes = group_box(tracks, sel, out_times)
if sel and len(sel) > 1 and boxes is not None:
need_w = 1.35 * np.nanmedian(boxes[:, 2] - boxes[:, 0])
max_w = min(frame_size[0], frame_size[1] * W / H)
if need_w > max_w * 1.05:
print(f'Note: the {len(sel)} selected people span ~{need_w:.0f} px but a {args.format} crop of this '
f'video is at most {max_w:.0f} px wide; framing on their centre of mass '
f'(try --format horizontal to fit the whole group)')
cam, tracked = compute_camera_path(boxes, frame_size, (W, H), out_fps, fill, anchor_y,
args.max_zoom, args.smooth)
tracked = tracked and kps is not None
if len(cam) == 1:
cam = np.repeat(cam, n, 0)
# ── 4. render ──
dial_max = args.dial_max or float(max(10, 2 * math.ceil(np.max(an['jps'], initial=0) * 1.1 / 2)))
overlay_args = (args.format, (W, H), an, out_fps, args.title, dial_max, not args.no_pace)
overlay = Overlay(*overlay_args) # main-process copy: summary stats only
enc, enc_name = encoder_args(ffmpeg, args.encoder, args.x264_preset, args.crf,
threads=max(1, cpu_count() // max(1, min(workers, n // 30 or 1))))
print(f'Encoder: {enc_name}')
cfg = dict(video=args.video, mirror=args.mirror, out_src=out_src, cam=cam,
kps=kps if (args.skeleton and tracked) else None, size=(W, H), fps=out_fps, enc=enc, ffmpeg=ffmpeg,
overlay=overlay_args, interp=cv2.INTER_LINEAR if turbo else cv2.INTER_CUBIC)
T.mark('setup render')
render_video(cfg, n, workers, output, None if args.no_audio else args.video,
progress=report(0.68, 1.0, 'Rendering video'), pool=pool)
T.mark('render')
print(f'\nSaved -> {output} ({W}x{H}, {out_fps:.2f} fps, {n} frames, final count {an["final"]})')
print(T.summary())
return dict(output=output, final=an['final'], event=event, beep_requested=args.beep, framing=framing,
jumpers=len(sel), misses=len(an['miss_times']), skipped=len(an['skip_times']),
duration=(an['jump_end'] - an['jump_start']) / out_fps, avg_jps=float(overlay.avg_jps),
peak_jps=float(np.max(an['jps'])) if len(an['jps']) else 0.0, splits=overlay.split_final,
split_seconds=overlay.split_s, size=(W, H), fps=out_fps, frames=n)
if __name__ == '__main__':
main()
"""
claude_anim_v2.py — broadcast-ready NextJump overlay video from the ONNX export.
Pipeline
1. Pass 1 (analysis): decode the video once, building
- 256x256 model frames for the ONNX counting model, sampled at 60 fps
(the training rate - feeding e.g. 120 fps footage makes it over-count), and
- MoveNet-Thunder keypoints (with crop-region tracking) for framing.
2. Sliding-window ONNX inference -> period length, periodicity, marks, phase.
3. Keypoints -> a smoothed "virtual camera" crop path that keeps the jumper
(and the rope arc) well framed at the target aspect ratio.
4. Pass 2 (render): decode again, crop/zoom each frame along the camera path,
composite broadcast graphics and pipe straight into ffmpeg (audio copied).
Example
python claude_anim_v2.py --video ali.mp4 --format vertical
python claude_anim_v2.py --video 30s_wr_male_new_b_60fps.mp4 \
--format horizontal --title "30s SPEED | WORLD RECORD"
"""
import os
import sys
import math
import time
import shutil
import tempfile
import argparse
import subprocess
import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont
# scipy is imported lazily: the render worker processes never need it (each import costs memory)
def gaussian_filter1d(*args, **kwargs):
from scipy.ndimage import gaussian_filter1d as f
return f(*args, **kwargs)
def median_filter(*args, **kwargs):
from scipy.ndimage import median_filter as f
return f(*args, **kwargs)
def cpu_count():
"""CPUs this process may actually use (container / affinity limits, e.g. a 2-vCPU Space on a big host)."""
try:
return max(1, len(os.sched_getaffinity(0)))
except AttributeError:
return max(1, os.cpu_count() or 1)
def _running_median(x, size):
"""1-D median filter with edge replication (numpy only; used by the overlay in render workers)."""
if len(x) == 0:
return x
h = size // 2
xp = np.pad(np.asarray(x, float), h, mode='edge')
return np.median(np.lib.stride_tricks.sliding_window_view(xp, size), axis=1)
ROOT = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.join(ROOT, 'movenet-pytorch'))
# Output order of checkpoint/onnx/nextjump_all_128.onnx. The graph's output
# *names* are shifted by one (masks were prepended to RepNet.forward's return
# tuple), so always index by position, never by name.
OUT_MASKS, OUT_PERIOD, OUT_PERIODICITY, OUT_MARKS, OUT_EVENT, OUT_FOOT, OUT_PHASE = range(7)
# COCO-17 keypoint indices (MoveNet order)
NOSE, L_SHOULDER, R_SHOULDER, L_HIP, R_HIP, L_ANKLE, R_ANKLE = 0, 5, 6, 11, 12, 15, 16
SKELETON = [(0, 1), (0, 2), (1, 3), (2, 4), (5, 6), (5, 7), (7, 9), (6, 8), (8, 10),
(5, 11), (6, 12), (11, 12), (11, 13), (13, 15), (12, 14), (14, 16)]
# ── Palette (matches claude_anim.py) ──
C_BG = (13, 17, 23)
C_ACCENT = (88, 166, 255)
C_FEET = (255, 157, 0)
C_ROPE = (78, 205, 196)
C_GOOD = (63, 185, 80)
C_BAD = (248, 81, 73)
C_TEXT = (230, 237, 243)
C_MUTED = (139, 148, 158)
MODEL_FPS = 60.0 # the counting model was trained on 60 fps clips (see eval.py data dirs)
FORMATS = {'vertical': (1080, 1920), 'horizontal': (1920, 1080)}
def sigmoid(x):
return 1 / (1 + np.exp(-x))
def hex_rgb(h):
h = h.lstrip('#')
return tuple(int(h[i:i + 2], 16) for i in (0, 2, 4))
def find_ffmpeg():
# NEXTJUMP_FFMPEG lets you point at a newer build (e.g. one whose NVENC works with your driver)
override = os.environ.get('NEXTJUMP_FFMPEG')
if override and os.path.isfile(override):