Image Segmentation
PyTorch
English
painting-vision-robotics-kit
semantic-segmentation
robotics
edge-ai
construction-ai
autonomous-painting
wall-painting-robot
paint-coverage-estimation
building-facade
drywall
skirting-detection
window-detection
lidar
depth-validation
segformer
deeplabv3
Eval Results (legacy)
| #!/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() | |