#!/usr/bin/env python3 """Offline re-tracking of cached detections, built for 1 fps. Why not ByteTrack: supervision computes max_time_lost = int(frame_rate / 30.0 * lost_track_buffer) so the service's `--frame-rate 1 --lost-track-buffer 60` gave max_time_lost = 2 FRAMES. At 1 fps that is a 2 second occlusion tolerance, which is why the A1 baseline produced 707 ids for ~2.5 people (median track life 5 s). Raising lost_track_buffer does almost nothing -- the divide-by-30 eats it. ByteTrack's Kalman filter and IoU-first matching also assume ~30 fps; at 1 fps a walking person moves ~0.7 bbox widths per frame and IoU is simply 0. Two stages instead: 1. Conservative association (short max_age) using predicted position, IoU, centre distance and box-size agreement -> high-purity tracklets. 2. Tracklet stitching across gaps up to `stitch_max_gap_sec`, gated by a physical speed limit and torso-appearance correlation, solved as a one-successor-per-tracklet assignment. Splitting it this way keeps stage 1 from jumping between people during a crossing, while stage 2 still recovers identity across long occlusions. """ from __future__ import annotations import argparse import json from pathlib import Path import numpy as np from scipy.optimize import linear_sum_assignment STAGE1_MAX_AGE = 3 # frames a tracklet may coast before stage 1 closes it IOU_W, DIST_W, SIZE_W = 0.45, 0.40, 0.15 STAGE1_MAX_COST = 0.72 def iou(a, b): ix1, iy1 = max(a[0], b[0]), max(a[1], b[1]) ix2, iy2 = min(a[2], b[2]), min(a[3], b[3]) iw, ih = max(0.0, ix2 - ix1), max(0.0, iy2 - iy1) inter = iw * ih ua = (a[2] - a[0]) * (a[3] - a[1]) + (b[2] - b[0]) * (b[3] - b[1]) - inter return inter / ua if ua > 0 else 0.0 def centre(b): return np.array([(b[0] + b[2]) / 2.0, (b[1] + b[3]) / 2.0]) def main() -> int: ap = argparse.ArgumentParser() ap.add_argument("--detections", required=True) ap.add_argument("--appearance", required=True) ap.add_argument("--zones", required=True) ap.add_argument("--out", required=True) ap.add_argument("--report", default="") args = ap.parse_args() zcfg = json.loads(Path(args.zones).read_text(encoding="utf-8")) tr = zcfg.get("tracker_rules", {}) max_gap = int(tr.get("stitch_max_gap_sec", 12)) max_speed = float(tr.get("stitch_max_speed_px_per_s", 260.0)) min_corr = float(tr.get("stitch_min_appearance_corr", 0.45)) frames = [] with Path(args.detections).open(encoding="utf-8") as fh: for line in fh: frames.append(json.loads(line)) n = len(frames) az = np.load(args.appearance) desc, index = az["desc"], az["index"] desc_of: dict[tuple[int, int], np.ndarray] = { (int(t), int(d)): desc[i] for i, (t, d) in enumerate(index) } # ---- stage 1: conservative association ------------------------------- # live: id -> dict(last_t, box, vel, hist[list of (t,box,det_idx)]) live: dict[int, dict] = {} closed: list[dict] = [] next_id = 1 for fi, fr in enumerate(frames): t = fr["t_local"] dets = fr["detections"] ids = list(live.keys()) if dets and ids: cost = np.ones((len(ids), len(dets)), np.float32) for i, tid in enumerate(ids): tk = live[tid] dt = t - tk["last_t"] pred = np.array(tk["box"], np.float32) c = centre(pred) pc = c + tk["vel"] * dt pred = pred + np.array([pc[0] - c[0], pc[1] - c[1], pc[0] - c[0], pc[1] - c[1]], np.float32) th = tk["box"][3] - tk["box"][1] for j, d in enumerate(dets): b = d["bbox_xyxy"] dist = float(np.linalg.norm(centre(b) - centre(pred))) gate = max_speed * max(dt, 1) if dist > gate: continue dh = b[3] - b[1] size = abs(dh - th) / max(th, dh, 1.0) cost[i, j] = ( IOU_W * (1.0 - iou(pred, b)) + DIST_W * min(dist / gate, 1.0) + SIZE_W * min(size, 1.0) ) ri, ci = linear_sum_assignment(cost) used_d = set() for i, j in zip(ri, ci): if cost[i, j] >= STAGE1_MAX_COST: continue tid = ids[i] tk = live[tid] b = dets[j]["bbox_xyxy"] dt = max(t - tk["last_t"], 1) tk["vel"] = (centre(b) - centre(tk["box"])) / dt tk["box"], tk["last_t"] = b, t tk["hist"].append((t, b, j, dets[j].get("confidence", 0.0))) used_d.add(j) else: used_d = set() for j, d in enumerate(dets): if j in used_d: continue live[next_id] = { "last_t": t, "box": d["bbox_xyxy"], "vel": np.zeros(2, np.float32), "hist": [(t, d["bbox_xyxy"], j, d.get("confidence", 0.0))], } next_id += 1 for tid in [k for k, v in live.items() if t - v["last_t"] > STAGE1_MAX_AGE]: closed.append({"id": tid, **live.pop(tid)}) for tid, v in live.items(): closed.append({"id": tid, **v}) tracklets = sorted(closed, key=lambda tk: tk["hist"][0][0]) print(f"stage1 tracklets={len(tracklets)} (baseline ByteTrack was 707)") # ---- stage 2: stitch across occlusions ------------------------------- def appearance(tk): vs = [desc_of[(t, j)] for t, _b, j, _c in tk["hist"] if (t, j) in desc_of] vs = [v for v in vs if v.sum() > 0] return np.mean(vs, axis=0) if vs else None for tk in tracklets: tk["t0"], tk["t1"] = tk["hist"][0][0], tk["hist"][-1][0] tk["b0"], tk["b1"] = tk["hist"][0][1], tk["hist"][-1][1] tk["app"] = appearance(tk) N = len(tracklets) BIG = 1e6 cost = np.full((N, N), BIG, np.float32) for i, a in enumerate(tracklets): for j, b in enumerate(tracklets): if i == j: continue dt = b["t0"] - a["t1"] if dt <= 0 or dt > max_gap: continue dist = float(np.linalg.norm(centre(b["b0"]) - centre(a["b1"]))) if dist > max_speed * dt: continue ha, hb = a["b1"][3] - a["b1"][1], b["b0"][3] - b["b0"][1] if abs(ha - hb) / max(ha, hb, 1) > 0.45: continue corr = 0.0 if a["app"] is not None and b["app"] is not None: va, vb = a["app"] - a["app"].mean(), b["app"] - b["app"].mean() den = float(np.linalg.norm(va) * np.linalg.norm(vb)) corr = float(va @ vb / den) if den > 0 else 0.0 if corr < min_corr: continue cost[i, j] = 0.45 * (dist / (max_speed * dt)) + 0.40 * (1.0 - corr) + 0.15 * (dt / max_gap) ri, ci = linear_sum_assignment(cost) succ = {int(i): int(j) for i, j in zip(ri, ci) if cost[i, j] < BIG} parent = list(range(N)) def find(x): while parent[x] != x: parent[x] = parent[parent[x]] x = parent[x] return x for i, j in succ.items(): a, b = find(i), find(j) if a != b: parent[b] = a groups: dict[int, list[int]] = {} for i in range(N): groups.setdefault(find(i), []).append(i) print(f"stage2 stitched -> {len(groups)} identities") final_id = {} for k, (_root, members) in enumerate(sorted(groups.items(), key=lambda kv: min(tracklets[m]["t0"] for m in kv[1])), start=1): for m in members: final_id[m] = f"p_{k:03d}" per_frame: dict[int, list[dict]] = {t: [] for t in range(n)} for i, tk in enumerate(tracklets): for t, b, _j, conf in tk["hist"]: per_frame[t].append( {"track_id": final_id[i], "label": "person", "confidence": conf, "bbox_xyxy": b} ) with Path(args.out).open("w", encoding="utf-8") as out: for fr in frames: t = fr["t_local"] out.write( json.dumps( {"t_local": t, "t_abs": fr["t_abs"], "people": per_frame[t]}, ensure_ascii=False ) + "\n" ) lens = [sum(len(tracklets[m]["hist"]) for m in mem) for mem in groups.values()] spans = [ max(tracklets[m]["t1"] for m in mem) - min(tracklets[m]["t0"] for m in mem) + 1 for mem in groups.values() ] rep = { "stage1_tracklets": len(tracklets), "final_identities": len(groups), "baseline_bytetrack_ids": 707, "median_track_len_sec": int(np.median(lens)), "mean_track_len_sec": round(float(np.mean(lens)), 1), "p90_track_span_sec": int(np.percentile(spans, 90)), "max_track_span_sec": int(max(spans)), "ids_lasting_ge_60s": int(sum(1 for s in spans if s >= 60)), } print(json.dumps(rep, indent=2)) if args.report: Path(args.report).write_text(json.dumps(rep, indent=2) + "\n", encoding="utf-8") return 0 if __name__ == "__main__": raise SystemExit(main())