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>
44 lines
1.6 KiB
Python
44 lines
1.6 KiB
Python
import numpy as np
|
|
|
|
from spelunkai_training.targets import Box, encode_targets, gaussian_radius
|
|
|
|
|
|
def test_gaussian_radius_is_positive_for_a_reasonable_box():
|
|
assert gaussian_radius(height=20, width=15) > 0
|
|
|
|
|
|
def test_encode_targets_places_peak_at_box_center():
|
|
box = Box(class_index=0, x=8, y=16, width=8, height=8) # center = (12, 20)
|
|
targets = encode_targets([box], num_classes=1, image_width=64, image_height=64, output_stride=4)
|
|
|
|
# center in output space: (12/4, 20/4) = (3, 5) -> (x=3, y=5)
|
|
assert targets.heatmap.shape == (1, 16, 16)
|
|
assert targets.heatmap[0, 5, 3] == 1.0
|
|
assert targets.heatmap.max() == 1.0
|
|
|
|
|
|
def test_encode_targets_sets_wh_and_mask_only_at_center():
|
|
box = Box(class_index=0, x=8, y=16, width=8, height=12)
|
|
targets = encode_targets([box], num_classes=1, image_width=64, image_height=64, output_stride=4)
|
|
|
|
assert targets.wh[0, 5, 3] == 8
|
|
assert targets.wh[1, 5, 3] == 12
|
|
assert targets.mask[5, 3] == 1.0
|
|
assert targets.mask.sum() == 1.0
|
|
|
|
|
|
def test_encode_targets_rejects_out_of_range_class_index():
|
|
box = Box(class_index=5, x=0, y=0, width=4, height=4)
|
|
try:
|
|
encode_targets([box], num_classes=2, image_width=32, image_height=32, output_stride=4)
|
|
assert False, "expected ValueError"
|
|
except ValueError:
|
|
pass
|
|
|
|
|
|
def test_encode_targets_with_no_boxes_is_all_zero():
|
|
targets = encode_targets([], num_classes=3, image_width=32, image_height=32, output_stride=4)
|
|
assert targets.heatmap.shape == (3, 8, 8)
|
|
assert np.all(targets.heatmap == 0)
|
|
assert np.all(targets.mask == 0)
|