Upload attention_unet+convnext model, code, and model card

README.md

@@ -0,0 +1,160 @@
+---
+license: apache-2.0
+library_name: pytorch
+tags:
+  - image-segmentation
+  - medical
+  - polyp
+  - colonoscopy
+  - unet
+  - attention-unet
+  - convnext
+datasets:
+  - andreribeiro87/kvasir-seg-augmented
+metrics:
+  - dice
+  - iou
+pipeline_tag: image-segmentation
+---
+
+# Attention U-Net ConvNeXt — Polyp Segmentation (Best Test Dice)
+
+Binary polyp segmentation model trained on [Kvasir-SEG](https://huggingface.co/datasets/Angelou0516/kvasir-seg).
+**Highest Dice score on the test set** (0.9411) among all 24 architecture × backbone combinations evaluated in
+the UNet-A benchmark sweep.
+
+## Model Description
+
+| Property | Value |
+|---|---|
+| Architecture | **Attention U-Net** (gate-based skip connections) |
+| Backbone | **ConvNeXt-Tiny** (ImageNet pre-trained via `timm`) |
+| Input size | 256 × 256 × 3 |
+| Output | 256 × 256 × 1 logit map (sigmoid → binary mask) |
+| Parameters | ~133 MB |
+| Loss | BCEDice (α = 0.5) |
+
+### Architecture Details
+
+**Attention U-Net** (Oktay et al., 2018) augments the standard encoder-decoder with *attention gates* on every
+skip connection. The gate computes a spatial attention coefficient from the decoder query and the encoder key,
+suppressing activations in irrelevant background regions and focusing the model on lesion boundaries.
+
+The **ConvNeXt-Tiny** backbone (pre-trained on ImageNet-1k) provides five resolution levels of feature maps.
+ConvNeXt's depthwise convolution design gives excellent feature quality with relatively low memory usage,
+making it well-suited for high-resolution segmentation.
+
+## Test Set Results
+
+Evaluated on the fixed 53-image test partition of Kvasir-SEG (50 % of the original validation split, seed 42):
+
+| Metric | Value |
+|---|---|
+| **Dice** | **0.9411** |
+| **IoU** | **0.8888** |
+| F1 | 0.9411 |
+| Precision | 0.9618 |
+| Recall | 0.9213 |
+| Accuracy | 0.9803 |
+
+## Sweep Leaderboard (all 24 models)
+
+| Rank | Model | Test Dice | Test IoU |
+|------|-------|-----------|---------|
+| **1** | **attention_unet_convnext (this model)** | **0.9411** | **0.8888** |
+| 2 | unet3plus_convnext | 0.9395 | 0.8859 |
+| 3 | unet_convnext | 0.9383 | 0.8838 |
+| 4 | resunet_efficientnet | 0.9338 | 0.8759 |
+| 5 | unet3plus_efficientnet | 0.9335 | 0.8753 |
+
+## Training Configuration
+
+- **Optimiser:** AdamW
+- **Learning rate:** 1e-3 (cosine decay, 5 % warmup)
+- **Weight decay:** 1e-4
+- **Batch size:** 64
+- **Epochs:** 50
+- **FP16:** enabled (A100 GPU)
+- **Loss:** BCEDice
+- **Dataset:** Kvasir-SEG augmented (5,368 train / 53 val / 53 test)
+- **Augmentation:** random H/V flips, ±30° rotation, brightness/contrast/saturation ±20 %
+
+> Note: This model was **not** subject to Optuna HPO. It is the raw sweep checkpoint.
+> For the HPO-tuned variant see [andreribeiro87/unet3plus-efficientnet-kvasir-seg](https://huggingface.co/andreribeiro87/unet3plus-efficientnet-kvasir-seg).
+
+## How to Use
+
+This model uses a custom PyTorch architecture. The model code is included in the repository.
+
+### Installation
+
+```bash
+pip install torch torchvision timm safetensors
+```
+
+### Inference
+
+```python
+import torch
+from safetensors.torch import load_file
+from torchvision.transforms import functional as TF
+from PIL import Image
+
+# 1. Clone or download the repo files
+# git clone https://huggingface.co/andreribeiro87/attention-unet-convnext-kvasir-seg
+# cd attention-unet-convnext-kvasir-seg
+
+# 2. Import the model
+from models import create_model
+
+# 3. Instantiate and load weights
+model = create_model("attention_unet", backbone="convnext")
+state_dict = load_file("model.safetensors")
+model.load_state_dict(state_dict)
+model.eval()
+
+# 4. Preprocess an image
+image = Image.open("your_colonoscopy_image.jpg").convert("RGB")
+x = TF.to_tensor(TF.resize(image, [256, 256])).unsqueeze(0)  # (1, 3, 256, 256)
+
+# 5. Predict
+with torch.no_grad():
+    logit = model(x)                          # (1, 1, 256, 256)
+    mask = (logit.sigmoid() > 0.5).squeeze()  # bool tensor (256, 256)
+
+# 6. Convert to PIL
+pred_mask = TF.to_pil_image(mask.float())
+```
+
+## Citation
+
+If you use this model or dataset, please cite the original Kvasir-SEG paper:
+
+```bibtex
+@inproceedings{jha2020kvasir,
+  title     = {Kvasir-SEG: A Segmented Polyp Dataset},
+  author    = {Jha, Debesh and Smedsrud, Pia H and Riegler, Michael A and Halvorsen, P{a}l
+               and de Lange, Thomas and Johansen, Dag and Johansen, H{a}vard D},
+  booktitle = {MultiMedia Modeling (MMM)},
+  year      = {2020}
+}
+```
+
+```bibtex
+@article{oktay2018attention,
+  title   = {Attention U-Net: Learning Where to Look for the Pancreas},
+  author  = {Oktay, Ozan and Schlemper, Jo and Folgoc, Loic Le and Lee, Matthew and Heinrich,
+             Mattias and Misawa, Kazunari and Mori, Kensaku and McDonagh, Steven and
+             Hammerla, Nils Y and Kainz, Bernhard and others},
+  journal = {arXiv preprint arXiv:1804.03999},
+  year    = {2018}
+}
+```
+
+## Limitations
+
+- Trained and evaluated exclusively on **Kvasir-SEG** (single-centre, single-modality).
+  Performance may degrade on other colonoscopy datasets or imaging conditions.
+- Binary segmentation only; does not distinguish between polyp types or severity.
+- Input resolution is fixed at **256 × 256**; very small polyps may not be fully captured.
+- **Not validated for clinical use.** This is a research model.

config.json

@@ -0,0 +1,27 @@
+{
+  "architecture": "attention_unet",
+  "backbone": "convnext",
+  "img_size": 256,
+  "in_channels": 3,
+  "out_channels": 1,
+  "training": {
+    "loss": "bce_dice",
+    "lr": 0.001,
+    "weight_decay": 0.0001,
+    "warmup_ratio": 0.05,
+    "scheduler": "cosine",
+    "batch_size": 64,
+    "epochs": 50,
+    "optimiser": "AdamW"
+  },
+  "test_metrics": {
+    "dice": 0.9411,
+    "iou": 0.8888,
+    "f1": 0.9411,
+    "precision": 0.9618,
+    "recall": 0.9213,
+    "accuracy": 0.9803
+  },
+  "dataset": "andreribeiro87/kvasir-seg-augmented",
+  "note": "Best test-Dice model from architecture sweep (no HPO applied)"
+}

loss.py

@@ -0,0 +1,94 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class DiceFocalLoss(nn.Module):
+    """Combined Dice + Focal loss for class-imbalanced binary segmentation.
+
+    Focal loss down-weights easy negatives, forcing the network to focus on
+    hard/uncertain pixels — the main failure mode when plateauing near 0.1.
+
+    Args:
+        alpha:        Focal weighting factor for positive class (0.25 typical).
+        gamma:        Focal modulating exponent.  Higher = more focus on hard
+                      pixels.  Tune in [0.5, 5.0] via Optuna.
+        dice_weight:  Weight of the Dice component.
+        focal_weight: Weight of the Focal component.
+        smooth:       Laplace smoothing for Dice denominator.
+    """
+
+    def __init__(
+        self,
+        alpha: float = 0.25,
+        gamma: float = 2.0,
+        dice_weight: float = 0.5,
+        focal_weight: float = 0.5,
+        smooth: float = 1e-6,
+    ) -> None:
+        super().__init__()
+        self.alpha = alpha
+        self.gamma = gamma
+        self.dice_weight = dice_weight
+        self.focal_weight = focal_weight
+        self.smooth = smooth
+
+    def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
+        # ---- Focal part ------------------------------------------------
+        bce = F.binary_cross_entropy_with_logits(y_pred, y_true.float(), reduction="none")
+        pt = torch.exp(-bce)
+        focal = self.alpha * (1.0 - pt) ** self.gamma * bce
+        focal_loss = focal.mean()
+
+        # ---- Dice part -------------------------------------------------
+        pred_sig = torch.sigmoid(y_pred)
+        inter = (pred_sig * y_true).sum(dim=(2, 3))
+        dice_loss = 1.0 - (2.0 * inter + self.smooth) / (
+            pred_sig.sum(dim=(2, 3)) + y_true.sum(dim=(2, 3)) + self.smooth
+        )
+        dice_loss = dice_loss.mean()
+
+        return self.dice_weight * dice_loss + self.focal_weight * focal_loss
+
+
+class DiceLoss(nn.Module):
+
+    def __init__(self, smooth: float = 1.0) -> None:
+        super().__init__()
+        self.smooth = smooth
+
+    def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
+        assert y_pred.size() == y_true.size()
+        y_pred = y_pred[:, 0].contiguous().view(-1)
+        y_true = y_true[:, 0].contiguous().view(-1)
+        intersection = (y_pred * y_true).sum()
+        dsc = (2.0 * intersection + self.smooth) / (
+            y_pred.sum() + y_true.sum() + self.smooth
+        )
+        return 1.0 - dsc
+
+
+class BCEDiceLoss(nn.Module):
+
+    def __init__(
+        self,
+        smooth: float = 1.0,
+        bce_weight: float = 0.5,
+        dice_weight: float = 0.5,
+        label_smoothing: float = 0.0,
+    ) -> None:
+        super().__init__()
+        self.bce_weight = bce_weight
+        self.dice_weight = dice_weight
+        self.label_smoothing = label_smoothing
+        self.dice = DiceLoss(smooth=smooth)
+
+    def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
+        if self.label_smoothing > 0.0:
+            # Smooth labels towards 0.5: prevents overconfident BCE
+            y_bce = y_true * (1.0 - self.label_smoothing) + self.label_smoothing * 0.5
+        else:
+            y_bce = y_true
+        bce = F.binary_cross_entropy_with_logits(y_pred, y_bce)
+        dice = self.dice(torch.sigmoid(y_pred), y_true)
+        return self.bce_weight * bce + self.dice_weight * dice

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b0c29289995adfc04f79dc1c5a0f1c0e0f85b3b19f11f0fab4ee0dd9699a2a5d
+size 139209476

models/__init__.py

@@ -0,0 +1,62 @@
+import torch.nn as nn
+
+from .attention_unet import AttentionUNet
+from .resunet import ResUNet
+from .transunet import TransUNet
+from .unet import UNet
+from .unet3plus import UNet3Plus
+from .unetplusplus import UNetPlusPlus
+
+# ---------------------------------------------------------------------------
+# Architecture registry
+# ---------------------------------------------------------------------------
+
+ARCHITECTURES: dict[str, type] = {
+    "unet": UNet,
+    "unetplusplus": UNetPlusPlus,
+    "unet3plus": UNet3Plus,
+    "attention_unet": AttentionUNet,
+    "resunet": ResUNet,
+    "transunet": TransUNet,
+}
+
+# Backbones available for all custom hierarchical architectures
+BACKBONES: list[str] = ["default", "efficientnet", "convnext", "swin", "siglip"]
+
+# Valid (architecture, backbone) combinations
+VALID_COMBINATIONS: dict[str, list[str]] = {
+    arch: BACKBONES for arch in ARCHITECTURES
+}
+
+# All architecture names
+ALL_ARCHITECTURES: list[str] = list(ARCHITECTURES)
+
+
+# ---------------------------------------------------------------------------
+# create_model — unified factory
+# ---------------------------------------------------------------------------
+
+def create_model(
+    architecture: str,
+    backbone: str,
+    in_channels: int = 3,
+    out_channels: int = 1,
+    **kwargs,
+) -> nn.Module:
+    valid_bbs = VALID_COMBINATIONS.get(architecture)
+    if valid_bbs is None:
+        raise ValueError(
+            f"Unknown architecture '{architecture}'. "
+            f"Choose from {ALL_ARCHITECTURES}"
+        )
+    if backbone not in valid_bbs:
+        raise ValueError(
+            f"Backbone '{backbone}' is not compatible with '{architecture}'. "
+            f"Valid choices: {valid_bbs}"
+        )
+    return ARCHITECTURES[architecture](
+        backbone_name=backbone,
+        in_channels=in_channels,
+        out_channels=out_channels,
+        **kwargs,
+    )

models/attention_unet.py

@@ -0,0 +1,44 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .backbones import create_backbone
+from .blocks import AttentionDecoderBlock
+
+
+class AttentionUNet(nn.Module):
+    def __init__(
+        self,
+        backbone_name: str = "default",
+        in_channels: int = 3,
+        out_channels: int = 1,
+        **_kw,
+    ):
+        super().__init__()
+        self.backbone = create_backbone(backbone_name, in_channels)
+        channels = self.backbone.out_channels
+        skip_channels, bottleneck_ch = channels[:-1], channels[-1]
+
+        self.decoders = nn.ModuleList()
+        in_ch = bottleneck_ch
+        for s_ch in reversed(skip_channels):
+            self.decoders.append(AttentionDecoderBlock(in_ch, s_ch, s_ch))
+            in_ch = s_ch
+
+        self.head = nn.Conv2d(in_ch, out_channels, kernel_size=1)
+
+    def forward(
+        self, x: torch.Tensor | None = None, pixel_values: torch.Tensor | None = None, **_kw,
+    ) -> torch.Tensor:
+        if pixel_values is not None:
+            x = pixel_values
+        input_size = x.shape[2:]
+
+        skips, x = self.backbone(x)
+        for dec, skip in zip(self.decoders, reversed(skips)):
+            x = dec(x, skip)
+
+        x = self.head(x)
+        if x.shape[2:] != input_size:
+            x = F.interpolate(x, size=input_size, mode="bilinear", align_corners=False)
+        return x

models/backbones.py

@@ -0,0 +1,317 @@
+import ssl
+from abc import ABC, abstractmethod
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision.models import (
+    ConvNeXt_Tiny_Weights,
+    EfficientNet_B0_Weights,
+    Swin_T_Weights,
+    convnext_tiny,
+    efficientnet_b0,
+    swin_t,
+)
+from torchvision.models.feature_extraction import create_feature_extractor
+
+from .blocks import ConvBlock
+
+
+def _disable_ssl_verification():
+    """Create an unverified SSL context for downloading pretrained weights."""
+    return ssl._create_unverified_context()
+
+
+# Workaround for SSL certificate verification issues when downloading pretrained weights
+_original_create_default_https_context = ssl._create_default_https_context
+
+
+def _enable_unverified_ssl():
+    ssl._create_default_https_context = _disable_ssl_verification
+
+
+def _restore_ssl():
+    ssl._create_default_https_context = _original_create_default_https_context
+
+
+class Backbone(ABC):
+    """All backbones return (skip_features high→low, bottleneck)."""
+
+    @property
+    @abstractmethod
+    def out_channels(self) -> list[int]:
+        """Channel counts from highest-res skip to bottleneck (last element)."""
+        ...
+
+    @abstractmethod
+    def forward(self, x: torch.Tensor) -> tuple[list[torch.Tensor], torch.Tensor]:
+        ...
+
+
+class DefaultBackbone(Backbone, nn.Module):
+    """Plain conv encoder identical to the classic U-Net."""
+
+    def __init__(self, in_channels: int = 3, block_cls: type | None = None, **_kw):
+        nn.Module.__init__(self)
+        block = block_cls or ConvBlock
+        self.enc1 = block(in_channels, 64)
+        self.enc2 = block(64, 128)
+        self.enc3 = block(128, 256)
+        self.enc4 = block(256, 512)
+        self.bottleneck = block(512, 1024)
+        self.pool = nn.MaxPool2d(2, 2)
+
+    @property
+    def out_channels(self) -> list[int]:
+        return [64, 128, 256, 512, 1024]
+
+    def forward(self, x: torch.Tensor):
+        skips = []
+        for enc in [self.enc1, self.enc2, self.enc3, self.enc4]:
+            x = enc(x)
+            skips.append(x)
+            x = self.pool(x)
+        return skips, self.bottleneck(x)
+
+
+class EfficientNetBackbone(Backbone, nn.Module):
+    """EfficientNet-B0 pretrained encoder (ImageNet)."""
+
+    _RETURN_NODES = {
+        "features.1": "s1",   # H/2,  16 ch
+        "features.2": "s2",   # H/4,  24 ch
+        "features.3": "s3",   # H/8,  40 ch
+        "features.5": "s4",   # H/16, 112 ch
+        "features.8": "bottleneck",  # H/32, 1280 ch
+    }
+
+    def __init__(self, in_channels: int = 3, **_kw):
+        nn.Module.__init__(self)
+        _enable_unverified_ssl()
+        try:
+            base = efficientnet_b0(weights=EfficientNet_B0_Weights.DEFAULT)
+        finally:
+            _restore_ssl()
+        if in_channels != 3:
+            old = base.features[0][0]
+            base.features[0][0] = nn.Conv2d(
+                in_channels, old.out_channels, old.kernel_size,
+                old.stride, old.padding, bias=False,
+            )
+        self.body = create_feature_extractor(base, return_nodes=self._RETURN_NODES)
+
+    @property
+    def out_channels(self) -> list[int]:
+        return [16, 24, 40, 112, 1280]
+
+    def forward(self, x: torch.Tensor):
+        f = self.body(x)
+        return [f["s1"], f["s2"], f["s3"], f["s4"]], f["bottleneck"]
+
+
+class ConvNeXTBackbone(Backbone, nn.Module):
+    """ConvNeXT-Tiny pretrained encoder (ImageNet)."""
+
+    _RETURN_NODES = {
+        "features.1": "s1",   # H/4,  96 ch
+        "features.3": "s2",   # H/8,  192 ch
+        "features.5": "s3",   # H/16, 384 ch
+        "features.7": "bottleneck",  # H/32, 768 ch
+    }
+
+    def __init__(self, in_channels: int = 3, **_kw):
+        nn.Module.__init__(self)
+        _enable_unverified_ssl()
+        try:
+            base = convnext_tiny(weights=ConvNeXt_Tiny_Weights.DEFAULT)
+        finally:
+            _restore_ssl()
+        if in_channels != 3:
+            old = base.features[0][0]
+            base.features[0][0] = nn.Conv2d(
+                in_channels, old.out_channels, old.kernel_size,
+                old.stride, old.padding, bias=False,
+            )
+        self.body = create_feature_extractor(base, return_nodes=self._RETURN_NODES)
+
+    @property
+    def out_channels(self) -> list[int]:
+        return [96, 192, 384, 768]
+
+    def forward(self, x: torch.Tensor):
+        f = self.body(x)
+        return [f["s1"], f["s2"], f["s3"]], f["bottleneck"]
+
+
+class SwinBackbone(Backbone, nn.Module):
+    """Swin Transformer Tiny pretrained encoder (ImageNet).
+
+    Produces hierarchical features at four scales — identical channel widths
+    to ConvNeXt-Tiny (96 → 192 → 384 → 768) — so it slots into every
+    existing decoder without any code changes.
+
+    Input (B, 3, H, W) → skips [(B,96,H/4,W/4), (B,192,H/8,W/8),
+                                  (B,384,H/16,W/16)], bottleneck (B,768,H/32,W/32)
+
+    Note: torchvision Swin outputs tensors in (B, H, W, C) layout;
+    the backbone permutes them to the standard (B, C, H, W) before returning.
+    """
+
+    _RETURN_NODES = {
+        "features.1": "s1",         # H/4,  96 ch
+        "features.3": "s2",         # H/8,  192 ch
+        "features.5": "s3",         # H/16, 384 ch
+        "features.7": "bottleneck", # H/32, 768 ch
+    }
+
+    def __init__(self, in_channels: int = 3, **_kw):
+        nn.Module.__init__(self)
+        _enable_unverified_ssl()
+        try:
+            base = swin_t(weights=Swin_T_Weights.DEFAULT)
+        finally:
+            _restore_ssl()
+        if in_channels != 3:
+            # Replace the first patch-embedding conv
+            old = base.features[0][0]
+            base.features[0][0] = nn.Conv2d(
+                in_channels, old.out_channels,
+                kernel_size=old.kernel_size, stride=old.stride,
+                padding=old.padding, bias=False,
+            )
+        self.body = create_feature_extractor(base, return_nodes=self._RETURN_NODES)
+
+    @property
+    def out_channels(self) -> list[int]:
+        return [96, 192, 384, 768]
+
+    def forward(self, x: torch.Tensor):
+        f = self.body(x)
+        # Swin-T stores features as (B, H, W, C) → convert to (B, C, H, W)
+        s1 = f["s1"].permute(0, 3, 1, 2).contiguous()
+        s2 = f["s2"].permute(0, 3, 1, 2).contiguous()
+        s3 = f["s3"].permute(0, 3, 1, 2).contiguous()
+        bn = f["bottleneck"].permute(0, 3, 1, 2).contiguous()
+        return [s1, s2, s3], bn
+
+
+class SigLIPBackbone(Backbone, nn.Module):
+    """SigLIP-Base/16 (93 M) pretrained vision encoder as a flat-ViT backbone.
+
+    Gemma 3 uses the larger ``google/siglip-so400m-patch14-384`` (400 M); this
+    class defaults to the practical ``google/siglip-base-patch16-224`` variant
+    (93 M) that fits comfortably alongside a UNet decoder. Swap MODEL_ID for
+    the Gemma 3 variant when VRAM permits.
+
+    Architecture note
+    -----------------
+    SigLIP is a pure Vision Transformer — every layer produces tokens at the
+    **same** spatial resolution (H/16 × W/16 = 14×14 for 224-px input). This
+    backbone therefore returns four feature maps that are **all at 14×14** but
+    at different semantic depths (layers 3 / 6 / 9 / 12 of the 12-layer ViT):
+
+        [z3, z6, z9]  →  skips  (each B, 768, 14, 14)
+        z12           →  bottleneck (B, 768, 14, 14)
+
+    This is intentionally different from the hierarchical CNN/Swin backbones.
+    Use the ``ViTUNet`` architecture (``models/vit_unet.py``) which handles the
+    flat-resolution skips by projecting and bilinearly upsampling them to match
+    each decoder stage. Standard UNet/AttentionUNet/ResUNet/TransUNet decoders
+    will NOT work correctly with this backbone.
+
+    Gemma 3 variant
+    ---------------
+    Replace MODEL_ID with ``"google/siglip-so400m-patch14-384"`` and set
+    ``PATCH_SIZE = 14``, ``INPUT_SIZE = 384`` to use the exact Gemma 3 encoder.
+    You will also need to adjust ``EXTRACT_LAYERS`` (the model has 27 layers).
+    """
+
+    MODEL_ID = "google/siglip-base-patch16-224"
+    PATCH_SIZE = 16
+    INPUT_SIZE = 224
+    # EXTRACT_LAYERS: quarter-intervals of the 12-layer ViT (1-indexed after embedding)
+    EXTRACT_LAYERS = (3, 6, 9, 12)
+
+    def __init__(self, in_channels: int = 3, **_kw):
+        nn.Module.__init__(self)
+        # google/siglip-base-patch16-224 is a full CLIP-style model (vision +
+        # text).  Loading with SiglipVisionModel.from_pretrained fails because
+        # the repo config is a SiglipConfig, not a SiglipVisionConfig.
+        # Solution: load the full SiglipModel and keep only the vision encoder.
+        from transformers import SiglipModel
+
+        full = SiglipModel.from_pretrained(self.MODEL_ID)
+        self._vit = full.vision_model        # SiglipVisionTransformer
+        self._hidden_dim: int = full.config.vision_config.hidden_size   # 768
+        del full                             # free text encoder weights
+
+        self._patch_grid = self.INPUT_SIZE // self.PATCH_SIZE           # 14
+
+        if in_channels != 3:
+            old = self._vit.embeddings.patch_embedding
+            self._vit.embeddings.patch_embedding = nn.Conv2d(
+                in_channels, self._hidden_dim,
+                kernel_size=self.PATCH_SIZE, stride=self.PATCH_SIZE, bias=False,
+            )
+
+        # Register forward hooks on the target transformer layers to capture
+        # intermediate hidden states.  This is version-agnostic — it works
+        # regardless of whether the transformers library honours
+        # `output_hidden_states=True` via its **kwargs API.
+        self._hooked: dict[int, torch.Tensor] = {}
+        self._hook_handles: list = []
+        for layer_idx in self.EXTRACT_LAYERS:
+            layer = self._vit.encoder.layers[layer_idx - 1]  # 0-indexed
+            handle = layer.register_forward_hook(self._make_hook(layer_idx))
+            self._hook_handles.append(handle)
+
+    def _make_hook(self, idx: int):
+        def _hook(module, inp, out):
+            # SiglipEncoderLayer returns a tuple; the first element is the
+            # hidden state tensor (B, N_patches, hidden_dim).
+            self._hooked[idx] = out[0] if isinstance(out, tuple) else out
+        return _hook
+
+    @property
+    def out_channels(self) -> list[int]:
+        return [self._hidden_dim] * len(self.EXTRACT_LAYERS)  # [768, 768, 768, 768]
+
+    def forward(self, x: torch.Tensor):
+        B = x.shape[0]
+        if x.shape[2] != self.INPUT_SIZE or x.shape[3] != self.INPUT_SIZE:
+            x = F.interpolate(
+                x, size=(self.INPUT_SIZE, self.INPUT_SIZE),
+                mode="bilinear", align_corners=False,
+            )
+
+        self._hooked.clear()
+        self._vit(pixel_values=x)     # hooks populate self._hooked
+
+        G = self._patch_grid           # 14
+        features = []
+        for layer_idx in self.EXTRACT_LAYERS:
+            hs = self._hooked[layer_idx]   # (B, 196, 768)
+            feat = (
+                hs.reshape(B, G, G, self._hidden_dim)
+                .permute(0, 3, 1, 2)
+                .contiguous()
+            )                              # (B, 768, 14, 14)
+            features.append(feat)
+
+        # z3=features[0] (early/local), z12=features[3] (final/global)
+        return features[:-1], features[-1]  # skips, bottleneck
+
+
+BACKBONE_REGISTRY: dict[str, type] = {
+    "default": DefaultBackbone,
+    "efficientnet": EfficientNetBackbone,
+    "convnext": ConvNeXTBackbone,
+    "swin": SwinBackbone,
+    "siglip": SigLIPBackbone,
+}
+
+
+def create_backbone(name: str, in_channels: int = 3, **kwargs) -> nn.Module:
+    if name not in BACKBONE_REGISTRY:
+        raise ValueError(f"Unknown backbone '{name}'. Choose from {list(BACKBONE_REGISTRY)}")
+    return BACKBONE_REGISTRY[name](in_channels=in_channels, **kwargs)

models/blocks.py

@@ -0,0 +1,140 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class ConvBlock(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int):
+        super().__init__()
+        self.block = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, 3, padding=1, bias=False),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(out_channels, out_channels, 3, padding=1, bias=False),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.block(x)
+
+
+class ResBlock(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int):
+        super().__init__()
+        self.conv = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, 3, padding=1, bias=False),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(out_channels, out_channels, 3, padding=1, bias=False),
+            nn.BatchNorm2d(out_channels),
+        )
+        self.shortcut = (
+            nn.Sequential(
+                nn.Conv2d(in_channels, out_channels, 1, bias=False),
+                nn.BatchNorm2d(out_channels),
+            )
+            if in_channels != out_channels
+            else nn.Identity()
+        )
+        self.relu = nn.ReLU(inplace=True)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.relu(self.conv(x) + self.shortcut(x))
+
+
+class AttentionGate(nn.Module):
+    """Additive attention gate: re-weights skip features using the decoder gate signal."""
+
+    def __init__(self, gate_channels: int, skip_channels: int):
+        super().__init__()
+        inter = max(1, skip_channels // 2)
+        self.W_gate = nn.Sequential(
+            nn.Conv2d(gate_channels, inter, 1, bias=False),
+            nn.BatchNorm2d(inter),
+        )
+        self.W_skip = nn.Sequential(
+            nn.Conv2d(skip_channels, inter, 1, bias=False),
+            nn.BatchNorm2d(inter),
+        )
+        self.psi = nn.Sequential(
+            nn.Conv2d(inter, 1, 1, bias=False),
+            nn.BatchNorm2d(1),
+            nn.Sigmoid(),
+        )
+        self.relu = nn.ReLU(inplace=True)
+
+    def forward(self, gate: torch.Tensor, skip: torch.Tensor) -> torch.Tensor:
+        g = self.W_gate(gate)
+        s = self.W_skip(skip)
+        if g.shape[2:] != s.shape[2:]:
+            g = F.interpolate(g, size=s.shape[2:], mode="bilinear", align_corners=False)
+        return skip * self.psi(self.relu(g + s))
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, dim: int, num_heads: int = 8, mlp_ratio: float = 4.0, dropout: float = 0.1):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(dim)
+        self.attn = nn.MultiheadAttention(dim, num_heads, dropout=dropout, batch_first=True)
+        self.norm2 = nn.LayerNorm(dim)
+        self.mlp = nn.Sequential(
+            nn.Linear(dim, int(dim * mlp_ratio)),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(int(dim * mlp_ratio), dim),
+            nn.Dropout(dropout),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        h = self.norm1(x)
+        h, _ = self.attn(h, h, h)
+        x = x + h
+        x = x + self.mlp(self.norm2(x))
+        return x
+
+
+# ---------------------------------------------------------------------------
+# Decoder blocks (one per architecture variant)
+# ---------------------------------------------------------------------------
+
+
+class UNetDecoderBlock(nn.Module):
+    def __init__(self, in_ch: int, skip_ch: int, out_ch: int):
+        super().__init__()
+        self.up = nn.ConvTranspose2d(in_ch, out_ch, kernel_size=2, stride=2)
+        self.conv = ConvBlock(out_ch + skip_ch, out_ch)
+
+    def forward(self, x: torch.Tensor, skip: torch.Tensor) -> torch.Tensor:
+        x = self.up(x)
+        if x.shape[2:] != skip.shape[2:]:
+            x = F.interpolate(x, size=skip.shape[2:], mode="bilinear", align_corners=False)
+        return self.conv(torch.cat([x, skip], dim=1))
+
+
+class AttentionDecoderBlock(nn.Module):
+    def __init__(self, in_ch: int, skip_ch: int, out_ch: int):
+        super().__init__()
+        self.up = nn.ConvTranspose2d(in_ch, out_ch, kernel_size=2, stride=2)
+        self.attn = AttentionGate(gate_channels=out_ch, skip_channels=skip_ch)
+        self.conv = ConvBlock(out_ch + skip_ch, out_ch)
+
+    def forward(self, x: torch.Tensor, skip: torch.Tensor) -> torch.Tensor:
+        x = self.up(x)
+        if x.shape[2:] != skip.shape[2:]:
+            x = F.interpolate(x, size=skip.shape[2:], mode="bilinear", align_corners=False)
+        skip = self.attn(gate=x, skip=skip)
+        return self.conv(torch.cat([x, skip], dim=1))
+
+
+class ResDecoderBlock(nn.Module):
+    def __init__(self, in_ch: int, skip_ch: int, out_ch: int):
+        super().__init__()
+        self.up = nn.ConvTranspose2d(in_ch, out_ch, kernel_size=2, stride=2)
+        self.conv = ResBlock(out_ch + skip_ch, out_ch)
+
+    def forward(self, x: torch.Tensor, skip: torch.Tensor) -> torch.Tensor:
+        x = self.up(x)
+        if x.shape[2:] != skip.shape[2:]:
+            x = F.interpolate(x, size=skip.shape[2:], mode="bilinear", align_corners=False)
+        return self.conv(torch.cat([x, skip], dim=1))

models/resunet.py

@@ -0,0 +1,47 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .backbones import create_backbone
+from .blocks import ResBlock, ResDecoderBlock
+
+
+class ResUNet(nn.Module):
+    def __init__(
+        self,
+        backbone_name: str = "default",
+        in_channels: int = 3,
+        out_channels: int = 1,
+        **_kw,
+    ):
+        super().__init__()
+        backbone_kwargs: dict = {}
+        if backbone_name == "default":
+            backbone_kwargs["block_cls"] = ResBlock
+        self.backbone = create_backbone(backbone_name, in_channels, **backbone_kwargs)
+        channels = self.backbone.out_channels
+        skip_channels, bottleneck_ch = channels[:-1], channels[-1]
+
+        self.decoders = nn.ModuleList()
+        in_ch = bottleneck_ch
+        for s_ch in reversed(skip_channels):
+            self.decoders.append(ResDecoderBlock(in_ch, s_ch, s_ch))
+            in_ch = s_ch
+
+        self.head = nn.Conv2d(in_ch, out_channels, kernel_size=1)
+
+    def forward(
+        self, x: torch.Tensor | None = None, pixel_values: torch.Tensor | None = None, **_kw,
+    ) -> torch.Tensor:
+        if pixel_values is not None:
+            x = pixel_values
+        input_size = x.shape[2:]
+
+        skips, x = self.backbone(x)
+        for dec, skip in zip(self.decoders, reversed(skips)):
+            x = dec(x, skip)
+
+        x = self.head(x)
+        if x.shape[2:] != input_size:
+            x = F.interpolate(x, size=input_size, mode="bilinear", align_corners=False)
+        return x

models/transunet.py

@@ -0,0 +1,71 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .backbones import create_backbone
+from .blocks import TransformerBlock, UNetDecoderBlock
+
+
+class TransUNet(nn.Module):
+    def __init__(
+        self,
+        backbone_name: str = "default",
+        in_channels: int = 3,
+        out_channels: int = 1,
+        img_size: int = 256,
+        transformer_dim: int = 512,
+        num_heads: int = 8,
+        num_layers: int = 6,
+        mlp_ratio: float = 4.0,
+        **_kw,
+    ):
+        super().__init__()
+        self.backbone = create_backbone(backbone_name, in_channels)
+        channels = self.backbone.out_channels
+        skip_channels, bottleneck_ch = channels[:-1], channels[-1]
+
+        with torch.no_grad():
+            dummy = torch.zeros(1, in_channels, img_size, img_size)
+            _, dummy_bn = self.backbone(dummy)
+            num_patches = dummy_bn.shape[2] * dummy_bn.shape[3]
+
+        self.proj_in = nn.Conv2d(bottleneck_ch, transformer_dim, kernel_size=1)
+        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, transformer_dim))
+        nn.init.trunc_normal_(self.pos_embed, std=0.02)
+        self.transformer = nn.Sequential(
+            *[TransformerBlock(transformer_dim, num_heads, mlp_ratio) for _ in range(num_layers)]
+        )
+        self.proj_out = nn.Conv2d(transformer_dim, bottleneck_ch, kernel_size=1)
+
+        self.decoders = nn.ModuleList()
+        in_ch = bottleneck_ch
+        for s_ch in reversed(skip_channels):
+            self.decoders.append(UNetDecoderBlock(in_ch, s_ch, s_ch))
+            in_ch = s_ch
+
+        self.head = nn.Conv2d(in_ch, out_channels, kernel_size=1)
+
+    def forward(
+        self, x: torch.Tensor | None = None, pixel_values: torch.Tensor | None = None, **_kw,
+    ) -> torch.Tensor:
+        if pixel_values is not None:
+            x = pixel_values
+        input_size = x.shape[2:]
+
+        skips, bottleneck = self.backbone(x)
+
+        B, _C, H, W = bottleneck.shape
+        t = self.proj_in(bottleneck)               # B, D, H, W
+        t = t.flatten(2).transpose(1, 2)           # B, N, D
+        t = t + self.pos_embed
+        t = self.transformer(t)                    # B, N, D
+        t = t.transpose(1, 2).view(B, -1, H, W)   # B, D, H, W
+        x = self.proj_out(t)                       # B, bottleneck_ch, H, W
+
+        for dec, skip in zip(self.decoders, reversed(skips)):
+            x = dec(x, skip)
+
+        x = self.head(x)
+        if x.shape[2:] != input_size:
+            x = F.interpolate(x, size=input_size, mode="bilinear", align_corners=False)
+        return x

models/unet.py

@@ -0,0 +1,44 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .backbones import create_backbone
+from .blocks import UNetDecoderBlock
+
+
+class UNet(nn.Module):
+    def __init__(
+        self,
+        backbone_name: str = "default",
+        in_channels: int = 3,
+        out_channels: int = 1,
+        **_kw,
+    ):
+        super().__init__()
+        self.backbone = create_backbone(backbone_name, in_channels)
+        channels = self.backbone.out_channels
+        skip_channels, bottleneck_ch = channels[:-1], channels[-1]
+
+        self.decoders = nn.ModuleList()
+        in_ch = bottleneck_ch
+        for s_ch in reversed(skip_channels):
+            self.decoders.append(UNetDecoderBlock(in_ch, s_ch, s_ch))
+            in_ch = s_ch
+
+        self.head = nn.Conv2d(in_ch, out_channels, kernel_size=1)
+
+    def forward(
+        self, x: torch.Tensor | None = None, pixel_values: torch.Tensor | None = None, **_kw,
+    ) -> torch.Tensor:
+        if pixel_values is not None:
+            x = pixel_values
+        input_size = x.shape[2:]
+
+        skips, x = self.backbone(x)
+        for dec, skip in zip(self.decoders, reversed(skips)):
+            x = dec(x, skip)
+
+        x = self.head(x)
+        if x.shape[2:] != input_size:
+            x = F.interpolate(x, size=input_size, mode="bilinear", align_corners=False)
+        return x

models/unet3plus.py

@@ -0,0 +1,120 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .backbones import create_backbone
+from .blocks import ConvBlock
+
+
+class UNet3Plus(nn.Module):
+    """UNet 3+ with full-scale skip connections.
+
+    Every decoder node aggregates features from ALL encoder levels (not just
+    the matching scale), giving each node simultaneous access to fine-grained
+    detail and deep semantic context. Each incoming stream is projected to
+    `inter_ch` channels before concatenation so total input channels are
+    predictable regardless of backbone channel widths.
+
+    Reference: Huang et al., "UNet 3+: A Full-Scale Connected UNet for
+    Medical Image Segmentation", ICASSP 2020.
+    """
+
+    def __init__(
+        self,
+        backbone_name: str = "default",
+        in_channels: int = 3,
+        out_channels: int = 1,
+        inter_ch: int = 64,
+        **_kw,
+    ):
+        super().__init__()
+        self.backbone = create_backbone(backbone_name, in_channels)
+        all_ch = list(self.backbone.out_channels)  # [skip0, ..., skip_{D-1}, bottleneck]
+        D = len(all_ch) - 1                        # number of skip levels = decoder levels
+        self._D = D
+        self._inter_ch = inter_ch
+        self._dec_ch = (D + 1) * inter_ch          # fixed output channels for every decoder node
+
+        # --- projections: each incoming stream is independently projected to inter_ch ---
+        # enc_projs[k][j]: encoder/bottleneck level j → inter_ch, for use at decoder k
+        self.enc_projs = nn.ModuleList([
+            nn.ModuleList([
+                nn.Sequential(
+                    nn.Conv2d(all_ch[j], inter_ch, 1, bias=False),
+                    nn.BatchNorm2d(inter_ch),
+                    nn.ReLU(inplace=True),
+                )
+                for j in range(D + 1)
+            ])
+            for k in range(D)
+        ])
+
+        # dec_projs[k][i]: prior decoder level i → inter_ch, for use at decoder k  (i < k)
+        # All prior decoder outputs have self._dec_ch channels.
+        self.dec_projs = nn.ModuleList([
+            nn.ModuleList([
+                nn.Sequential(
+                    nn.Conv2d(self._dec_ch, inter_ch, 1, bias=False),
+                    nn.BatchNorm2d(inter_ch),
+                    nn.ReLU(inplace=True),
+                )
+                for _ in range(k)           # one projection per prior decoder level
+            ])
+            for k in range(D)
+        ])
+
+        # Fusion conv at each decoder level: (D+1+k)*inter_ch → dec_ch
+        self.dec_convs = nn.ModuleList([
+            ConvBlock((D + 1 + k) * inter_ch, self._dec_ch)
+            for k in range(D)
+        ])
+
+        self.head = nn.Conv2d(self._dec_ch, out_channels, kernel_size=1)
+
+    def forward(
+        self,
+        x: torch.Tensor | None = None,
+        pixel_values: torch.Tensor | None = None,
+        **_kw,
+    ) -> torch.Tensor:
+        if pixel_values is not None:
+            x = pixel_values
+        input_size = x.shape[2:]
+
+        skips, bottleneck = self.backbone(x)
+        D = self._D
+        # enc_feats: shallowest encoder first, bottleneck last  (D+1 tensors)
+        enc_feats = list(skips) + [bottleneck]
+
+        dec_outputs: list[torch.Tensor] = []
+
+        # Build decoder from deepest (k=0) to shallowest (k=D-1).
+        # Decoder k targets the same spatial resolution as skips[D-1-k].
+        for k in range(D):
+            target = skips[D - 1 - k].shape[2:]
+            parts: list[torch.Tensor] = []
+
+            # All encoder / bottleneck streams
+            for j, feat in enumerate(enc_feats):
+                feat = _resize(feat, target)
+                parts.append(self.enc_projs[k][j](feat))
+
+            # All prior decoder streams (deeper → current scale = always upsample)
+            for i, df in enumerate(dec_outputs):
+                feat = F.interpolate(df, size=target, mode="bilinear", align_corners=False)
+                parts.append(self.dec_projs[k][i](feat))
+
+            dec_outputs.append(self.dec_convs[k](torch.cat(parts, dim=1)))
+
+        out = self.head(dec_outputs[-1])
+        if out.shape[2:] != input_size:
+            out = F.interpolate(out, size=input_size, mode="bilinear", align_corners=False)
+        return out
+
+
+def _resize(feat: torch.Tensor, target: tuple[int, int] | torch.Size) -> torch.Tensor:
+    if feat.shape[2:] == torch.Size(target):
+        return feat
+    if feat.shape[2] > target[0]:
+        return F.adaptive_max_pool2d(feat, target)
+    return F.interpolate(feat, size=target, mode="bilinear", align_corners=False)

models/unetplusplus.py

@@ -0,0 +1,83 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .backbones import create_backbone
+from .blocks import ConvBlock
+
+
+class UNetPlusPlus(nn.Module):
+    """UNet++ with dense nested skip connections.
+
+    Each node x[i][j] aggregates all previous same-scale nodes x[i][0..j-1]
+    plus an upsampled feature from the level below x[i+1][j-1], enabling
+    the decoder to learn progressively richer skip representations before
+    the final prediction at x[0][D].
+    See this: https://arxiv.org/abs/1807.10165
+    """
+
+    def __init__(
+        self,
+        backbone_name: str = "default",
+        in_channels: int = 3,
+        out_channels: int = 1,
+        **_kw,
+    ):
+        super().__init__()
+        self.backbone = create_backbone(backbone_name, in_channels)
+        channels = self.backbone.out_channels
+        skip_channels = list(channels[:-1])  # shallowest → deepest
+        bottleneck_ch = channels[-1]
+        self._D = len(skip_channels)
+
+        # all_ch[i] = channel count of the raw encoder node x[i][0]
+        all_ch = skip_channels + [bottleneck_ch]
+
+        # Dense nodes: x[i][j] for j > 0.
+        # x[i][j] concatenates: j same-scale predecessors + 1 upsampled from below.
+        # All intermediate node outputs keep skip_channels[i] channels.
+        self.nodes = nn.ModuleDict()
+        for j in range(1, self._D + 1):
+            for i in range(self._D - j + 1):
+                # channels coming from same scale: j tensors each with skip_channels[i] channels
+                from_same = j * skip_channels[i]
+                # channels coming from one level deeper (upsampled)
+                from_below = skip_channels[i + 1] if j > 1 else all_ch[i + 1]
+                self.nodes[f"{i}_{j}"] = ConvBlock(from_same + from_below, skip_channels[i])
+
+        self.head = nn.Conv2d(skip_channels[0], out_channels, kernel_size=1)
+
+    def forward(
+        self,
+        x: torch.Tensor | None = None,
+        pixel_values: torch.Tensor | None = None,
+        **_kw,
+    ) -> torch.Tensor:
+        if pixel_values is not None:
+            x = pixel_values
+        input_size = x.shape[2:]
+
+        skips, bottleneck = self.backbone(x)
+        D = self._D
+
+        # Initialise node cache with encoder outputs
+        cache: dict[tuple[int, int], torch.Tensor] = {}
+        for i, s in enumerate(skips):
+            cache[(i, 0)] = s
+        cache[(D, 0)] = bottleneck
+
+        # Fill the dense grid column by column (increasing j)
+        for j in range(1, D + 1):
+            for i in range(D - j + 1):
+                prev = [cache[(i, k)] for k in range(j)]
+                target_size = prev[0].shape[2:]
+                below = F.interpolate(
+                    cache[(i + 1, j - 1)], size=target_size,
+                    mode="bilinear", align_corners=False,
+                )
+                cache[(i, j)] = self.nodes[f"{i}_{j}"](torch.cat(prev + [below], dim=1))
+
+        out = self.head(cache[(0, D)])
+        if out.shape[2:] != input_size:
+            out = F.interpolate(out, size=input_size, mode="bilinear", align_corners=False)
+        return out