spelunkai/training

Detection Model Training

Per-set training pipeline for the anchor-free, center-heatmap-based (CenterNet-style) CNN detectors — one model per label set (Enemy, Items, Traps, ...). Custom, small/ efficient architectures sized around the overall 66ms/tick inference budget.

See CLAUDE.md §3.3–3.4 for full requirements.

Status: core v1 implemented — target encoding, model, loss, dataset loader (pulls a promoted dataset version from labeling/backend), and a training script. Not yet implemented/verified: actual multi-epoch training on a real dataset (needs jai's GPU and real labeled data — nothing to train on yet), KPI evaluation (precision/recall/mAP-equivalent), and inference/export (ONNX etc. for the runtime loop). The exact loss weighting and architecture sizing are a reasonable v1 default, not a tuned final answer — CLAUDE.md's own roadmap flags the exact formulation as an open item; treat this as the starting point to iterate from once real training runs are possible.

How it works

• targets.py — encodes ground-truth boxes into training targets: a per-class Gaussian center-point heatmap (radius scaled to box size via the standard CornerNet/CenterNet formula, so overlap penalties roughly track IoU) plus a width/height regression target and a mask marking which pixels are actual object centers.
• model.py — CenterNetDetector: a small conv backbone (stride 4 output, matching OUTPUT_STRIDE) with two 1x1-conv heads — a per-class heatmap (sigmoid) and a width/height regression head — exactly the two outputs CLAUDE.md §3.3 specifies, nothing extra (no separate offset head).
• losses.py — modified focal loss for the heatmap + masked L1 loss for width/height (masked to object-center pixels only), combined with the standard CenterNet wh_weight=0.1.
• client.py — pulls a promoted dataset version (GET /dataset-versions/{id}) and its set's class list (GET /sets/{id}) from the labeling backend over plain HTTP (stdlib urllib, no extra dependency for two GET requests).
• dataset.py — CenterNetDataset: a torch.utils.data.Dataset that reads each frame's image straight from local disk (training runs on the same machine/ FRAMES_ROOT as the labeling backend, so no need to re-download images) and encodes its labels via targets.py.
• train.py — wires it all together: fetch → DataLoader → train loop → checkpoint (checkpoints/last.pt) after every epoch.

Setup

python -m venv .venv
source .venv/bin/activate
pip install -e ".[dev]"

On a machine without an NVIDIA GPU (e.g. for running the tests), install the CPU build explicitly to avoid pulling a multi-GB CUDA wheel:

pip install torch --index-url https://download.pytorch.org/whl/cpu
pip install -e ".[dev]"

On jai, install the CUDA build matching its driver/CUDA version instead (see pytorch.org).

Training

spelunkai-train \
  --backend-url http://127.0.0.1:8000 \
  --set-id 1 --dataset-version-id 1 \
  --images-root /path/to/FRAMES_ROOT \
  --output-dir checkpoints

Testing

pytest

Everything here is testable on CPU with synthetic data and doesn't need the labeling backend running: target encoding (peak placement, wh/mask correctness), the model's forward-pass output shapes, loss behavior (zero for a perfect prediction, masked correctly), the dataset-version JSON parsing, and CenterNetDataset.__getitem__ against a tiny generated image. Not verified here: an actual multi-epoch training run converging on real data — that needs jai's GPU and a real promoted dataset.