diff --git a/training/README.md b/training/README.md index 1603c6a..e08b74c 100644 --- a/training/README.md +++ b/training/README.md @@ -2,18 +2,82 @@ 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. Includes -KPI evaluation (precision, recall, mAP-equivalent for the heatmap formulation) against -held-out labeled data. +efficient architectures sized around the overall 66ms/tick inference budget. See CLAUDE.md §3.3–3.4 for full requirements. -**Status:** not yet implemented. +**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 . +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. diff --git a/training/pyproject.toml b/training/pyproject.toml index f2ac441..38286fd 100644 --- a/training/pyproject.toml +++ b/training/pyproject.toml @@ -3,7 +3,17 @@ name = "spelunkai-training" version = "0.0.0" description = "SpelunkAI detector training pipeline (anchor-free, per-set CNNs)" requires-python = ">=3.10" -dependencies = [] +dependencies = [ + "torch>=2.0", + "numpy>=1.24", + "pillow>=10.0", +] + +[project.optional-dependencies] +dev = ["pytest>=8"] + +[project.scripts] +spelunkai-train = "spelunkai_training.train:main" [build-system] requires = ["setuptools>=68"] diff --git a/training/src/spelunkai_training/client.py b/training/src/spelunkai_training/client.py new file mode 100644 index 0000000..8576bdc --- /dev/null +++ b/training/src/spelunkai_training/client.py @@ -0,0 +1,85 @@ +"""Fetches a promoted dataset version (and its set's class list) from the +labeling backend. Uses the stdlib HTTP client since this is a couple of +simple GET requests - not worth adding a dependency for. +""" +from __future__ import annotations + +import json +import urllib.request +from dataclasses import dataclass +from typing import Dict, List + + +@dataclass +class LabelData: + sub_class_id: int + x: float + y: float + width: float + height: float + + +@dataclass +class FrameData: + id: int + session_name: str + frame_index: int + image_path: str + width: int + height: int + labels: List[LabelData] + + +@dataclass +class DatasetVersionData: + id: int + set_id: int + name: str + frames: List[FrameData] + + +def parse_dataset_version(payload: Dict) -> DatasetVersionData: + frames = [] + for entry in payload["frames"]: + frame = entry["frame"] + labels = [ + LabelData(sub_class_id=l["sub_class_id"], x=l["x"], y=l["y"], width=l["width"], height=l["height"]) + for l in entry["labels"] + ] + frames.append(FrameData( + id=frame["id"], + session_name=frame["session_name"], + frame_index=frame["frame_index"], + image_path=frame["image_path"], + width=frame["width"], + height=frame["height"], + labels=labels, + )) + return DatasetVersionData(id=payload["id"], set_id=payload["set_id"], name=payload["name"], frames=frames) + + +def build_class_mapping(sets_payload: Dict) -> Dict[int, int]: + """sub_class_id -> stable class index (0..N-1), ordered by sub_class_id so + the mapping doesn't shift between training runs on different dataset + versions of the same set.""" + sub_class_ids = sorted( + sub_class["id"] + for main_class in sets_payload["main_classes"] + for sub_class in main_class["sub_classes"] + ) + return {sub_class_id: index for index, sub_class_id in enumerate(sub_class_ids)} + + +def _get_json(url: str) -> Dict: + with urllib.request.urlopen(url, timeout=30) as resp: + return json.loads(resp.read().decode("utf-8")) + + +def fetch_dataset_version(base_url: str, version_id: int) -> DatasetVersionData: + payload = _get_json(f"{base_url.rstrip('/')}/dataset-versions/{version_id}") + return parse_dataset_version(payload) + + +def fetch_class_mapping(base_url: str, set_id: int) -> Dict[int, int]: + payload = _get_json(f"{base_url.rstrip('/')}/sets/{set_id}") + return build_class_mapping(payload) diff --git a/training/src/spelunkai_training/dataset.py b/training/src/spelunkai_training/dataset.py new file mode 100644 index 0000000..14926fc --- /dev/null +++ b/training/src/spelunkai_training/dataset.py @@ -0,0 +1,64 @@ +"""PyTorch Dataset wrapping a fetched dataset version: loads each frame's +image from local disk - training runs on the same machine/FRAMES_ROOT as the +labeling backend, so images are read directly rather than re-downloaded over +HTTP - and encodes its labels into CenterNet targets via `targets.py`. +""" +from __future__ import annotations + +from pathlib import Path +from typing import Dict + +import numpy as np +import torch +from PIL import Image +from torch.utils.data import Dataset + +from .client import DatasetVersionData +from .model import OUTPUT_STRIDE +from .targets import Box, encode_targets + + +class CenterNetDataset(Dataset): + def __init__( + self, + dataset_version: DatasetVersionData, + images_root: Path, + class_index_by_sub_class_id: Dict[int, int], + output_stride: int = OUTPUT_STRIDE, + ): + self.frames = dataset_version.frames + self.images_root = Path(images_root) + self.class_index_by_sub_class_id = class_index_by_sub_class_id + self.output_stride = output_stride + self.num_classes = len(set(class_index_by_sub_class_id.values())) + + def __len__(self) -> int: + return len(self.frames) + + def __getitem__(self, index: int) -> Dict[str, torch.Tensor]: + frame = self.frames[index] + image = Image.open(self.images_root / frame.image_path).convert("RGB") + image_tensor = torch.from_numpy(np.asarray(image, dtype=np.float32) / 255.0).permute(2, 0, 1) + + boxes = [ + Box( + class_index=self.class_index_by_sub_class_id[label.sub_class_id], + x=label.x, y=label.y, width=label.width, height=label.height, + ) + for label in frame.labels + if label.sub_class_id in self.class_index_by_sub_class_id + ] + targets = encode_targets( + boxes=boxes, + num_classes=self.num_classes, + image_width=frame.width, + image_height=frame.height, + output_stride=self.output_stride, + ) + + return { + "image": image_tensor, + "heatmap": torch.from_numpy(targets.heatmap), + "wh": torch.from_numpy(targets.wh), + "mask": torch.from_numpy(targets.mask), + } diff --git a/training/src/spelunkai_training/losses.py b/training/src/spelunkai_training/losses.py new file mode 100644 index 0000000..a2dc54a --- /dev/null +++ b/training/src/spelunkai_training/losses.py @@ -0,0 +1,51 @@ +"""Loss functions for the CenterNet-style detector: a modified focal loss +for the heatmap (Law & Deng 2018 / Zhou et al. 2019) and a masked L1 loss +for width/height regression, applied only at ground-truth object centers. +""" +from __future__ import annotations + +from typing import Dict, Tuple + +import torch + + +def focal_loss(pred: torch.Tensor, target: torch.Tensor, alpha: float = 2.0, beta: float = 4.0) -> torch.Tensor: + """Pixel-wise modified focal loss, normalized by the number of objects (peaks == 1).""" + eps = 1e-6 + pred = pred.clamp(eps, 1 - eps) + + pos_mask = target.eq(1).float() + neg_mask = target.lt(1).float() + neg_weights = torch.pow(1 - target, beta) + + pos_loss = torch.log(pred) * torch.pow(1 - pred, alpha) * pos_mask + neg_loss = torch.log(1 - pred) * torch.pow(pred, alpha) * neg_weights * neg_mask + + num_pos = pos_mask.sum() + loss = -(pos_loss.sum() + neg_loss.sum()) + if num_pos > 0: + loss = loss / num_pos + return loss + + +def masked_l1_loss(pred: torch.Tensor, target: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: + """L1 loss over `pred`/`target` (both `(B, 2, H, W)`), counted only where `mask` (`(B, H, W)`) is 1.""" + mask = mask.unsqueeze(1) # (B, 1, H, W), broadcasts over the 2 wh channels + num_valid = mask.sum().clamp(min=1.0) + loss = torch.abs(pred - target) * mask + return loss.sum() / num_valid + + +def detection_loss( + pred_heatmap: torch.Tensor, + pred_wh: torch.Tensor, + target_heatmap: torch.Tensor, + target_wh: torch.Tensor, + mask: torch.Tensor, + wh_weight: float = 0.1, +) -> Tuple[torch.Tensor, Dict[str, float]]: + """Combined loss; `wh_weight=0.1` matches the standard CenterNet size-loss weighting.""" + hm_loss = focal_loss(pred_heatmap, target_heatmap) + wh_loss = masked_l1_loss(pred_wh, target_wh, mask) + total = hm_loss + wh_weight * wh_loss + return total, {"heatmap_loss": hm_loss.item(), "wh_loss": wh_loss.item(), "total_loss": total.item()} diff --git a/training/src/spelunkai_training/model.py b/training/src/spelunkai_training/model.py new file mode 100644 index 0000000..11b3055 --- /dev/null +++ b/training/src/spelunkai_training/model.py @@ -0,0 +1,69 @@ +"""A small anchor-free, center-heatmap-based detector (CenterNet-style), per +CLAUDE.md §3.3: predicts a per-class center-point heatmap plus direct +width/height regression from each detected center. + +Deliberately small/efficient - this is a v1 architecture sized with the +66ms/tick budget in mind, given multiple such models may run per tick. +Tune channel widths/depth from here once real timing numbers are measured +on the target GPU (RTX 3070 Ti); the architecture itself (stride, head +shapes) is what the rest of the pipeline depends on, not these specific +widths. +""" +from __future__ import annotations + +from typing import Tuple + +import torch +from torch import nn + +OUTPUT_STRIDE = 4 + + +class ConvBlock(nn.Module): + def __init__(self, in_channels: int, out_channels: int, stride: int = 1): + super().__init__() + self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False) + self.bn = nn.BatchNorm2d(out_channels) + self.act = nn.ReLU(inplace=True) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.act(self.bn(self.conv(x))) + + +class CenterNetDetector(nn.Module): + def __init__(self, num_classes: int, base_channels: int = 32): + super().__init__() + c = base_channels + self.backbone = nn.Sequential( + ConvBlock(3, c // 2, stride=1), + ConvBlock(c // 2, c, stride=2), # stride 2 + ConvBlock(c, c * 2, stride=2), # stride 4 (= OUTPUT_STRIDE) + ConvBlock(c * 2, c * 2, stride=1), + ConvBlock(c * 2, c * 2, stride=1), + ) + + head_channels = c * 2 + self.heatmap_head = nn.Sequential( + ConvBlock(head_channels, head_channels, stride=1), + nn.Conv2d(head_channels, num_classes, kernel_size=1), + ) + self.wh_head = nn.Sequential( + ConvBlock(head_channels, head_channels, stride=1), + nn.Conv2d(head_channels, 2, kernel_size=1), + ) + + self._init_heatmap_bias() + + def _init_heatmap_bias(self) -> None: + # Standard CenterNet trick: bias the heatmap head's last conv so + # predictions start near-zero everywhere (most pixels are + # background) instead of ~0.5, which noticeably stabilizes early + # focal-loss training. + final_conv = self.heatmap_head[-1] + nn.init.constant_(final_conv.bias, -2.19) # sigmoid(-2.19) ~= 0.1 + + def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + features = self.backbone(x) + heatmap = torch.sigmoid(self.heatmap_head(features)) + wh = self.wh_head(features) + return heatmap, wh diff --git a/training/src/spelunkai_training/targets.py b/training/src/spelunkai_training/targets.py new file mode 100644 index 0000000..5d9d2db --- /dev/null +++ b/training/src/spelunkai_training/targets.py @@ -0,0 +1,108 @@ +"""CenterNet-style target encoding: converts pixel-space bounding boxes into +a per-class center-point heatmap plus width/height regression targets, per +CLAUDE.md §3.3 ("a center-point heatmap per class + direct width/height +regression from each detected center"). + +Uses the standard CornerNet/CenterNet gaussian-radius formulation (Law & +Deng 2018; Zhou et al. 2019), so overlapping-box penalties in the heatmap +loss roughly track IoU instead of using an arbitrary fixed-radius blob. +""" +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import List + +import numpy as np + + +@dataclass +class Box: + class_index: int + x: float + y: float + width: float + height: float + + +@dataclass +class Targets: + heatmap: np.ndarray # (num_classes, out_h, out_w), float32, in [0, 1] + wh: np.ndarray # (2, out_h, out_w), float32, pixel units (image space) + mask: np.ndarray # (out_h, out_w), float32, 1 at object centers else 0 + + +def gaussian_radius(height: float, width: float, min_overlap: float = 0.7) -> float: + a1 = 1.0 + b1 = height + width + c1 = width * height * (1 - min_overlap) / (1 + min_overlap) + r1 = (b1 + math.sqrt(b1 ** 2 - 4 * a1 * c1)) / 2 + + a2 = 4.0 + b2 = 2 * (height + width) + c2 = (1 - min_overlap) * width * height + r2 = (b2 + math.sqrt(b2 ** 2 - 4 * a2 * c2)) / 2 + + a3 = 4.0 * min_overlap + b3 = -2 * min_overlap * (height + width) + c3 = (min_overlap - 1) * width * height + r3 = (b3 + math.sqrt(b3 ** 2 - 4 * a3 * c3)) / 2 + + return min(r1, r2, r3) + + +def _gaussian_2d(diameter: int, sigma: float) -> np.ndarray: + radius = (diameter - 1) / 2 + y, x = np.ogrid[-radius:radius + 1, -radius:radius + 1] + gaussian = np.exp(-(x * x + y * y) / (2 * sigma * sigma)) + gaussian[gaussian < np.finfo(gaussian.dtype).eps * gaussian.max()] = 0 + return gaussian + + +def _draw_gaussian(heatmap: np.ndarray, center_x: int, center_y: int, radius: int) -> None: + diameter = 2 * radius + 1 + gaussian = _gaussian_2d(diameter, sigma=diameter / 6) + + out_h, out_w = heatmap.shape + left, right = min(center_x, radius), min(out_w - center_x, radius + 1) + top, bottom = min(center_y, radius), min(out_h - center_y, radius + 1) + if left + right <= 0 or top + bottom <= 0: + return + + masked_heatmap = heatmap[center_y - top:center_y + bottom, center_x - left:center_x + right] + masked_gaussian = gaussian[radius - top:radius + bottom, radius - left:radius + right] + np.maximum(masked_heatmap, masked_gaussian, out=masked_heatmap) + + +def encode_targets( + boxes: List[Box], + num_classes: int, + image_width: int, + image_height: int, + output_stride: int, +) -> Targets: + out_w = image_width // output_stride + out_h = image_height // output_stride + + heatmap = np.zeros((num_classes, out_h, out_w), dtype=np.float32) + wh = np.zeros((2, out_h, out_w), dtype=np.float32) + mask = np.zeros((out_h, out_w), dtype=np.float32) + + for box in boxes: + if not (0 <= box.class_index < num_classes): + raise ValueError(f"class_index {box.class_index} out of range for num_classes={num_classes}") + + center_x_out = (box.x + box.width / 2) / output_stride + center_y_out = (box.y + box.height / 2) / output_stride + cx, cy = int(center_x_out), int(center_y_out) + if not (0 <= cx < out_w and 0 <= cy < out_h): + continue # center falls outside the feature map - shouldn't normally happen + + radius = max(0, int(gaussian_radius(box.height / output_stride, box.width / output_stride))) + _draw_gaussian(heatmap[box.class_index], cx, cy, radius) + + wh[0, cy, cx] = box.width + wh[1, cy, cx] = box.height + mask[cy, cx] = 1.0 + + return Targets(heatmap=heatmap, wh=wh, mask=mask) diff --git a/training/src/spelunkai_training/train.py b/training/src/spelunkai_training/train.py new file mode 100644 index 0000000..d50d307 --- /dev/null +++ b/training/src/spelunkai_training/train.py @@ -0,0 +1,68 @@ +"""Training entry point: fetch a promoted dataset version from the labeling +backend, train the CenterNet-style detector on it, checkpoint after every +epoch. Meant to run on `jai` (GPU) - CPU works too, just much slower. +""" +from __future__ import annotations + +import argparse +from pathlib import Path + +import torch +from torch.utils.data import DataLoader + +from .client import fetch_class_mapping, fetch_dataset_version +from .dataset import CenterNetDataset +from .losses import detection_loss +from .model import CenterNetDetector + + +def main() -> None: + parser = argparse.ArgumentParser(prog="spelunkai-train") + parser.add_argument("--backend-url", default="http://127.0.0.1:8000") + parser.add_argument("--dataset-version-id", type=int, required=True) + parser.add_argument("--set-id", type=int, required=True) + parser.add_argument( + "--images-root", type=Path, required=True, + help="Local FRAMES_ROOT the labeling backend serves images from", + ) + parser.add_argument("--output-dir", type=Path, default=Path("checkpoints")) + parser.add_argument("--epochs", type=int, default=50) + parser.add_argument("--batch-size", type=int, default=8) + parser.add_argument("--lr", type=float, default=1e-3) + parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") + args = parser.parse_args() + + class_index_by_sub_class_id = fetch_class_mapping(args.backend_url, args.set_id) + dataset_version = fetch_dataset_version(args.backend_url, args.dataset_version_id) + dataset = CenterNetDataset(dataset_version, args.images_root, class_index_by_sub_class_id) + loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True, num_workers=2) + + model = CenterNetDetector(num_classes=dataset.num_classes).to(args.device) + optimizer = torch.optim.Adam(model.parameters(), lr=args.lr) + + args.output_dir.mkdir(parents=True, exist_ok=True) + + for epoch in range(args.epochs): + model.train() + epoch_loss = 0.0 + for batch in loader: + images = batch["image"].to(args.device) + target_heatmap = batch["heatmap"].to(args.device) + target_wh = batch["wh"].to(args.device) + mask = batch["mask"].to(args.device) + + pred_heatmap, pred_wh = model(images) + loss, _ = detection_loss(pred_heatmap, pred_wh, target_heatmap, target_wh, mask) + + optimizer.zero_grad() + loss.backward() + optimizer.step() + epoch_loss += loss.item() + + avg_loss = epoch_loss / max(1, len(loader)) + print(f"epoch {epoch + 1}/{args.epochs} - loss {avg_loss:.4f}") + torch.save(model.state_dict(), args.output_dir / "last.pt") + + +if __name__ == "__main__": + main() diff --git a/training/tests/test_client.py b/training/tests/test_client.py new file mode 100644 index 0000000..4df8229 --- /dev/null +++ b/training/tests/test_client.py @@ -0,0 +1,44 @@ +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} diff --git a/training/tests/test_dataset.py b/training/tests/test_dataset.py new file mode 100644 index 0000000..3e1a681 --- /dev/null +++ b/training/tests/test_dataset.py @@ -0,0 +1,69 @@ +from pathlib import Path + +from PIL import Image + +from spelunkai_training.client import DatasetVersionData, FrameData, LabelData +from spelunkai_training.dataset import CenterNetDataset +from spelunkai_training.model import OUTPUT_STRIDE + + +def _write_fake_image(path: Path, width: int, height: int) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + Image.new("RGB", (width, height), color=(10, 20, 30)).save(path) + + +def test_dataset_getitem_shapes(tmp_path): + images_root = tmp_path / "frames" + _write_fake_image(images_root / "run01_frames" / "frame_000000.png", width=64, height=64) + + dataset_version = DatasetVersionData( + id=1, set_id=1, name="v1", + frames=[ + FrameData( + id=1, session_name="run01", frame_index=0, + image_path="run01_frames/frame_000000.png", width=64, height=64, + labels=[LabelData(sub_class_id=7, x=8, y=8, width=16, height=16)], + ), + ], + ) + + dataset = CenterNetDataset( + dataset_version=dataset_version, + images_root=images_root, + class_index_by_sub_class_id={7: 0}, + ) + + assert len(dataset) == 1 + sample = dataset[0] + + out_size = 64 // OUTPUT_STRIDE + assert sample["image"].shape == (3, 64, 64) + assert sample["heatmap"].shape == (1, out_size, out_size) + assert sample["wh"].shape == (2, out_size, out_size) + assert sample["mask"].shape == (out_size, out_size) + assert sample["mask"].sum().item() == 1.0 + + +def test_dataset_ignores_labels_for_unknown_sub_classes(tmp_path): + images_root = tmp_path / "frames" + _write_fake_image(images_root / "run01_frames" / "frame_000000.png", width=32, height=32) + + dataset_version = DatasetVersionData( + id=1, set_id=1, name="v1", + frames=[ + FrameData( + id=1, session_name="run01", frame_index=0, + image_path="run01_frames/frame_000000.png", width=32, height=32, + labels=[LabelData(sub_class_id=999, x=0, y=0, width=4, height=4)], + ), + ], + ) + + dataset = CenterNetDataset( + dataset_version=dataset_version, + images_root=images_root, + class_index_by_sub_class_id={7: 0}, + ) + + sample = dataset[0] + assert sample["mask"].sum().item() == 0.0 diff --git a/training/tests/test_losses.py b/training/tests/test_losses.py new file mode 100644 index 0000000..7c75bf1 --- /dev/null +++ b/training/tests/test_losses.py @@ -0,0 +1,58 @@ +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) diff --git a/training/tests/test_model.py b/training/tests/test_model.py new file mode 100644 index 0000000..333c0df --- /dev/null +++ b/training/tests/test_model.py @@ -0,0 +1,34 @@ +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 diff --git a/training/tests/test_targets.py b/training/tests/test_targets.py new file mode 100644 index 0000000..ab9baba --- /dev/null +++ b/training/tests/test_targets.py @@ -0,0 +1,43 @@ +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)