#!/usr/bin/env python3 """Training augmentation/loss tests. Skipped automatically when torch is unavailable. Run with ``python3 test_train_aug.py``. Checks that multi-scale crops keep image/mask alignment, that padding is labelled ignore (never ``other``), and that the Dice term, EMA, and parameter groups behave. """ import random def main(): try: import torch except ImportError: print("torch not installed; skipping training augmentation tests") return import numpy as np from PIL import Image from label_schema import IGNORE from models import build_model from train import ModelEma, dice_loss, letterbox, parameter_groups, scale_crop # A left/right split image whose mask matches it pixel for pixel. img = np.zeros((300, 200, 3), np.uint8) img[:, 100:] = 255 mask = np.zeros((300, 200), np.uint8) mask[:, 100:] = 2 image, mask_im, material = Image.fromarray(img), Image.fromarray(mask), Image.fromarray(np.full((300, 200), 255, np.uint8)) random.seed(0) for _ in range(20): ci, cm, cmat = scale_crop(image, mask_im, material, size=256, crop=128, scale_range=(0.5, 1.5)) assert ci.size == cm.size == cmat.size == (128, 128) a, m = np.asarray(ci)[..., 0], np.asarray(cm) content = m != IGNORE # Bright pixels are wall, dark pixels are other, except at the resize seam. assert ((a[content & (m == 2)] > 100).mean() > 0.95) and ((a[content & (m == 0)] < 150).mean() > 0.95) assert set(np.unique(m)) <= {0, 2, IGNORE} print("ok scale_crop keeps image/mask aligned") # Small images are padded with ignore, not "other". ci, cm, _ = scale_crop(image, mask_im, material, size=64, crop=128, scale_range=(1.0, 1.0)) m = np.asarray(cm) assert (m == IGNORE).any() and m[0, 0] != IGNORE and m[-1, -1] == IGNORE print("ok crop padding is ignore") from train import ignore_fill boxed = np.asarray(letterbox(mask_im, 300, ignore_fill(mask_im))) assert (boxed[:, :40] == IGNORE).all(), "letterbox bars must be ignore" print("ok letterbox bars are ignore") # Dice: perfect logits ~0 loss, wrong logits ~1; inactive classes are excluded. target = torch.tensor([[[0, 2], [2, IGNORE]]]) good = torch.nn.functional.one_hot(target.clamp(0, 9), 10).permute(0, 3, 1, 2).float() * 50 weights = torch.zeros(10) weights[[0, 2]] = 1 assert dice_loss(good, target, weights) < 0.05 assert dice_loss(-good, target, weights) > 0.5 print("ok dice loss") model = build_model("deeplabv3_mobilenet_v3_large", 10) groups = parameter_groups(model, 1e-4, 10, 0.0) assert groups[1]["lr"] == 1e-3 and groups[1]["params"] and groups[0]["params"] head_ids = {id(p) for p in groups[1]["params"]} assert all(id(p) in head_ids for p in model.drywall_head.parameters()) assert all(id(p) in head_ids for p in model.condition_head.parameters()) assert all(id(p) in head_ids for p in model.segmenter.classifier.parameters()) print("ok head parameter group") ema = ModelEma(model, 0.5) with torch.no_grad(): for p in model.parameters(): p.add_(1.0) before = next(ema.module.parameters()).clone() ema.update(model, step=1000) after = next(ema.module.parameters()) assert torch.allclose(after, before + 0.5, atol=1e-5) print("ok EMA update") # End to end on a tiny synthetic set: wallpaper paste keeps masks valid, and # evaluate() scores every class/material with a status and per-source groups. import tempfile from pathlib import Path from torch.utils.data import DataLoader from label_schema import CLASSES, MATERIAL_CLASSES from make_synthetic_rooms import write_dataset from train import WallDataset, evaluate with tempfile.TemporaryDirectory() as tmp: write_dataset(tmp, {"train": 4, "val": 14}, seed=2, size=(96, 128)) random.seed(1) train_ds = WallDataset(tmp, "train", 96, augment=True, crop=64, wallpaper_paste=1.0) for i in range(len(train_ds)): x, y, m, c = train_ds[i] assert x.shape == (3, 64, 64) and y.shape == m.shape == c.shape == (64, 64) assert not ((c != IGNORE) & ~((y >= 1) & (y <= 3))).any(), "condition only on wall" wall = (y == 1) | (y == 2) assert not ((m != IGNORE) & ~((y >= 1) & (y <= 3))).any() assert ((m == 2) & wall).any() or not wall.any(), "pasted wallpaper should be labelled" val_ds = WallDataset(tmp, "val", 96) assert val_ds.coverage_report()["classes_missing"] == [] assert val_ds.coverage_report()["materials_missing"] == [] scores = evaluate(build_model("deeplabv3_mobilenet_v3_large", 10).eval(), DataLoader(val_ds, batch_size=4), torch.device("cpu"), groups=val_ds.groups(Path(tmp) / "metadata.csv")) assert set(scores["per_class"]) == set(CLASSES) and set(scores["material"]) == set(MATERIAL_CLASSES) assert scores["classes_without_ground_truth"] == [] and scores["materials_without_ground_truth"] == [] assert all(scores["per_class"][c]["iou"] is not None for c in CLASSES) assert list(scores["by_group"]) == ["environment:indoor", "source:synthetic"] assert scores["images"] == 14 and scores["ece_semantic"] is not None assert set(scores["condition"]) == {"sound", "paint_failure", "rust"} print("ok wallpaper paste + full-coverage evaluate") # Temperature scaling recovers a known overconfidence factor. from train import fit_temperature class Overconfident(torch.nn.Module): def forward(self, x): n = x.shape[0] g = torch.Generator().manual_seed(0) logits = torch.randn(n, 10, 16, 16, generator=g) return {"semantic": logits * 4.0, "drywall": torch.zeros(n, 3, 16, 16)} g = torch.Generator().manual_seed(0) logits = torch.randn(8, 10, 16, 16, generator=g) labels = torch.distributions.Categorical(logits=logits.permute(0, 2, 3, 1)).sample() loader = [(torch.zeros(8, 3, 16, 16), labels, torch.full((8, 16, 16), IGNORE), torch.full((8, 16, 16), IGNORE))] t_sem, t_mat = fit_temperature(Overconfident(), loader, torch.device("cpu"), pixels_per_image=256) assert 3.0 < t_sem < 5.5 and t_mat == 1.0, (t_sem, t_mat) print("ok temperature scaling") print("training augmentation tests passed") if __name__ == "__main__": main()