test_train_aug.py

#!/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()