Jonas
07e75247c4
Add the anchor-free, center-heatmap detector CLAUDE.md §3.3 specifies: - targets.py: encodes ground-truth boxes into a per-class Gaussian heatmap + wh regression target + center mask, using the standard CornerNet/CenterNet gaussian-radius formulation. - model.py: a small conv backbone (stride 4) with heatmap (sigmoid) and wh regression heads - exactly the two outputs CLAUDE.md specifies, sized with the 66ms/tick budget in mind. - losses.py: modified focal loss (heatmap) + masked L1 (wh), combined with the standard CenterNet wh_weight=0.1. - client.py / dataset.py: pull a promoted dataset version + its set's class list from the labeling backend and turn it into a torch.utils.data.Dataset, reading images from local disk (same machine as FRAMES_ROOT). - train.py: wires it into a basic DataLoader -> train loop -> per-epoch checkpoint. The exact loss weighting/architecture sizing is a reasonable, standard v1 default, not a tuned final answer - CLAUDE.md's own roadmap flags the exact formulation as still open; this is the starting point to iterate from. 16/16 unit tests pass on CPU with synthetic data (target/model/loss correctness, dataset-version parsing). Beyond that, ran a real end-to-end smoke test: live labeling backend -> promoted dataset version -> `spelunkai-train` actually training one real epoch against it and writing a checkpoint. Not verified: multi-epoch convergence on real data, which needs jai's GPU and real labeled frames. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
84 lines
3.6 KiB
Markdown
84 lines
3.6 KiB
Markdown
# 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.
|