"""Model factory: shared semantic + drywall-material heads over several backbones. The original prototype hard-coded MobileNetV3-Large/DeepLabV3. That is small and fast but well behind modern open segmentation models. This module keeps the two-head design (semantic coverage + independent drywall substrate) but lets you choose the encoder: * ``deeplabv3_mobilenet_v3_large`` - the original; edge/latency oriented. * ``deeplabv3_resnet50`` / ``deeplabv3_resnet101`` - stronger torchvision baselines (ASPP + ResNet), no extra dependencies. * ``segformer_b0`` .. ``segformer_b5`` - modern transformer segmentation from Hugging Face (requires ``transformers``); the recommended accuracy option. All variants return ``{"semantic": (B, C, H, W), "drywall": (B, M, H, W)}`` at input resolution, where M is the number of material classes (3 with wallpaper; 2 for checkpoints trained before wallpaper existed), so ``train.py`` and ``predict.py`` are architecture-agnostic. Models built with condition classes also return ``"condition": (B, K, H, W)`` - surface condition (sound / paint failure / rust) from a third head on the same features; checkpoints from before that head load without it. Normalisation is shared (ImageNet mean/std), which SegFormer also expects. """ import torch from torch import nn from torch.nn import functional as F from torchvision.models import MobileNet_V3_Large_Weights, ResNet50_Weights, ResNet101_Weights from torchvision.models.segmentation import ( deeplabv3_mobilenet_v3_large, deeplabv3_resnet50, deeplabv3_resnet101, ) from label_schema import CONDITION_CLASSES, LEGACY_MATERIAL_CLASSES, MATERIAL_CLASSES # torchvision DeepLabV3 variants: name -> (builder, pretrained backbone weights). DEEPLAB_ARCHITECTURES = { "deeplabv3_mobilenet_v3_large": (deeplabv3_mobilenet_v3_large, MobileNet_V3_Large_Weights.DEFAULT), "deeplabv3_resnet50": (deeplabv3_resnet50, ResNet50_Weights.DEFAULT), "deeplabv3_resnet101": (deeplabv3_resnet101, ResNet101_Weights.DEFAULT), } # SegFormer variants: name -> Hugging Face id of a pretrained encoder/decoder. SEGFORMER_ARCHITECTURES = { "segformer_b0": "nvidia/segformer-b0-finetuned-ade-512-512", "segformer_b1": "nvidia/segformer-b1-finetuned-ade-512-512", "segformer_b2": "nvidia/segformer-b2-finetuned-ade-512-512", "segformer_b3": "nvidia/segformer-b3-finetuned-ade-512-512", "segformer_b4": "nvidia/segformer-b4-finetuned-ade-512-512", # b5 has no 512-input ADE20K checkpoint on the Hub; the only b5 release is 640. "segformer_b5": "nvidia/segformer-b5-finetuned-ade-640-640", } DEFAULT_ARCH = "deeplabv3_mobilenet_v3_large" def available_architectures(): return sorted(list(DEEPLAB_ARCHITECTURES) + list(SEGFORMER_ARCHITECTURES)) def _drywall_head(channels, num_material_classes, hidden=128, dropout=0.1): return nn.Sequential( nn.Conv2d(channels, hidden, kernel_size=3, padding=1, bias=False), nn.BatchNorm2d(hidden), nn.ReLU(inplace=True), nn.Dropout2d(dropout), nn.Conv2d(hidden, num_material_classes, kernel_size=1), ) class WallPaintNet(nn.Module): """torchvision DeepLabV3 semantic head plus an independent drywall head.""" def __init__(self, num_semantic_classes, pretrained_backbone=False, arch=DEFAULT_ARCH, num_material_classes=len(MATERIAL_CLASSES), num_condition_classes=len(CONDITION_CLASSES)): super().__init__() if arch not in DEEPLAB_ARCHITECTURES: raise ValueError(f"unknown DeepLab architecture {arch!r}; choose from {sorted(DEEPLAB_ARCHITECTURES)}") builder, default_weights = DEEPLAB_ARCHITECTURES[arch] weights = default_weights if pretrained_backbone else None self.arch = arch self.segmenter = builder(weights=None, weights_backbone=weights, num_classes=num_semantic_classes) was_training = self.segmenter.training self.segmenter.eval() with torch.inference_mode(): features = self.segmenter.backbone(torch.zeros(1, 3, 128, 128))["out"] self.segmenter.train(was_training) self.drywall_head = _drywall_head(features.shape[1], num_material_classes) self.condition_head = (_drywall_head(features.shape[1], num_condition_classes) if num_condition_classes else None) def forward(self, x): size = x.shape[-2:] features = self.segmenter.backbone(x) semantic = self.segmenter.classifier(features["out"]) drywall = self.drywall_head(features["out"]) out = { "semantic": F.interpolate(semantic, size=size, mode="bilinear", align_corners=False), "drywall": F.interpolate(drywall, size=size, mode="bilinear", align_corners=False), } if self.condition_head is not None: out["condition"] = F.interpolate(self.condition_head(features["out"]), size=size, mode="bilinear", align_corners=False) return out class SegformerPaintNet(nn.Module): """Hugging Face SegFormer semantic head plus an independent drywall head. SegFormer predicts logits at 1/4 resolution; the last encoder hidden state feeds the drywall head. Heavy enough to matter, but a genuine modern transformer backbone rather than a 2021 MobileNet. """ def __init__(self, num_semantic_classes, pretrained_backbone=False, arch="segformer_b2", num_material_classes=len(MATERIAL_CLASSES), num_condition_classes=len(CONDITION_CLASSES)): super().__init__() if arch not in SEGFORMER_ARCHITECTURES: raise ValueError(f"unknown SegFormer architecture {arch!r}; choose from {sorted(SEGFORMER_ARCHITECTURES)}") try: from transformers import SegformerConfig, SegformerForSemanticSegmentation except ImportError as exc: # pragma: no cover - optional dependency raise ImportError("SegFormer architectures require `pip install transformers`") from exc hf_id = SEGFORMER_ARCHITECTURES[arch] self.arch = arch if pretrained_backbone: self.segmenter = SegformerForSemanticSegmentation.from_pretrained( hf_id, num_labels=num_semantic_classes, ignore_mismatched_sizes=True) else: config = SegformerConfig.from_pretrained(hf_id, num_labels=num_semantic_classes) self.segmenter = SegformerForSemanticSegmentation(config) self.drywall_head = _drywall_head(self.segmenter.config.hidden_sizes[-1], num_material_classes) # Peeling edges and rust streaks are fine detail: the condition head reads the # highest-resolution encoder stage (1/4 scale) rather than the coarsest one. self.condition_head = (_drywall_head(self.segmenter.config.hidden_sizes[0], num_condition_classes) if num_condition_classes else None) def forward(self, x): size = x.shape[-2:] outputs = self.segmenter(pixel_values=x, output_hidden_states=True) semantic = F.interpolate(outputs.logits, size=size, mode="bilinear", align_corners=False) drywall = self.drywall_head(outputs.hidden_states[-1]) drywall = F.interpolate(drywall, size=size, mode="bilinear", align_corners=False) out = {"semantic": semantic, "drywall": drywall} if self.condition_head is not None: out["condition"] = F.interpolate(self.condition_head(outputs.hidden_states[0]), size=size, mode="bilinear", align_corners=False) return out def build_model(arch, num_semantic_classes, pretrained_backbone=False, num_material_classes=len(MATERIAL_CLASSES), num_condition_classes=len(CONDITION_CLASSES)): """Construct the model for ``arch`` (see :func:`available_architectures`); 0 condition classes = no condition head.""" if arch in DEEPLAB_ARCHITECTURES: return WallPaintNet(num_semantic_classes, pretrained_backbone, arch, num_material_classes, num_condition_classes) if arch in SEGFORMER_ARCHITECTURES: return SegformerPaintNet(num_semantic_classes, pretrained_backbone, arch, num_material_classes, num_condition_classes) raise ValueError(f"unknown architecture {arch!r}; choose from {available_architectures()}") def checkpoint_material_classes(checkpoint): """Material classes a checkpoint was trained with (pre-wallpaper checkpoints have two).""" names = tuple(checkpoint.get("material_classes") or ()) if names not in (MATERIAL_CLASSES, LEGACY_MATERIAL_CLASSES): raise ValueError(f"checkpoint material classes {names!r} are not a known schema; retrain with drywall_masks") return names def checkpoint_condition_classes(checkpoint): """Condition classes a checkpoint was trained with; () for checkpoints from before the condition head.""" names = tuple(checkpoint.get("condition_classes") or ()) if names not in ((), CONDITION_CLASSES): raise ValueError(f"checkpoint condition classes {names!r} are not a known schema") return names def build_from_checkpoint(checkpoint, num_semantic_classes): """Rebuild and load the model a checkpoint describes: 2 or 3 materials, with or without a condition head.""" model = build_model(str(checkpoint.get("arch", DEFAULT_ARCH)), num_semantic_classes, pretrained_backbone=False, num_material_classes=len(checkpoint_material_classes(checkpoint)), num_condition_classes=len(checkpoint_condition_classes(checkpoint))) model.load_state_dict(checkpoint["model"]) return model def warm_start(model, checkpoint): """Initialise ``model`` from a checkpoint of the same arch, growing the material head if needed. A pre-wallpaper checkpoint has a 2-way material classifier; its rows are copied into the new 3-way head and the wallpaper row keeps its fresh init, so facade/wall knowledge survives while wallpaper is learned. A checkpoint without a condition head leaves that head at its fresh init. Returns the names of tensors that could not be copied (empty when everything matched). """ target = model.state_dict() skipped = [] for name, value in checkpoint["model"].items(): if name not in target: skipped.append(name) elif target[name].shape == value.shape: target[name] = value elif name.startswith("drywall_head.") and value.shape[1:] == target[name].shape[1:] \ and value.shape[0] < target[name].shape[0]: target[name][:value.shape[0]] = value else: skipped.append(name) model.load_state_dict(target) return skipped