# 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](https://pytorch.org/get-started/locally/)).

## 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.