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>
35 lines
954 B
Python
35 lines
954 B
Python
import torch
|
|
|
|
from spelunkai_training.model import OUTPUT_STRIDE, CenterNetDetector
|
|
|
|
|
|
def test_forward_pass_output_shapes():
|
|
model = CenterNetDetector(num_classes=3, base_channels=8)
|
|
model.eval()
|
|
|
|
x = torch.rand(2, 3, 64, 64)
|
|
with torch.no_grad():
|
|
heatmap, wh = model(x)
|
|
|
|
expected_size = 64 // OUTPUT_STRIDE
|
|
assert heatmap.shape == (2, 3, expected_size, expected_size)
|
|
assert wh.shape == (2, 2, expected_size, expected_size)
|
|
|
|
|
|
def test_heatmap_output_is_in_unit_range():
|
|
model = CenterNetDetector(num_classes=2, base_channels=8)
|
|
model.eval()
|
|
|
|
x = torch.rand(1, 3, 32, 32)
|
|
with torch.no_grad():
|
|
heatmap, _ = model(x)
|
|
|
|
assert heatmap.min() >= 0.0
|
|
assert heatmap.max() <= 1.0
|
|
|
|
|
|
def test_real_capture_resolution_is_divisible_by_output_stride():
|
|
# Spelunky Classic HD's fixed capture resolution (CLAUDE.md §3.1).
|
|
assert 1280 % OUTPUT_STRIDE == 0
|
|
assert 720 % OUTPUT_STRIDE == 0
|