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>
59 lines
1.8 KiB
Python
59 lines
1.8 KiB
Python
import pytest
|
|
import torch
|
|
|
|
from spelunkai_training.losses import detection_loss, focal_loss, masked_l1_loss
|
|
|
|
|
|
def test_focal_loss_is_zero_for_a_perfect_prediction():
|
|
target = torch.zeros(1, 1, 8, 8)
|
|
target[0, 0, 3, 3] = 1.0
|
|
pred = target.clone()
|
|
|
|
loss = focal_loss(pred, target)
|
|
assert torch.isclose(loss, torch.tensor(0.0), atol=1e-6)
|
|
|
|
|
|
def test_focal_loss_is_positive_for_a_wrong_prediction():
|
|
target = torch.zeros(1, 1, 8, 8)
|
|
target[0, 0, 3, 3] = 1.0
|
|
pred = torch.full_like(target, 0.1)
|
|
|
|
loss = focal_loss(pred, target)
|
|
assert loss.item() > 0
|
|
|
|
|
|
def test_masked_l1_loss_only_counts_masked_pixels():
|
|
pred = torch.zeros(1, 2, 4, 4)
|
|
target = torch.zeros(1, 2, 4, 4)
|
|
mask = torch.zeros(1, 4, 4)
|
|
|
|
# mismatch outside the mask - should not affect the loss
|
|
pred[0, :, 0, 0] = 100.0
|
|
loss = masked_l1_loss(pred, target, mask)
|
|
assert torch.isclose(loss, torch.tensor(0.0))
|
|
|
|
# mismatch inside the mask - should affect the loss. Both wh channels
|
|
# contribute (|5-3|=2 each), summed per masked pixel then divided by the
|
|
# number of masked pixels (1) -> 4, not 2.
|
|
mask[0, 1, 1] = 1.0
|
|
pred[0, :, 1, 1] = 5.0
|
|
target[0, :, 1, 1] = 3.0
|
|
loss = masked_l1_loss(pred, target, mask)
|
|
assert torch.isclose(loss, torch.tensor(4.0))
|
|
|
|
|
|
def test_detection_loss_combines_both_terms():
|
|
heatmap_target = torch.zeros(1, 1, 8, 8)
|
|
heatmap_target[0, 0, 2, 2] = 1.0
|
|
wh_target = torch.zeros(1, 2, 8, 8)
|
|
wh_target[0, :, 2, 2] = 10.0
|
|
mask = torch.zeros(1, 8, 8)
|
|
mask[0, 2, 2] = 1.0
|
|
|
|
pred_heatmap = torch.full_like(heatmap_target, 0.3)
|
|
pred_wh = torch.zeros_like(wh_target)
|
|
|
|
total, parts = detection_loss(pred_heatmap, pred_wh, heatmap_target, wh_target, mask)
|
|
assert total.item() > 0
|
|
assert parts["total_loss"] == pytest.approx(parts["heatmap_loss"] + 0.1 * parts["wh_loss"], rel=1e-4)
|