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>
45 lines
1.3 KiB
Python
45 lines
1.3 KiB
Python
from spelunkai_training.client import build_class_mapping, parse_dataset_version
|
|
|
|
|
|
def test_parse_dataset_version():
|
|
payload = {
|
|
"id": 1,
|
|
"set_id": 2,
|
|
"name": "v1",
|
|
"frames": [
|
|
{
|
|
"frame": {
|
|
"id": 10, "session_name": "run01", "frame_index": 0,
|
|
"image_path": "run01_frames/frame_000000.png", "width": 1280, "height": 720,
|
|
},
|
|
"labels": [
|
|
{"sub_class_id": 5, "x": 1.0, "y": 2.0, "width": 3.0, "height": 4.0},
|
|
],
|
|
},
|
|
],
|
|
}
|
|
|
|
result = parse_dataset_version(payload)
|
|
|
|
assert result.id == 1
|
|
assert result.set_id == 2
|
|
assert result.name == "v1"
|
|
assert len(result.frames) == 1
|
|
frame = result.frames[0]
|
|
assert frame.image_path == "run01_frames/frame_000000.png"
|
|
assert len(frame.labels) == 1
|
|
assert frame.labels[0].sub_class_id == 5
|
|
|
|
|
|
def test_build_class_mapping_is_stable_and_ordered_by_id():
|
|
payload = {
|
|
"main_classes": [
|
|
{"id": 1, "name": "Enemy", "sub_classes": [{"id": 5, "name": "Snake"}, {"id": 2, "name": "Bat"}]},
|
|
{"id": 2, "name": "Hazard", "sub_classes": [{"id": 9, "name": "Spikes"}]},
|
|
],
|
|
}
|
|
|
|
mapping = build_class_mapping(payload)
|
|
|
|
assert mapping == {2: 0, 5: 1, 9: 2}
|