Upload models.py with huggingface_hub

models.py

@@ -0,0 +1,130 @@
+"""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, 2, H, W)}`` at
+input resolution, so ``train.py`` and ``predict.py`` are architecture-agnostic.
+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,
+)
+
+# 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",
+    "segformer_b5": "nvidia/segformer-b5-finetuned-ade-512-512",
+}
+DEFAULT_ARCH = "deeplabv3_mobilenet_v3_large"
+
+
+def available_architectures():
+    return sorted(list(DEEPLAB_ARCHITECTURES) + list(SEGFORMER_ARCHITECTURES))
+
+
+def _drywall_head(channels, 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, 2, 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):
+        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])
+
+    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"])
+        return {
+            "semantic": F.interpolate(semantic, size=size, mode="bilinear", align_corners=False),
+            "drywall": F.interpolate(drywall, size=size, mode="bilinear", align_corners=False),
+        }
+
+
+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"):
+        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])
+
+    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)
+        return {"semantic": semantic, "drywall": drywall}
+
+
+def build_model(arch, num_semantic_classes, pretrained_backbone=False):
+    """Construct the dual-head model for ``arch`` (see :func:`available_architectures`)."""
+    if arch in DEEPLAB_ARCHITECTURES:
+        return WallPaintNet(num_semantic_classes, pretrained_backbone, arch)
+    if arch in SEGFORMER_ARCHITECTURES:
+        return SegformerPaintNet(num_semantic_classes, pretrained_backbone, arch)
+    raise ValueError(f"unknown architecture {arch!r}; choose from {available_architectures()}")