Upload test_train_aug.py with huggingface_hub

test_train_aug.py

@@ -63,6 +63,7 @@ def main():
     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")
 
@@ -76,6 +77,56 @@ def main():
     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")