Upload folder using huggingface_hub

code/framework/common/subset.py

@@ -0,0 +1,40 @@
+"""Deterministic low-data subset selection.
+
+The SAME function is used by (a) the segmentation trainer's real-data loader and
+(b) the pixel-diffusion generator's data loader, so that a low-data experiment at
+fraction f trains the generator on EXACTLY the real images the segmenter sees —
+no leakage, no mismatch, fully reproducible.
+
+Selection is driven only by (fraction, fraction_seed) and the SORTED list of
+items, so it is independent of the training seed: the 3 segmentation seeds all
+use the identical real subset, only their weight init differs.
+"""
+from __future__ import annotations
+
+import math
+import random
+from typing import List, Sequence, TypeVar
+
+T = TypeVar("T")
+
+
+def select_fraction(items: Sequence[T], fraction: float, seed: int = 0) -> List[T]:
+    """Return a deterministic subset of `items` of size ceil(fraction*N).
+
+    items: any sequence (e.g. list of (image_path, mask_path) tuples). It is
+           sorted by repr() first so order of discovery never affects the subset.
+    fraction: in (0, 1]. >=1 returns all items (sorted copy).
+    seed: subset seed (NOT the training seed). Fixed across training seeds.
+    """
+    ordered = sorted(items, key=lambda x: repr(x))
+    n = len(ordered)
+    if fraction >= 1.0 or n == 0:
+        return ordered
+    if fraction <= 0.0:
+        raise ValueError(f"fraction must be in (0,1], got {fraction}")
+    k = max(1, math.ceil(fraction * n))
+    rng = random.Random(seed)
+    idx = list(range(n))
+    rng.shuffle(idx)
+    keep = sorted(idx[:k])
+    return [ordered[i] for i in keep]

code/framework/config.py

@@ -0,0 +1,116 @@
+"""Unified experiment configuration.
+
+A single dataclass drives every run. Values can come from (in priority order):
+  1. command-line flags (argparse)  2. a YAML file (--config)  3. dataclass defaults.
+
+The same config object is used by train.py / test.py so that a training run and
+its evaluation are guaranteed to agree on dataset, model, image size, etc.
+"""
+from __future__ import annotations
+
+import argparse
+import dataclasses
+from dataclasses import dataclass, field, asdict
+from typing import Optional, List
+
+import yaml
+
+
+@dataclass
+class Config:
+    # ---- experiment identity ----
+    exp_name: str = "default"                 # results/<exp_name>/<dataset>/<arch>/seed<seed>/
+    seed: int = 0
+
+    # ---- data ----
+    data_root: str = "dataset/processed_unified"
+    dataset: str = "cvc_clinicdb"             # folder name under data_root
+    protocol: str = "official"                # e.g. official / fold01 ...
+    in_channels: int = 0                      # 0 = auto-detect from metadata/first image
+    num_classes: int = 0                      # 0 = auto-detect from metadata/masks (incl. background)
+    img_size: int = 256                       # square resize target (Swin/TransUNet need 224)
+    # extra synthetic (image,mask) pairs to MERGE into the train split.
+    # Points at a dir laid out like a split: <synth_train_dir>/{images,masks}/.
+    synth_train_dir: str = ""                 # "" = real data only (no generative augmentation)
+
+    # low-data regime: train on a deterministic fraction of the REAL train split.
+    # The SAME subset is used by the PixDiff generator (framework.common.subset), so
+    # real and real+synth arms never disagree. fraction_seed is FIXED across the 3
+    # training seeds (only weight init varies), so the data subset stays constant.
+    train_fraction: float = 1.0               # 1.0 = full data; e.g. 0.1 = 10%
+    fraction_seed: int = 0                    # subset seed (independent of `seed`)
+
+    # ---- augmentation (conventional baseline tier) ----
+    aug: str = "standard"                     # none | standard | strong  (albumentations online)
+    aug_backend: str = "albumentations"       # albumentations | monai
+    normalize: str = "auto"                   # auto(imagenet for RGB, 0.5 for gray) | imagenet | none
+
+    # ---- model ----
+    arch: str = "unet"                        # see models/registry.py REGISTRY
+    encoder: str = "resnet34"                 # SMP encoder name (ignored by non-SMP archs)
+    encoder_weights: str = "imagenet"         # imagenet | none
+    pretrained_ckpt: str = ""                 # ViT/Swin pretrain for transunet/swinunet (optional)
+
+    # ---- optimization ----
+    epochs: int = 100
+    batch_size: int = 16                      # per-GPU batch size
+    lr: float = 1e-4
+    weight_decay: float = 1e-4
+    optimizer: str = "adamw"                  # adamw | sgd
+    scheduler: str = "poly"                   # poly | cosine | none
+    warmup_epochs: int = 0
+    loss: str = "ce_dice"                     # ce_dice | ce | dice
+    num_workers: int = 8
+    grad_clip: float = 0.0                    # 0 = disabled
+
+    # ---- precision / hardware ----
+    amp: str = "bf16"                         # bf16(A100+) | fp16(V100) | fp32
+    # DDP is driven by torchrun env vars (RANK/WORLD_SIZE/LOCAL_RANK); nothing to set here.
+
+    # ---- evaluation / logging ----
+    val_interval: int = 5                     # epochs between validations
+    min_epochs: int = 0                       # never early-stop before this many epochs
+    patience: int = 0                          # early-stop after this many epochs w/o val improvement (0 = off)
+    save_interval: int = 0                    # 0 = only save best + last
+    include_background: bool = False          # include class 0 in reported Dice/IoU
+    compute_hd95: bool = True
+    out_root: str = "results"
+    resume: str = ""                          # path to checkpoint to resume from
+    visualize: bool = True                    # save overlays at test time
+    vis_max: int = 32                         # max number of overlay images to save
+
+    def out_dir(self) -> str:
+        return f"{self.out_root}/{self.exp_name}/{self.dataset}_{self.protocol}/{self.arch}/seed{self.seed}"
+
+    def to_yaml(self, path: str) -> None:
+        with open(path, "w") as f:
+            yaml.safe_dump(asdict(self), f, sort_keys=False, allow_unicode=True)
+
+    @classmethod
+    def from_args(cls, argv: Optional[List[str]] = None) -> "Config":
+        # First pass: only grab --config so YAML can set defaults that flags then override.
+        pre = argparse.ArgumentParser(add_help=False)
+        pre.add_argument("--config", type=str, default="")
+        known, _ = pre.parse_known_args(argv)
+
+        base = cls()
+        if known.config:
+            with open(known.config) as f:
+                ydata = yaml.safe_load(f) or {}
+            base = dataclasses.replace(base, **{k: v for k, v in ydata.items()
+                                                if k in {f.name for f in dataclasses.fields(cls)}})
+
+        p = argparse.ArgumentParser(parents=[pre],
+                                    description="SegGen unified segmentation framework")
+        for f in dataclasses.fields(cls):
+            default = getattr(base, f.name)
+            if f.type is bool or isinstance(default, bool):
+                # support --flag / --no-flag
+                p.add_argument(f"--{f.name}", dest=f.name, action="store_true", default=default)
+                p.add_argument(f"--no-{f.name}", dest=f.name, action="store_false")
+            else:
+                p.add_argument(f"--{f.name}", type=type(default) if default is not None else str,
+                               default=default)
+        ns = p.parse_args(argv)
+        kwargs = {f.name: getattr(ns, f.name) for f in dataclasses.fields(cls)}
+        return cls(**kwargs)

code/framework/config.py.bak

@@ -0,0 +1,109 @@
+"""Unified experiment configuration.
+
+A single dataclass drives every run. Values can come from (in priority order):
+  1. command-line flags (argparse)  2. a YAML file (--config)  3. dataclass defaults.
+
+The same config object is used by train.py / test.py so that a training run and
+its evaluation are guaranteed to agree on dataset, model, image size, etc.
+"""
+from __future__ import annotations
+
+import argparse
+import dataclasses
+from dataclasses import dataclass, field, asdict
+from typing import Optional, List
+
+import yaml
+
+
+@dataclass
+class Config:
+    # ---- experiment identity ----
+    exp_name: str = "default"                 # results/<exp_name>/<dataset>/<arch>/seed<seed>/
+    seed: int = 0
+
+    # ---- data ----
+    data_root: str = "dataset/processed_unified"
+    dataset: str = "cvc_clinicdb"             # folder name under data_root
+    protocol: str = "official"                # e.g. official / fold01 ...
+    in_channels: int = 0                      # 0 = auto-detect from metadata/first image
+    num_classes: int = 0                      # 0 = auto-detect from metadata/masks (incl. background)
+    img_size: int = 256                       # square resize target (Swin/TransUNet need 224)
+    # extra synthetic (image,mask) pairs to MERGE into the train split.
+    # Points at a dir laid out like a split: <synth_train_dir>/{images,masks}/.
+    synth_train_dir: str = ""                 # "" = real data only (no generative augmentation)
+
+    # ---- augmentation (conventional baseline tier) ----
+    aug: str = "standard"                     # none | standard | strong  (albumentations online)
+    aug_backend: str = "albumentations"       # albumentations | monai
+    normalize: str = "auto"                   # auto(imagenet for RGB, 0.5 for gray) | imagenet | none
+
+    # ---- model ----
+    arch: str = "unet"                        # see models/registry.py REGISTRY
+    encoder: str = "resnet34"                 # SMP encoder name (ignored by non-SMP archs)
+    encoder_weights: str = "imagenet"         # imagenet | none
+    pretrained_ckpt: str = ""                 # ViT/Swin pretrain for transunet/swinunet (optional)
+
+    # ---- optimization ----
+    epochs: int = 100
+    batch_size: int = 16                      # per-GPU batch size
+    lr: float = 1e-4
+    weight_decay: float = 1e-4
+    optimizer: str = "adamw"                  # adamw | sgd
+    scheduler: str = "poly"                   # poly | cosine | none
+    warmup_epochs: int = 0
+    loss: str = "ce_dice"                     # ce_dice | ce | dice
+    num_workers: int = 8
+    grad_clip: float = 0.0                    # 0 = disabled
+
+    # ---- precision / hardware ----
+    amp: str = "bf16"                         # bf16(A100+) | fp16(V100) | fp32
+    # DDP is driven by torchrun env vars (RANK/WORLD_SIZE/LOCAL_RANK); nothing to set here.
+
+    # ---- evaluation / logging ----
+    val_interval: int = 5                     # epochs between validations
+    min_epochs: int = 0                       # never early-stop before this many epochs
+    patience: int = 0                          # early-stop after this many epochs w/o val improvement (0 = off)
+    save_interval: int = 0                    # 0 = only save best + last
+    include_background: bool = False          # include class 0 in reported Dice/IoU
+    compute_hd95: bool = True
+    out_root: str = "results"
+    resume: str = ""                          # path to checkpoint to resume from
+    visualize: bool = True                    # save overlays at test time
+    vis_max: int = 32                         # max number of overlay images to save
+
+    def out_dir(self) -> str:
+        return f"{self.out_root}/{self.exp_name}/{self.dataset}_{self.protocol}/{self.arch}/seed{self.seed}"
+
+    def to_yaml(self, path: str) -> None:
+        with open(path, "w") as f:
+            yaml.safe_dump(asdict(self), f, sort_keys=False, allow_unicode=True)
+
+    @classmethod
+    def from_args(cls, argv: Optional[List[str]] = None) -> "Config":
+        # First pass: only grab --config so YAML can set defaults that flags then override.
+        pre = argparse.ArgumentParser(add_help=False)
+        pre.add_argument("--config", type=str, default="")
+        known, _ = pre.parse_known_args(argv)
+
+        base = cls()
+        if known.config:
+            with open(known.config) as f:
+                ydata = yaml.safe_load(f) or {}
+            base = dataclasses.replace(base, **{k: v for k, v in ydata.items()
+                                                if k in {f.name for f in dataclasses.fields(cls)}})
+
+        p = argparse.ArgumentParser(parents=[pre],
+                                    description="SegGen unified segmentation framework")
+        for f in dataclasses.fields(cls):
+            default = getattr(base, f.name)
+            if f.type is bool or isinstance(default, bool):
+                # support --flag / --no-flag
+                p.add_argument(f"--{f.name}", dest=f.name, action="store_true", default=default)
+                p.add_argument(f"--no-{f.name}", dest=f.name, action="store_false")
+            else:
+                p.add_argument(f"--{f.name}", type=type(default) if default is not None else str,
+                               default=default)
+        ns = p.parse_args(argv)
+        kwargs = {f.name: getattr(ns, f.name) for f in dataclasses.fields(cls)}
+        return cls(**kwargs)

code/framework/data/loaders.py

@@ -0,0 +1,41 @@
+"""Build datasets / dataloaders from a Config, consistent across train & test."""
+from __future__ import annotations
+
+from torch.utils.data import DataLoader
+from torch.utils.data.distributed import DistributedSampler
+
+from .unified_dataset import UnifiedSegDataset
+from .transforms import build_transform
+from ..engine.distributed import is_dist
+
+
+def build_dataset(cfg, split: str) -> UnifiedSegDataset:
+    train = (split == "train")
+    synth = cfg.synth_train_dir if train else ""
+    # construct without transform first so in_channels/num_classes auto-detect runs
+    ds = UnifiedSegDataset(
+        data_root=cfg.data_root, dataset=cfg.dataset, protocol=cfg.protocol, split=split,
+        transform=None, in_channels=cfg.in_channels, num_classes=cfg.num_classes,
+        synth_dir=synth,
+        train_fraction=cfg.train_fraction, fraction_seed=cfg.fraction_seed,
+    )
+    ds.transform = build_transform(cfg.img_size, ds.in_channels, train=train,
+                                   aug=cfg.aug, normalize=cfg.normalize)
+    return ds
+
+
+def build_loader(cfg, split: str, ds: UnifiedSegDataset) -> DataLoader:
+    train = (split == "train")
+    sampler = None
+    if is_dist():
+        sampler = DistributedSampler(ds, shuffle=train, drop_last=train)
+    return DataLoader(
+        ds,
+        batch_size=cfg.batch_size,
+        shuffle=(train and sampler is None),
+        sampler=sampler,
+        num_workers=cfg.num_workers,
+        pin_memory=True,
+        drop_last=(train and sampler is None),
+        persistent_workers=cfg.num_workers > 0,
+    )

code/framework/data/loaders.py.bak

@@ -0,0 +1,40 @@
+"""Build datasets / dataloaders from a Config, consistent across train & test."""
+from __future__ import annotations
+
+from torch.utils.data import DataLoader
+from torch.utils.data.distributed import DistributedSampler
+
+from .unified_dataset import UnifiedSegDataset
+from .transforms import build_transform
+from ..engine.distributed import is_dist
+
+
+def build_dataset(cfg, split: str) -> UnifiedSegDataset:
+    train = (split == "train")
+    synth = cfg.synth_train_dir if train else ""
+    # construct without transform first so in_channels/num_classes auto-detect runs
+    ds = UnifiedSegDataset(
+        data_root=cfg.data_root, dataset=cfg.dataset, protocol=cfg.protocol, split=split,
+        transform=None, in_channels=cfg.in_channels, num_classes=cfg.num_classes,
+        synth_dir=synth,
+    )
+    ds.transform = build_transform(cfg.img_size, ds.in_channels, train=train,
+                                   aug=cfg.aug, normalize=cfg.normalize)
+    return ds
+
+
+def build_loader(cfg, split: str, ds: UnifiedSegDataset) -> DataLoader:
+    train = (split == "train")
+    sampler = None
+    if is_dist():
+        sampler = DistributedSampler(ds, shuffle=train, drop_last=train)
+    return DataLoader(
+        ds,
+        batch_size=cfg.batch_size,
+        shuffle=(train and sampler is None),
+        sampler=sampler,
+        num_workers=cfg.num_workers,
+        pin_memory=True,
+        drop_last=(train and sampler is None),
+        persistent_workers=cfg.num_workers > 0,
+    )

code/framework/data/transforms.py

@@ -0,0 +1,78 @@
+"""Albumentations-based transform builder (the conventional-augmentation tier).
+
+Key correctness guarantees for segmentation:
+  * masks always use NEAREST interpolation (integer class ids never blended);
+    every geometric transform sets mask_interpolation=cv2.INTER_NEAREST.
+  * image and mask receive the SAME random spatial parameters (Albumentations
+    applies one transform jointly to image= and mask=).
+
+`aug` presets:  none (resize+normalize only) | standard | strong.
+Returns a callable: (image HWC uint8, mask HW int) -> (FloatTensor[C,H,W], LongTensor[H,W]).
+"""
+from __future__ import annotations
+
+from typing import Tuple
+
+import cv2
+import numpy as np
+import torch
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
+
+
+_IMAGENET_MEAN = (0.485, 0.456, 0.406)
+_IMAGENET_STD = (0.229, 0.224, 0.225)
+
+
+def _normalize(in_channels: int, normalize: str) -> A.Normalize:
+    if normalize == "none":
+        mean = (0.0,) * in_channels
+        std = (1.0,) * in_channels
+    elif normalize == "imagenet" and in_channels == 3:
+        mean, std = _IMAGENET_MEAN, _IMAGENET_STD
+    else:  # auto
+        if in_channels == 3:
+            mean, std = _IMAGENET_MEAN, _IMAGENET_STD
+        else:
+            mean, std = (0.5,) * in_channels, (0.5,) * in_channels
+    return A.Normalize(mean=mean, std=std, max_pixel_value=255.0)
+
+
+def build_transform(img_size: int, in_channels: int, train: bool,
+                    aug: str = "standard", normalize: str = "auto"):
+    N = cv2.INTER_NEAREST
+    L = cv2.INTER_LINEAR
+    ops = []
+
+    if train and aug != "none":
+        ops += [
+            A.Resize(img_size, img_size, interpolation=L, mask_interpolation=N),
+            A.HorizontalFlip(p=0.5),
+            A.VerticalFlip(p=0.5),
+            A.Affine(scale=(0.9, 1.1), translate_percent=(0.0, 0.05),
+                     rotate=(-15, 15), interpolation=L, mask_interpolation=N, p=0.5),
+        ]
+        if aug == "strong":
+            ops += [
+                A.ElasticTransform(alpha=30, sigma=6, interpolation=L,
+                                   mask_interpolation=N, p=0.3),
+                A.GridDistortion(num_steps=5, distort_limit=0.2, interpolation=L,
+                                 mask_interpolation=N, p=0.3),
+                A.RandomBrightnessContrast(p=0.5),
+                A.GaussNoise(p=0.2),
+            ]
+            if in_channels == 3:
+                ops.append(A.CLAHE(p=0.2))
+    else:
+        ops.append(A.Resize(img_size, img_size, interpolation=L, mask_interpolation=N))
+
+    ops += [_normalize(in_channels, normalize), ToTensorV2()]
+    compose = A.Compose(ops)
+
+    def _apply(image: np.ndarray, mask: np.ndarray) -> Tuple[torch.Tensor, torch.Tensor]:
+        out = compose(image=image, mask=mask)
+        img = out["image"].float()                 # C,H,W
+        msk = out["mask"].long()                    # H,W
+        return img, msk
+
+    return _apply

code/framework/data/unified_dataset.py

@@ -0,0 +1,187 @@
+"""Dataset reader for the standardized `processed_unified` layout.
+
+Expected layout (see dataset/SEGMENTATION_WORKSPACE_README.md):
+    <data_root>/<dataset>/<protocol>/<split>/images/*.png
+    <data_root>/<dataset>/<protocol>/<split>/masks/*.png
+    <data_root>/<dataset>/metadata.json          (optional, preferred)
+    <data_root>/<dataset>/manifest.jsonl         (optional)
+
+Returns per item: {"image": FloatTensor[C,H,W], "mask": LongTensor[H,W], "name": str}.
+
+Binary and multi-class masks are both supported: masks keep their integer class
+ids (0..C-1). Auto-detection of in_channels / num_classes falls back to scanning
+files when metadata is absent, so the loader is robust to missing metadata.
+"""
+from __future__ import annotations
+
+import json
+import os
+from glob import glob
+from typing import Optional, Callable, List, Tuple
+
+import numpy as np
+import cv2
+from torch.utils.data import Dataset
+
+
+_MODALITY_CHANNELS = {  # hint table; only used when metadata lacks in_channels
+    "rgb": 3, "fundus": 3, "colonoscopy": 3, "endoscopy": 3, "histopathology": 3,
+    "ultrasound": 1, "mri": 1, "ct": 1, "grayscale": 1,
+}
+
+# Documented class counts (incl. background). metadata.json on the server has no
+# num_classes field, so this table is the fast, reliable primary source; unknown
+# datasets fall back to a FULL scan of the mask set (accurate but slower).
+_KNOWN_NUM_CLASSES = {
+    "cvc_clinicdb": 2, "kvasir_seg": 2, "fives": 2, "busi": 2,
+    "refuge2": 3, "acdc_png": 4,
+    "idridd_segmentation": 6, "pannuke_semantic": 6,
+}
+
+
+def _read_metadata(data_root: str, dataset: str) -> dict:
+    path = os.path.join(data_root, dataset, "metadata.json")
+    if os.path.isfile(path):
+        try:
+            with open(path) as f:
+                return json.load(f)
+        except Exception:
+            return {}
+    return {}
+
+
+def _pair_from_manifest(split_dir: str, manifest: str) -> Optional[List[Tuple[str, str]]]:
+    if not os.path.isfile(manifest):
+        return None
+    pairs = []
+    base = os.path.dirname(manifest)
+    with open(manifest) as f:
+        for line in f:
+            line = line.strip()
+            if not line:
+                continue
+            rec = json.loads(line)
+            img = rec.get("image") or rec.get("image_path") or rec.get("img")
+            msk = rec.get("mask") or rec.get("mask_path") or rec.get("label")
+            if img is None or msk is None:
+                return None
+            # manifest paths may be relative to dataset root or absolute
+            ip = img if os.path.isabs(img) else os.path.join(base, img)
+            mp = msk if os.path.isabs(msk) else os.path.join(base, msk)
+            # only keep entries that fall under this split dir
+            if os.path.normpath(split_dir) in os.path.normpath(ip):
+                pairs.append((ip, mp))
+    return pairs or None
+
+
+def _pair_by_glob(split_dir: str) -> List[Tuple[str, str]]:
+    img_dir = os.path.join(split_dir, "images")
+    msk_dir = os.path.join(split_dir, "masks")
+    imgs = sorted(glob(os.path.join(img_dir, "*")))
+    pairs = []
+    for ip in imgs:
+        stem = os.path.splitext(os.path.basename(ip))[0]
+        # mask may share extension or be .png
+        cands = glob(os.path.join(msk_dir, stem + ".*"))
+        if not cands:
+            continue
+        pairs.append((ip, cands[0]))
+    return pairs
+
+
+def detect_in_channels(meta: dict, sample_img: Optional[str]) -> int:
+    if meta.get("in_channels"):
+        return int(meta["in_channels"])
+    mod = str(meta.get("modality", "")).lower()
+    for k, v in _MODALITY_CHANNELS.items():
+        if k in mod:
+            return v
+    if sample_img and os.path.isfile(sample_img):
+        im = cv2.imread(sample_img, cv2.IMREAD_UNCHANGED)
+        if im is not None and im.ndim == 3 and im.shape[2] >= 3:
+            return 3
+    return 1
+
+
+def detect_num_classes(meta: dict, mask_paths: List[str], dataset: str = "") -> int:
+    if dataset in _KNOWN_NUM_CLASSES:
+        return _KNOWN_NUM_CLASSES[dataset]
+    if meta.get("num_classes"):
+        return int(meta["num_classes"])
+    # unknown dataset: scan ALL masks so a rare class is never missed
+    vals = set()
+    for mp in mask_paths:
+        m = cv2.imread(mp, cv2.IMREAD_GRAYSCALE)
+        if m is not None:
+            vals.update(np.unique(m).tolist())
+    if not vals:
+        return 2
+    maxv = max(vals)
+    return int(maxv) + 1 if maxv >= 1 else 2
+
+
+class UnifiedSegDataset(Dataset):
+    def __init__(self, data_root: str, dataset: str, protocol: str, split: str,
+                 transform: Optional[Callable] = None,
+                 in_channels: int = 0, num_classes: int = 0,
+                 synth_dir: str = "",
+                 train_fraction: float = 1.0, fraction_seed: int = 0):
+        self.data_root = data_root
+        self.dataset = dataset
+        self.split = split
+        self.transform = transform
+
+        split_dir = os.path.join(data_root, dataset, protocol, split)
+        if not os.path.isdir(split_dir):
+            raise FileNotFoundError(
+                f"split dir not found: {split_dir}\n"
+                f"(data is prepared separately; see dataset/ scripts)")
+
+        manifest = os.path.join(data_root, dataset, "manifest.jsonl")
+        pairs = _pair_from_manifest(split_dir, manifest) or _pair_by_glob(split_dir)
+        if not pairs:
+            raise RuntimeError(f"no (image,mask) pairs found in {split_dir}")
+
+        # low-data: deterministically subsample REAL train pairs (the generator uses
+        # the SAME select_fraction, so no leakage) BEFORE merging any synthetic data.
+        if split == "train" and train_fraction < 1.0:
+            from ..common.subset import select_fraction
+            pairs = select_fraction(pairs, train_fraction, fraction_seed)
+
+        # optionally merge synthetic (image,mask) pairs into the (train) split
+        if synth_dir and os.path.isdir(synth_dir):
+            sp = _pair_by_glob(synth_dir if os.path.isdir(os.path.join(synth_dir, "images"))
+                               else os.path.dirname(synth_dir))
+            pairs = pairs + sp
+
+        self.pairs = pairs
+        meta = _read_metadata(data_root, dataset)
+        self.in_channels = in_channels or detect_in_channels(meta, pairs[0][0])
+        self.num_classes = num_classes or detect_num_classes(meta, [p[1] for p in pairs], dataset)
+
+    def __len__(self) -> int:
+        return len(self.pairs)
+
+    def _load_image(self, path: str) -> np.ndarray:
+        if self.in_channels == 1:
+            im = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
+            if im is None:
+                raise IOError(f"cannot read image {path}")
+            return im[:, :, None]                      # H,W,1
+        im = cv2.imread(path, cv2.IMREAD_COLOR)        # BGR
+        if im is None:
+            raise IOError(f"cannot read image {path}")
+        return cv2.cvtColor(im, cv2.COLOR_BGR2RGB)     # H,W,3
+
+    def __getitem__(self, idx: int):
+        ip, mp = self.pairs[idx]
+        image = self._load_image(ip)
+        mask = cv2.imread(mp, cv2.IMREAD_GRAYSCALE)
+        if mask is None:
+            raise IOError(f"cannot read mask {mp}")
+        mask = mask.astype(np.int64)
+
+        if self.transform is not None:
+            image, mask = self.transform(image, mask)
+        return {"image": image, "mask": mask,
+                "name": os.path.splitext(os.path.basename(ip))[0]}

code/framework/data/unified_dataset.py.bak

@@ -0,0 +1,180 @@
+"""Dataset reader for the standardized `processed_unified` layout.
+
+Expected layout (see dataset/SEGMENTATION_WORKSPACE_README.md):
+    <data_root>/<dataset>/<protocol>/<split>/images/*.png
+    <data_root>/<dataset>/<protocol>/<split>/masks/*.png
+    <data_root>/<dataset>/metadata.json          (optional, preferred)
+    <data_root>/<dataset>/manifest.jsonl         (optional)
+
+Returns per item: {"image": FloatTensor[C,H,W], "mask": LongTensor[H,W], "name": str}.
+
+Binary and multi-class masks are both supported: masks keep their integer class
+ids (0..C-1). Auto-detection of in_channels / num_classes falls back to scanning
+files when metadata is absent, so the loader is robust to missing metadata.
+"""
+from __future__ import annotations
+
+import json
+import os
+from glob import glob
+from typing import Optional, Callable, List, Tuple
+
+import numpy as np
+import cv2
+from torch.utils.data import Dataset
+
+
+_MODALITY_CHANNELS = {  # hint table; only used when metadata lacks in_channels
+    "rgb": 3, "fundus": 3, "colonoscopy": 3, "endoscopy": 3, "histopathology": 3,
+    "ultrasound": 1, "mri": 1, "ct": 1, "grayscale": 1,
+}
+
+# Documented class counts (incl. background). metadata.json on the server has no
+# num_classes field, so this table is the fast, reliable primary source; unknown
+# datasets fall back to a FULL scan of the mask set (accurate but slower).
+_KNOWN_NUM_CLASSES = {
+    "cvc_clinicdb": 2, "kvasir_seg": 2, "fives": 2, "busi": 2,
+    "refuge2": 3, "acdc_png": 4,
+    "idridd_segmentation": 6, "pannuke_semantic": 6,
+}
+
+
+def _read_metadata(data_root: str, dataset: str) -> dict:
+    path = os.path.join(data_root, dataset, "metadata.json")
+    if os.path.isfile(path):
+        try:
+            with open(path) as f:
+                return json.load(f)
+        except Exception:
+            return {}
+    return {}
+
+
+def _pair_from_manifest(split_dir: str, manifest: str) -> Optional[List[Tuple[str, str]]]:
+    if not os.path.isfile(manifest):
+        return None
+    pairs = []
+    base = os.path.dirname(manifest)
+    with open(manifest) as f:
+        for line in f:
+            line = line.strip()
+            if not line:
+                continue
+            rec = json.loads(line)
+            img = rec.get("image") or rec.get("image_path") or rec.get("img")
+            msk = rec.get("mask") or rec.get("mask_path") or rec.get("label")
+            if img is None or msk is None:
+                return None
+            # manifest paths may be relative to dataset root or absolute
+            ip = img if os.path.isabs(img) else os.path.join(base, img)
+            mp = msk if os.path.isabs(msk) else os.path.join(base, msk)
+            # only keep entries that fall under this split dir
+            if os.path.normpath(split_dir) in os.path.normpath(ip):
+                pairs.append((ip, mp))
+    return pairs or None
+
+
+def _pair_by_glob(split_dir: str) -> List[Tuple[str, str]]:
+    img_dir = os.path.join(split_dir, "images")
+    msk_dir = os.path.join(split_dir, "masks")
+    imgs = sorted(glob(os.path.join(img_dir, "*")))
+    pairs = []
+    for ip in imgs:
+        stem = os.path.splitext(os.path.basename(ip))[0]
+        # mask may share extension or be .png
+        cands = glob(os.path.join(msk_dir, stem + ".*"))
+        if not cands:
+            continue
+        pairs.append((ip, cands[0]))
+    return pairs
+
+
+def detect_in_channels(meta: dict, sample_img: Optional[str]) -> int:
+    if meta.get("in_channels"):
+        return int(meta["in_channels"])
+    mod = str(meta.get("modality", "")).lower()
+    for k, v in _MODALITY_CHANNELS.items():
+        if k in mod:
+            return v
+    if sample_img and os.path.isfile(sample_img):
+        im = cv2.imread(sample_img, cv2.IMREAD_UNCHANGED)
+        if im is not None and im.ndim == 3 and im.shape[2] >= 3:
+            return 3
+    return 1
+
+
+def detect_num_classes(meta: dict, mask_paths: List[str], dataset: str = "") -> int:
+    if dataset in _KNOWN_NUM_CLASSES:
+        return _KNOWN_NUM_CLASSES[dataset]
+    if meta.get("num_classes"):
+        return int(meta["num_classes"])
+    # unknown dataset: scan ALL masks so a rare class is never missed
+    vals = set()
+    for mp in mask_paths:
+        m = cv2.imread(mp, cv2.IMREAD_GRAYSCALE)
+        if m is not None:
+            vals.update(np.unique(m).tolist())
+    if not vals:
+        return 2
+    maxv = max(vals)
+    return int(maxv) + 1 if maxv >= 1 else 2
+
+
+class UnifiedSegDataset(Dataset):
+    def __init__(self, data_root: str, dataset: str, protocol: str, split: str,
+                 transform: Optional[Callable] = None,
+                 in_channels: int = 0, num_classes: int = 0,
+                 synth_dir: str = ""):
+        self.data_root = data_root
+        self.dataset = dataset
+        self.split = split
+        self.transform = transform
+
+        split_dir = os.path.join(data_root, dataset, protocol, split)
+        if not os.path.isdir(split_dir):
+            raise FileNotFoundError(
+                f"split dir not found: {split_dir}\n"
+                f"(data is prepared separately; see dataset/ scripts)")
+
+        manifest = os.path.join(data_root, dataset, "manifest.jsonl")
+        pairs = _pair_from_manifest(split_dir, manifest) or _pair_by_glob(split_dir)
+        if not pairs:
+            raise RuntimeError(f"no (image,mask) pairs found in {split_dir}")
+
+        # optionally merge synthetic (image,mask) pairs into the (train) split
+        if synth_dir and os.path.isdir(synth_dir):
+            sp = _pair_by_glob(synth_dir if os.path.isdir(os.path.join(synth_dir, "images"))
+                               else os.path.dirname(synth_dir))
+            pairs = pairs + sp
+
+        self.pairs = pairs
+        meta = _read_metadata(data_root, dataset)
+        self.in_channels = in_channels or detect_in_channels(meta, pairs[0][0])
+        self.num_classes = num_classes or detect_num_classes(meta, [p[1] for p in pairs], dataset)
+
+    def __len__(self) -> int:
+        return len(self.pairs)
+
+    def _load_image(self, path: str) -> np.ndarray:
+        if self.in_channels == 1:
+            im = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
+            if im is None:
+                raise IOError(f"cannot read image {path}")
+            return im[:, :, None]                      # H,W,1
+        im = cv2.imread(path, cv2.IMREAD_COLOR)        # BGR
+        if im is None:
+            raise IOError(f"cannot read image {path}")
+        return cv2.cvtColor(im, cv2.COLOR_BGR2RGB)     # H,W,3
+
+    def __getitem__(self, idx: int):
+        ip, mp = self.pairs[idx]
+        image = self._load_image(ip)
+        mask = cv2.imread(mp, cv2.IMREAD_GRAYSCALE)
+        if mask is None:
+            raise IOError(f"cannot read mask {mp}")
+        mask = mask.astype(np.int64)
+
+        if self.transform is not None:
+            image, mask = self.transform(image, mask)
+        return {"image": image, "mask": mask,
+                "name": os.path.splitext(os.path.basename(ip))[0]}

code/framework/efficiency.py

@@ -0,0 +1,106 @@
+"""Per-architecture efficiency table: #params, FLOPs (GMac), inference throughput.
+
+Representative setting (in_channels=3, num_classes=2). FLOPs via thop -> fvcore ->
+ptflops (whichever is installed); params and throughput always computed. Run once,
+on a GPU (A100). Output: results/<exp>/efficiency.{md,csv}.
+
+  python framework/efficiency.py --img_size 256 --out_root results --exp_name baselines
+"""
+from __future__ import annotations
+
+import os
+import sys
+import time
+import json
+import argparse
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+import torch
+
+from framework.models.registry import build_model, required_img_size
+
+ARCHS = ["unet", "unetpp", "deeplabv3plus", "attention_unet", "transunet", "swinunet"]
+
+
+def count_flops_gmac(model, x):
+    try:
+        from thop import profile
+        macs, _ = profile(model, inputs=(x,), verbose=False)
+        return macs / 1e9
+    except Exception:
+        pass
+    try:
+        from fvcore.nn import FlopCountAnalysis
+        return FlopCountAnalysis(model, x).total() / 1e9
+    except Exception:
+        pass
+    return float("nan")
+
+
+@torch.no_grad()
+def throughput(model, x, iters=50, warmup=10):
+    for _ in range(warmup):
+        model(x)
+    if x.is_cuda:
+        torch.cuda.synchronize()
+    t0 = time.time()
+    for _ in range(iters):
+        model(x)
+    if x.is_cuda:
+        torch.cuda.synchronize()
+    dt = time.time() - t0
+    return iters * x.size(0) / dt          # images / sec
+
+
+def encoder_for(arch):
+    if arch in ("unet", "unetpp", "deeplabv3plus"):
+        return "resnet50"
+    if arch == "transunet":
+        return "R50-ViT-B_16"
+    return "resnet34"
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--img_size", type=int, default=256)
+    ap.add_argument("--batch_size", type=int, default=8)
+    ap.add_argument("--out_root", default="results")
+    ap.add_argument("--exp_name", default="baselines")
+    args = ap.parse_args()
+
+    dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    rows = []
+    for arch in ARCHS:
+        sz = required_img_size(arch) or args.img_size
+        model = build_model(arch, in_channels=3, num_classes=2, img_size=sz,
+                            encoder=encoder_for(arch), encoder_weights="none").to(dev).eval()
+        params_m = sum(p.numel() for p in model.parameters()) / 1e6
+        x1 = torch.randn(1, 3, sz, sz, device=dev)
+        gmac = count_flops_gmac(model, x1)
+        xb = torch.randn(args.batch_size, 3, sz, sz, device=dev)
+        try:
+            ips = throughput(model, xb)
+        except Exception as e:
+            ips = float("nan"); print(f"[warn] throughput {arch}: {e}")
+        rows.append({"arch": arch, "img": sz, "params_M": round(params_m, 2),
+                     "gmac": round(gmac, 2) if gmac == gmac else None,
+                     "imgs_per_s": round(ips, 1) if ips == ips else None})
+        print(f"{arch:16s} img={sz} params={params_m:.2f}M GMac={gmac:.2f} {ips:.1f} img/s")
+        del model, x1, xb
+        if dev.type == "cuda":
+            torch.cuda.empty_cache()
+
+    base = os.path.join(args.out_root, args.exp_name)
+    os.makedirs(base, exist_ok=True)
+    with open(os.path.join(base, "efficiency.json"), "w") as f:
+        json.dump(rows, f, indent=2)
+    md = "| Method | Img | Params(M) | GMac | Img/s |\n|---|---|---|---|---|\n"
+    for r in rows:
+        md += f"| {r['arch']} | {r['img']} | {r['params_M']} | {r['gmac']} | {r['imgs_per_s']} |\n"
+    open(os.path.join(base, "efficiency.md"), "w").write(md)
+    print(md)
+    print(f"written {base}/efficiency.{{json,md}}")
+
+
+if __name__ == "__main__":
+    main()

code/framework/engine/distributed.py

@@ -0,0 +1,85 @@
+"""Minimal DDP helpers driven entirely by torchrun environment variables.
+
+Launch with:  torchrun --nproc_per_node=<N> framework/train.py ...
+Single-process (no torchrun) also works: world_size falls back to 1.
+"""
+from __future__ import annotations
+
+import os
+import random
+from typing import List, Any
+
+import numpy as np
+import torch
+import torch.distributed as dist
+
+
+def is_dist() -> bool:
+    return dist.is_available() and dist.is_initialized()
+
+
+def get_rank() -> int:
+    return dist.get_rank() if is_dist() else 0
+
+
+def get_world_size() -> int:
+    return dist.get_world_size() if is_dist() else 1
+
+
+def is_main() -> bool:
+    return get_rank() == 0
+
+
+def setup_distributed() -> int:
+    """Init the process group if launched under torchrun. Returns local_rank."""
+    if "RANK" in os.environ and "WORLD_SIZE" in os.environ:
+        local_rank = int(os.environ.get("LOCAL_RANK", 0))
+        torch.cuda.set_device(local_rank)
+        # bind this rank's device to the PG so collectives/barrier don't guess
+        # GPU 0 (avoids the NCCL "devices unknown" warning + potential hang)
+        try:
+            dist.init_process_group(backend="nccl", init_method="env://",
+                                    device_id=torch.device("cuda", local_rank))
+        except TypeError:  # older torch without device_id kwarg
+            dist.init_process_group(backend="nccl", init_method="env://")
+        dist.barrier(device_ids=[local_rank])
+        return local_rank
+    # single GPU / CPU fallback
+    if torch.cuda.is_available():
+        torch.cuda.set_device(0)
+    return 0
+
+
+def cleanup_distributed() -> None:
+    if is_dist():
+        dist.barrier()
+        dist.destroy_process_group()
+
+
+def all_gather_object(obj: Any) -> List[Any]:
+    """Gather arbitrary picklable objects from all ranks into a flat list."""
+    if not is_dist():
+        return [obj]
+    out: List[Any] = [None for _ in range(get_world_size())]
+    dist.all_gather_object(out, obj)
+    return out
+
+
+def set_seed(seed: int, rank: int = 0, deterministic: bool = False) -> None:
+    """Seed all RNGs. Each rank gets a distinct stream (seed + rank) so DDP
+    workers don't draw identical augmentation noise, while staying reproducible."""
+    s = seed + rank
+    random.seed(s)
+    np.random.seed(s)
+    torch.manual_seed(s)
+    torch.cuda.manual_seed_all(s)
+    if deterministic:
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False
+    else:
+        torch.backends.cudnn.benchmark = True
+
+
+def print_main(*args, **kwargs) -> None:
+    if is_main():
+        print(*args, **kwargs, flush=True)

code/framework/engine/evaluator.py

@@ -0,0 +1,74 @@
+"""Test-time evaluation: load best checkpoint, run on the test split, compute
+Dice / IoU / HD95 (per image -> mean +- SD), optionally save overlay visualizations
+and a metrics.json. Runs single-process (rank 0) for deterministic reporting.
+"""
+from __future__ import annotations
+
+import os
+import json
+
+import numpy as np
+import torch
+
+from ..data.loaders import build_dataset, build_loader
+from ..metrics.metrics import per_image_metrics, aggregate
+from ..visualize.overlay import save_overlay
+
+
+@torch.no_grad()
+def evaluate(cfg, model, device, ckpt_path: str = "") -> dict:
+    ds = build_dataset(cfg, "test")
+    num_classes = ds.num_classes
+    loader = build_loader(cfg, "test", ds)  # single-process: no DistributedSampler
+
+    ckpt_path = ckpt_path or os.path.join(cfg.out_dir(), "best.pth")
+    if os.path.isfile(ckpt_path):
+        ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
+        state = ckpt.get("model", ckpt)
+        model.load_state_dict(state)
+        print(f"[eval] loaded {ckpt_path}")
+    else:
+        print(f"[eval][warn] checkpoint not found: {ckpt_path} (evaluating current weights)")
+
+    model = model.to(device).eval()
+    use_amp = cfg.amp in ("bf16", "fp16")
+    amp_dtype = torch.bfloat16 if cfg.amp == "bf16" else torch.float16
+
+    records = []
+    vis_dir = os.path.join(cfg.out_dir(), "vis")
+    if cfg.visualize:
+        os.makedirs(vis_dir, exist_ok=True)
+    saved = 0
+
+    for batch in loader:
+        img = batch["image"].to(device, non_blocking=True)
+        msk = batch["mask"].numpy()
+        names = batch["name"]
+        with torch.autocast("cuda", dtype=amp_dtype, enabled=use_amp):
+            logits = model(img)
+        pred = logits.argmax(1).cpu().numpy()
+        for i in range(pred.shape[0]):
+            records.append(per_image_metrics(
+                pred[i], msk[i], num_classes,
+                include_background=cfg.include_background,
+                compute_hd95=cfg.compute_hd95))
+            if cfg.visualize and saved < cfg.vis_max:
+                save_overlay(img[i].cpu(), msk[i], pred[i], num_classes,
+                             os.path.join(vis_dir, f"{names[i]}.png"))
+                saved += 1
+
+    agg = aggregate(records)
+    out = {
+        "dataset": cfg.dataset, "protocol": cfg.protocol, "arch": cfg.arch,
+        "seed": cfg.seed, "num_classes": num_classes,
+        "metrics": agg,
+        "per_image": records,
+    }
+    out_path = os.path.join(cfg.out_dir(), "metrics.json")
+    with open(out_path, "w") as f:
+        json.dump(out, f, indent=2)
+    print(f"[eval] dice={agg['dice_mean']:.4f}+-{agg['dice_std']:.4f} "
+          f"iou={agg['iou_mean']:.4f} hd95={agg['hd95_mean']:.3f} assd={agg['assd_mean']:.3f} "
+          f"sens={agg['sensitivity_mean']:.4f} spec={agg['specificity_mean']:.4f} "
+          f"prec={agg['precision_mean']:.4f} -> {out_path}")
+    return out

code/framework/engine/losses.py

@@ -0,0 +1,33 @@
+"""Loss builder. Everything is treated as MULTICLASS (binary == 2 classes),
+which sidesteps the binary/multiclass mode pitfall and unifies all datasets.
+
+  ce_dice : CrossEntropy + multiclass Dice  (default, robust for medical seg)
+  ce      : CrossEntropy only
+  dice    : multiclass Dice only
+
+Inputs: logits [B,C,H,W], target [B,H,W] (long, ids 0..C-1).
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import segmentation_models_pytorch as smp
+
+
+class CEDiceLoss(nn.Module):
+    def __init__(self, mode: str = "ce_dice"):
+        super().__init__()
+        self.mode = mode
+        self.ce = nn.CrossEntropyLoss()
+        self.dice = smp.losses.DiceLoss(mode=smp.losses.MULTICLASS_MODE, from_logits=True)
+
+    def forward(self, logits: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+        if self.mode == "ce":
+            return self.ce(logits, target)
+        if self.mode == "dice":
+            return self.dice(logits, target)
+        return self.ce(logits, target) + self.dice(logits, target)
+
+
+def build_loss(name: str = "ce_dice") -> nn.Module:
+    return CEDiceLoss(name)

code/framework/engine/trainer.py

@@ -0,0 +1,189 @@
+"""DDP + AMP training loop.
+
+  * DDP: launched via torchrun; uses DistributedSampler. Single-GPU also works.
+  * AMP: bf16 (A100+) / fp16 (V100, with GradScaler) / fp32.
+  * Best checkpoint chosen by mean foreground Dice on the val split.
+"""
+from __future__ import annotations
+
+import os
+import math
+import json
+import time
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.nn.parallel import DistributedDataParallel as DDP
+
+from .distributed import (is_dist, is_main, get_rank, get_world_size,
+                          all_gather_object, print_main)
+from .losses import build_loss
+from ..metrics.metrics import per_image_metrics
+from ..data.loaders import build_dataset, build_loader
+
+
+_AMP_DTYPE = {"bf16": torch.bfloat16, "fp16": torch.float16}
+
+
+def build_optimizer(cfg, params):
+    if cfg.optimizer == "sgd":
+        return torch.optim.SGD(params, lr=cfg.lr, momentum=0.9,
+                               weight_decay=cfg.weight_decay, nesterov=True)
+    return torch.optim.AdamW(params, lr=cfg.lr, weight_decay=cfg.weight_decay)
+
+
+def lr_at(cfg, epoch: int) -> float:
+    if epoch < cfg.warmup_epochs:
+        return cfg.lr * (epoch + 1) / max(1, cfg.warmup_epochs)
+    e = epoch - cfg.warmup_epochs
+    total = max(1, cfg.epochs - cfg.warmup_epochs)
+    if cfg.scheduler == "poly":
+        return cfg.lr * (1 - e / total) ** 0.9
+    if cfg.scheduler == "cosine":
+        return cfg.lr * 0.5 * (1 + math.cos(math.pi * e / total))
+    return cfg.lr
+
+
+class Trainer:
+    def __init__(self, cfg, model: nn.Module, local_rank: int):
+        self.cfg = cfg
+        self.device = torch.device("cuda", local_rank) if torch.cuda.is_available() else torch.device("cpu")
+        self.local_rank = local_rank
+
+        self.train_ds = build_dataset(cfg, "train")
+        self.val_ds = build_dataset(cfg, "val")
+        self.num_classes = self.train_ds.num_classes
+        self.train_loader = build_loader(cfg, "train", self.train_ds)
+        self.val_loader = build_loader(cfg, "val", self.val_ds)
+
+        self.model = model.to(self.device)
+        if is_dist():
+            self.model = DDP(self.model, device_ids=[local_rank], output_device=local_rank,
+                             find_unused_parameters=False)
+
+        self.criterion = build_loss(cfg.loss).to(self.device)
+        self.optimizer = build_optimizer(cfg, self.model.parameters())
+        self.amp = cfg.amp
+        self.use_amp = self.amp in _AMP_DTYPE
+        self.scaler = torch.amp.GradScaler("cuda", enabled=(self.amp == "fp16"))
+        self.best = -1.0
+        self.start_epoch = 0
+        self.out_dir = cfg.out_dir()
+        if is_main():
+            os.makedirs(self.out_dir, exist_ok=True)
+            cfg.to_yaml(os.path.join(self.out_dir, "config.yaml"))
+        if cfg.resume:
+            self._load(cfg.resume)
+
+    # ---- checkpoint ----
+    def _bare(self):
+        return self.model.module if is_dist() else self.model
+
+    def _save(self, name: str, epoch: int):
+        if not is_main():
+            return
+        torch.save({
+            "epoch": epoch,
+            "model": self._bare().state_dict(),
+            "optimizer": self.optimizer.state_dict(),
+            "best": self.best,
+            "num_classes": self.num_classes,
+            "config": self.cfg.__dict__,
+        }, os.path.join(self.out_dir, name))
+
+    def _load(self, path: str):
+        ckpt = torch.load(path, map_location="cpu", weights_only=False)
+        self._bare().load_state_dict(ckpt["model"])
+        if "optimizer" in ckpt:
+            self.optimizer.load_state_dict(ckpt["optimizer"])
+        self.best = ckpt.get("best", -1.0)
+        self.start_epoch = ckpt.get("epoch", -1) + 1
+        print_main(f"[resume] from {path} at epoch {self.start_epoch}")
+
+    # ---- loops ----
+    def _autocast(self):
+        if self.use_amp:
+            return torch.autocast("cuda", dtype=_AMP_DTYPE[self.amp])
+        return torch.autocast("cuda", enabled=False)
+
+    def train_one_epoch(self, epoch: int):
+        self.model.train()
+        if is_dist():
+            self.train_loader.sampler.set_epoch(epoch)
+        for g in self.optimizer.param_groups:
+            g["lr"] = lr_at(self.cfg, epoch)
+
+        running, n = 0.0, 0
+        t0 = time.time()
+        for it, batch in enumerate(self.train_loader):
+            img = batch["image"].to(self.device, non_blocking=True)
+            msk = batch["mask"].to(self.device, non_blocking=True)
+            self.optimizer.zero_grad(set_to_none=True)
+            with self._autocast():
+                logits = self.model(img)
+                loss = self.criterion(logits, msk)
+            if self.amp == "fp16":
+                self.scaler.scale(loss).backward()
+                if self.cfg.grad_clip > 0:
+                    self.scaler.unscale_(self.optimizer)
+                    nn.utils.clip_grad_norm_(self.model.parameters(), self.cfg.grad_clip)
+                self.scaler.step(self.optimizer)
+                self.scaler.update()
+            else:
+                loss.backward()
+                if self.cfg.grad_clip > 0:
+                    nn.utils.clip_grad_norm_(self.model.parameters(), self.cfg.grad_clip)
+                self.optimizer.step()
+            running += loss.item() * img.size(0)
+            n += img.size(0)
+        if is_main():
+            print_main(f"[ep {epoch:03d}] loss={running/max(1,n):.4f} "
+                       f"lr={self.optimizer.param_groups[0]['lr']:.2e} "
+                       f"({time.time()-t0:.1f}s)")
+
+    @torch.no_grad()
+    def validate(self) -> float:
+        self.model.eval()
+        records = []
+        for batch in self.val_loader:
+            img = batch["image"].to(self.device, non_blocking=True)
+            msk = batch["mask"].numpy()
+            with self._autocast():
+                logits = self.model(img)
+            pred = logits.argmax(1).cpu().numpy()
+            for i in range(pred.shape[0]):
+                records.append(per_image_metrics(
+                    pred[i], msk[i], self.num_classes,
+                    include_background=self.cfg.include_background,
+                    compute_hd95=False))
+        gathered = all_gather_object(records)
+        flat = [r for part in gathered for r in part]
+        dices = np.array([r["dice"] for r in flat], dtype=np.float64)
+        dices = dices[~np.isnan(dices)]
+        return float(dices.mean()) if dices.size else 0.0
+
+    def fit(self):
+        best_epoch = self.start_epoch - 1
+        for epoch in range(self.start_epoch, self.cfg.epochs):
+            self.train_one_epoch(epoch)
+            do_val = ((epoch + 1) % self.cfg.val_interval == 0) or (epoch + 1 == self.cfg.epochs)
+            if do_val:
+                dice = self.validate()
+                if dice > self.best:
+                    self.best = dice
+                    best_epoch = epoch
+                    self._save("best.pth", epoch)
+                print_main(f"[ep {epoch:03d}] val_dice={dice:.4f} "
+                           f"(best={self.best:.4f} @ep{best_epoch})")
+                # early stopping: stop if val Dice hasn't improved for `patience` epochs
+                if (self.cfg.patience > 0 and (epoch + 1) >= self.cfg.min_epochs
+                        and (epoch - best_epoch) >= self.cfg.patience):
+                    print_main(f"[early-stop] no val improvement for {epoch - best_epoch} epochs "
+                               f"(patience={self.cfg.patience}); best={self.best:.4f} @ep{best_epoch}")
+                    self._save("last.pth", epoch)
+                    break
+            if self.cfg.save_interval and (epoch + 1) % self.cfg.save_interval == 0:
+                self._save(f"epoch{epoch+1}.pth", epoch)
+            self._save("last.pth", epoch)
+        print_main(f"[done] best val_dice={self.best:.4f} @ep{best_epoch} -> {self.out_dir}/best.pth")

code/framework/eval_at_res.py

@@ -0,0 +1,92 @@
+"""Re-score an existing framework checkpoint at a COMMON evaluation resolution R.
+
+Needed for resolution-fair comparison: a model that takes a fixed input size
+(SwinUNet=224, TransUNet=256) is run at its native input size, but its prediction
+and the GROUND TRUTH (loaded at native, not the 256-degraded dataloader copy) are
+both resized to R, and metrics are computed at R — matching how the conv methods
+(trained at R) and nnU-Net/U-Mamba (re-scored with --eval_size R) are evaluated.
+
+  python framework/eval_at_res.py --data_root <root> --dataset fives --protocol official \
+     --arch swinunet --seed 0 --eval_size 768 --exp_name baselines
+"""
+from __future__ import annotations
+
+import os
+import sys
+import json
+import argparse
+
+import numpy as np
+import cv2
+import torch
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+from framework.models.registry import build_model, required_img_size
+from framework.metrics.metrics import per_image_metrics, aggregate
+from framework.data.unified_dataset import UnifiedSegDataset
+from framework.data.transforms import build_transform
+
+
+def _to_R(arr, R):
+    return cv2.resize(arr.astype(np.uint8), (R, R), interpolation=cv2.INTER_NEAREST).astype(np.int64)
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--data_root", required=True)
+    ap.add_argument("--dataset", required=True)
+    ap.add_argument("--protocol", required=True)
+    ap.add_argument("--arch", required=True)
+    ap.add_argument("--encoder", default="resnet50")
+    ap.add_argument("--seed", type=int, required=True)
+    ap.add_argument("--eval_size", type=int, required=True, help="common resolution R")
+    ap.add_argument("--exp_name", default="baselines")
+    ap.add_argument("--out_root", default="results")
+    ap.add_argument("--normalize", default="auto")
+    args = ap.parse_args()
+
+    R = args.eval_size
+    model_res = required_img_size(args.arch) or R          # SwinUNet 224 / TransUNet 256 / conv -> R
+
+    ds = UnifiedSegDataset(args.data_root, args.dataset, args.protocol, "test", transform=None)
+    ds.transform = build_transform(model_res, ds.in_channels, train=False,
+                                   aug="none", normalize=args.normalize)
+    num_classes = ds.num_classes
+
+    out_dir = os.path.join(args.out_root, args.exp_name,
+                           f"{args.dataset}_{args.protocol}", args.arch, f"seed{args.seed}")
+    ckpt_path = os.path.join(out_dir, "best.pth")
+    if not os.path.isfile(ckpt_path):
+        raise SystemExit(f"checkpoint not found: {ckpt_path}")
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = build_model(args.arch, in_channels=ds.in_channels, num_classes=num_classes,
+                        img_size=model_res, encoder=args.encoder,
+                        encoder_weights="none", pretrained_ckpt="")
+    ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
+    model.load_state_dict(ckpt.get("model", ckpt))
+    model = model.to(device).eval()
+
+    records = []
+    with torch.no_grad():
+        for idx in range(len(ds)):
+            item = ds[idx]
+            img = item["image"].unsqueeze(0).to(device)         # 1,C,model_res,model_res
+            pred = model(img).argmax(1)[0].cpu().numpy()        # model_res x model_res
+            gt = cv2.imread(ds.pairs[idx][1], cv2.IMREAD_GRAYSCALE)  # native H x W, values 0..C-1
+            records.append(per_image_metrics(_to_R(pred, R), _to_R(gt, R), num_classes,
+                                             include_background=False, compute_hd95=True))
+
+    agg = aggregate(records)
+    out = {"dataset": args.dataset, "protocol": args.protocol, "arch": args.arch,
+           "seed": args.seed, "num_classes": num_classes,
+           "eval_size": R, "metrics": agg, "per_image": records}
+    os.makedirs(out_dir, exist_ok=True)
+    with open(os.path.join(out_dir, "metrics.json"), "w") as f:
+        json.dump(out, f, indent=2)
+    print(f"[eval_at_res] {args.dataset}/{args.protocol} {args.arch} seed{args.seed} @R={R}: "
+          f"n={len(records)} dice={agg['dice_mean']:.4f} hd95={agg['hd95_mean']:.2f} -> {out_dir}")
+
+
+if __name__ == "__main__":
+    main()

code/framework/eval_ckpt.py

@@ -0,0 +1,120 @@
+"""Evaluate a trained framework checkpoint on a chosen dataset/split, with the full
+metric set INCLUDING boundary metrics. Two uses, one script (no retraining):
+
+  * C7 (boundary fidelity, in-domain): point --eval_dataset at the same dataset to
+    recompute boundary-Dice / NSD from a saved best.pth.
+  * C4 (cross-center): set --eval_dataset/--eval_protocol to a DIFFERENT but
+    label-compatible dataset (e.g. train busi? no — train cvc_clinicdb, eval
+    kvasir_seg; both binary polyp). num_classes/in_channels must match.
+
+Run from project root (…/NPJ), env seggen:
+  CUDA_VISIBLE_DEVICES=5 python -m framework.eval_ckpt \
+      --ckpt results/baselines/cvc_clinicdb_official/unet/seed0/best.pth \
+      --arch unet --encoder resnet50 \
+      --data_root $DR --dataset cvc_clinicdb --protocol official \
+      --eval_dataset kvasir_seg --eval_protocol official \
+      --out_json results/crosscenter/cvc2kvasir_unet_seed0.json
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import sys
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+
+import numpy as np
+import torch
+from torch.utils.data import DataLoader
+
+from framework.config import Config  # noqa: F401  (kept for parity / future YAML use)
+from framework.models.registry import build_model, required_img_size
+from framework.data.unified_dataset import UnifiedSegDataset
+from framework.data.transforms import build_transform
+from framework.metrics.metrics import per_image_metrics
+from framework.metrics.boundary import boundary_metrics
+
+
+def get_args():
+    p = argparse.ArgumentParser("Evaluate a checkpoint (in-domain or cross-center) + boundary metrics")
+    p.add_argument("--ckpt", required=True)
+    p.add_argument("--arch", default="unet")
+    p.add_argument("--encoder", default="resnet50")
+    p.add_argument("--data_root", required=True)
+    p.add_argument("--dataset", required=True, help="dataset the ckpt was TRAINED on (sets in_ch/num_classes)")
+    p.add_argument("--protocol", required=True)
+    p.add_argument("--eval_dataset", default="", help="dataset to evaluate ON (default = --dataset)")
+    p.add_argument("--eval_protocol", default="", help="default = --protocol")
+    p.add_argument("--split", default="test")
+    p.add_argument("--img_size", type=int, default=256)
+    p.add_argument("--normalize", default="auto")
+    p.add_argument("--tol", type=float, default=2.0)
+    p.add_argument("--batch_size", type=int, default=16)
+    p.add_argument("--num_workers", type=int, default=6)
+    p.add_argument("--out_json", default="")
+    return p.parse_args()
+
+
+def main():
+    a = get_args()
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    img_size = required_img_size(a.arch) or a.img_size
+    eval_ds_name = a.eval_dataset or a.dataset
+    eval_proto = a.eval_protocol or a.protocol
+
+    # in_ch / num_classes are fixed by the TRAIN dataset (how the model was built)
+    train_probe = UnifiedSegDataset(a.data_root, a.dataset, a.protocol, "test", transform=None)
+    in_ch, n_cls = train_probe.in_channels, train_probe.num_classes
+
+    model = build_model(a.arch, in_channels=in_ch, num_classes=n_cls, img_size=img_size,
+                        encoder=a.encoder, encoder_weights="none", pretrained_ckpt="")
+    ckpt = torch.load(a.ckpt, map_location="cpu", weights_only=False)
+    model.load_state_dict(ckpt["model"])
+    model.to(device).eval()
+
+    tf = build_transform(img_size, in_ch, train=False, aug="none", normalize=a.normalize)
+    ds = UnifiedSegDataset(a.data_root, eval_ds_name, eval_proto, a.split,
+                           transform=tf, in_channels=in_ch, num_classes=n_cls)
+    if ds.num_classes != n_cls:
+        raise ValueError(f"label mismatch: train num_classes={n_cls} vs eval={ds.num_classes} "
+                         f"({a.dataset}->{eval_ds_name}) — not label-compatible for cross-center.")
+    cross = (eval_ds_name != a.dataset) or (eval_proto != a.protocol)
+    print(f"[eval] ckpt={os.path.basename(a.ckpt)} train={a.dataset}/{a.protocol} "
+          f"eval={eval_ds_name}/{eval_proto}/{a.split} cross_center={cross} "
+          f"in_ch={in_ch} num_classes={n_cls} n={len(ds)}", flush=True)
+
+    loader = DataLoader(ds, batch_size=a.batch_size, shuffle=False, num_workers=a.num_workers)
+    recs = []
+    for batch in loader:
+        img = batch["image"].to(device, non_blocking=True)
+        msk = batch["mask"].numpy()
+        with torch.no_grad():
+            pred = model(img).argmax(1).cpu().numpy()
+        for i in range(pred.shape[0]):
+            m = per_image_metrics(pred[i], msk[i], n_cls,
+                                  include_background=False, compute_hd95=True)
+            b = boundary_metrics(pred[i], msk[i], n_cls, tol=a.tol)
+            m.update(b)
+            recs.append(m)
+
+    keys = [k for k, val in recs[0].items() if isinstance(val, (int, float))]  # skip per_class dict
+    summary = {}
+    for k in keys:
+        v = np.array([r[k] for r in recs], dtype=np.float64)
+        v = v[~np.isnan(v)]
+        summary[k] = {"mean": round(float(v.mean()), 4) if v.size else None,
+                      "std": round(float(v.std()), 4) if v.size else None, "n": int(v.size)}
+    out = {"ckpt": a.ckpt, "train": f"{a.dataset}/{a.protocol}",
+           "eval": f"{eval_ds_name}/{eval_proto}/{a.split}", "cross_center": cross,
+           "num_images": len(ds), "metrics": summary}
+    print(json.dumps(out, indent=2), flush=True)
+    if a.out_json:
+        os.makedirs(os.path.dirname(os.path.abspath(a.out_json)) or ".", exist_ok=True)
+        with open(a.out_json, "w") as f:
+            json.dump(out, f, indent=2)
+        print(f"[eval] wrote {a.out_json}", flush=True)
+
+
+if __name__ == "__main__":
+    main()

code/framework/metrics/boundary.py

@@ -0,0 +1,68 @@
+"""Boundary-localized segmentation metrics (for the C7 boundary-fidelity analysis).
+
+Two standard measures, computed per foreground class then averaged:
+  * Normalized Surface Dice (NSD) @ tol: fraction of pred/gt surface points within
+    `tol` pixels of the other surface (Nikolov et al.); the medical-standard
+    boundary metric.
+  * Boundary-Dice @ tol: Dice between the tol-dilated pred and gt boundaries.
+
+Recomputable post-hoc from saved predictions/checkpoints, so adding it never
+requires retraining (see framework/eval_ckpt.py).
+"""
+from __future__ import annotations
+
+import numpy as np
+
+try:
+    from scipy.ndimage import binary_erosion, binary_dilation, distance_transform_edt
+    _HAVE_SCIPY = True
+except Exception:  # pragma: no cover
+    _HAVE_SCIPY = False
+
+
+def _surface(bin_mask: np.ndarray) -> np.ndarray:
+    return bin_mask ^ binary_erosion(bin_mask)
+
+
+def _nsd_binary(pred: np.ndarray, gt: np.ndarray, tol: float) -> float:
+    sp, sg = _surface(pred), _surface(gt)
+    ssum = sp.sum() + sg.sum()
+    if sp.sum() == 0 and sg.sum() == 0:
+        return 1.0
+    if ssum == 0:
+        return 0.0
+    dt_to_gt = distance_transform_edt(~sg)
+    dt_to_pred = distance_transform_edt(~sp)
+    pred_close = (dt_to_gt[sp] <= tol).sum()
+    gt_close = (dt_to_pred[sg] <= tol).sum()
+    return float((pred_close + gt_close) / ssum)
+
+
+def _bdice_binary(pred: np.ndarray, gt: np.ndarray, tol: int) -> float:
+    sp, sg = _surface(pred), _surface(gt)
+    if sp.sum() == 0 and sg.sum() == 0:
+        return 1.0
+    spd = binary_dilation(sp, iterations=int(tol))
+    sgd = binary_dilation(sg, iterations=int(tol))
+    denom = spd.sum() + sgd.sum()
+    if denom == 0:
+        return 0.0
+    return float(2.0 * (spd & sgd).sum() / denom)
+
+
+def boundary_metrics(pred: np.ndarray, gt: np.ndarray, num_classes: int,
+                     tol: float = 2.0) -> dict:
+    """Mean over foreground classes (1..num_classes-1). NaN if scipy missing."""
+    if not _HAVE_SCIPY:
+        return {"nsd": float("nan"), "boundary_dice": float("nan")}
+    nsds, bdices = [], []
+    for c in range(1, num_classes):
+        p, g = (pred == c), (gt == c)
+        if g.sum() == 0 and p.sum() == 0:
+            continue
+        nsds.append(_nsd_binary(p, g, tol))
+        bdices.append(_bdice_binary(p, g, int(round(tol))))
+    return {
+        "nsd": float(np.mean(nsds)) if nsds else float("nan"),
+        "boundary_dice": float(np.mean(bdices)) if bdices else float("nan"),
+    }

code/framework/metrics/metrics.py

@@ -0,0 +1,128 @@
+"""Unified segmentation metrics.
+
+Per-image (per-case), foreground-class macro-averaged unless noted:
+  * Dice (DSC)            overlap (headline)
+  * IoU (Jaccard)         overlap
+  * HD95                  95th-percentile Hausdorff distance (boundary)
+  * ASSD                  average symmetric surface distance (boundary)
+  * Sensitivity / Recall  TP/(TP+FN)
+  * Specificity           TN/(TN+FP)   (one-vs-rest, pixel-level)
+  * Precision             TP/(TP+FP)
+
+Convention: masks are integer maps 0..C-1 (0 = background); binary == 2 classes.
+Per-class values are also recorded (per_class[c]) for per-class paper tables.
+Surface metrics use MONAI if available, else medpy, else NaN.
+Aggregation: per-image -> mean±SD over the test set; report/aggregate.py then
+does mean±SD over seeds.
+"""
+from __future__ import annotations
+
+from typing import Dict, List
+import warnings
+
+import numpy as np
+
+_SURF_BACKEND = None
+_OVERLAP_KEYS = ("dice", "iou", "sensitivity", "specificity", "precision")
+_SURFACE_KEYS = ("hd95", "assd")
+SCALAR_KEYS = _OVERLAP_KEYS + _SURFACE_KEYS
+
+
+def _select_surface_backend():
+    global _SURF_BACKEND
+    if _SURF_BACKEND is not None:
+        return _SURF_BACKEND
+    try:
+        from monai.metrics import compute_hausdorff_distance, compute_average_surface_distance  # noqa
+        _SURF_BACKEND = "monai"
+    except Exception:
+        try:
+            from medpy.metric.binary import hd95, assd  # noqa
+            _SURF_BACKEND = "medpy"
+        except Exception:
+            _SURF_BACKEND = "none"
+            warnings.warn("Neither MONAI nor medpy available -> HD95/ASSD will be NaN.")
+    return _SURF_BACKEND
+
+
+def _surface_binary(pred: np.ndarray, gt: np.ndarray) -> Dict[str, float]:
+    """HD95 and ASSD for one binary 2D mask pair. Handles empty masks."""
+    backend = _select_surface_backend()
+    p_any, g_any = pred.any(), gt.any()
+    if not p_any and not g_any:
+        return {"hd95": 0.0, "assd": 0.0}
+    if not p_any or not g_any:
+        return {"hd95": float("nan"), "assd": float("nan")}
+    if backend == "monai":
+        import torch
+        from monai.metrics import compute_hausdorff_distance, compute_average_surface_distance
+        p = torch.from_numpy(pred[None, None].astype(np.uint8))
+        g = torch.from_numpy(gt[None, None].astype(np.uint8))
+        hd = compute_hausdorff_distance(p, g, percentile=95).item()
+        asd = compute_average_surface_distance(p, g, symmetric=True).item()
+        return {"hd95": float(hd), "assd": float(asd)}
+    if backend == "medpy":
+        from medpy.metric.binary import hd95 as _hd95, assd as _assd
+        pb, gb = pred.astype(bool), gt.astype(bool)
+        return {"hd95": float(_hd95(pb, gb)), "assd": float(_assd(pb, gb))}
+    return {"hd95": float("nan"), "assd": float("nan")}
+
+
+def _nanmean(xs):
+    xs = [x for x in xs if not (isinstance(x, float) and np.isnan(x))]
+    return float(np.mean(xs)) if xs else float("nan")
+
+
+def per_image_metrics(pred: np.ndarray, target: np.ndarray, num_classes: int,
+                      include_background: bool = False,
+                      compute_hd95: bool = True) -> Dict[str, object]:
+    """Metrics for a single image (2D int class maps). Returns foreground-macro
+    scalars plus a per_class breakdown. Classes absent in BOTH pred and gt are
+    skipped so they don't dilute the average."""
+    start = 0 if include_background else 1
+    acc = {k: [] for k in SCALAR_KEYS}
+    per_class: Dict[str, Dict[str, float]] = {}
+
+    for c in range(start, num_classes):
+        p = pred == c
+        g = target == c
+        if not p.any() and not g.any():
+            continue
+        tp = float(np.logical_and(p, g).sum())
+        fp = float(np.logical_and(p, ~g).sum())
+        fn = float(np.logical_and(~p, g).sum())
+        tn = float(np.logical_and(~p, ~g).sum())
+        cls = {
+            "dice": (2 * tp) / (2 * tp + fp + fn + 1e-8),
+            "iou": tp / (tp + fp + fn + 1e-8),
+            "sensitivity": tp / (tp + fn + 1e-8),
+            "specificity": tn / (tn + fp + 1e-8),
+            "precision": tp / (tp + fp + 1e-8),
+        }
+        if compute_hd95:
+            cls.update(_surface_binary(p, g))
+        else:
+            cls.update({"hd95": float("nan"), "assd": float("nan")})
+        per_class[str(c)] = cls
+        for k in SCALAR_KEYS:
+            acc[k].append(cls[k])
+
+    out: Dict[str, object] = {}
+    for k in _OVERLAP_KEYS:
+        out[k] = float(np.mean(acc[k])) if acc[k] else float("nan")
+    for k in _SURFACE_KEYS:
+        out[k] = _nanmean(acc[k]) if acc[k] else float("nan")
+    out["per_class"] = per_class
+    return out
+
+
+def aggregate(records: List[Dict[str, object]]) -> Dict[str, float]:
+    """Aggregate per-image metric dicts into mean/std over the set (per metric)."""
+    out: Dict[str, float] = {}
+    for key in SCALAR_KEYS:
+        vals = np.array([r[key] for r in records], dtype=np.float64)
+        vals = vals[~np.isnan(vals)]
+        out[f"{key}_mean"] = float(vals.mean()) if vals.size else float("nan")
+        out[f"{key}_std"] = float(vals.std()) if vals.size else float("nan")
+    out["n_images"] = float(len(records))
+    return out

code/framework/models/attention_unet.py

@@ -0,0 +1,89 @@
+"""Attention U-Net (Oktay et al., 2018) — additive attention gates on skips.
+
+Reimplemented cleanly inside the framework. The original sfczekalski/attention_unet
+repo is unmaintained (2020), unlicensed, has DRIVE-specific hardcoded padding and a
+commented-out training loop, so only the architecture idea is reused here. This is
+the faithful additive gating attention (NOT SMP's scSE), configurable for arbitrary
+in_channels / num_classes.
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+
+class ConvBlock(nn.Module):
+    def __init__(self, cin, cout):
+        super().__init__()
+        self.block = nn.Sequential(
+            nn.Conv2d(cin, cout, 3, padding=1, bias=False),
+            nn.BatchNorm2d(cout), nn.ReLU(inplace=True),
+            nn.Conv2d(cout, cout, 3, padding=1, bias=False),
+            nn.BatchNorm2d(cout), nn.ReLU(inplace=True),
+        )
+
+    def forward(self, x):
+        return self.block(x)
+
+
+class UpConv(nn.Module):
+    def __init__(self, cin, cout):
+        super().__init__()
+        self.up = nn.Sequential(
+            nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True),
+            nn.Conv2d(cin, cout, 3, padding=1, bias=False),
+            nn.BatchNorm2d(cout), nn.ReLU(inplace=True),
+        )
+
+    def forward(self, x):
+        return self.up(x)
+
+
+class AttentionGate(nn.Module):
+    """Additive attention gate: gating signal g (coarser) modulates skip x."""
+    def __init__(self, f_g, f_l, f_int):
+        super().__init__()
+        self.w_g = nn.Sequential(nn.Conv2d(f_g, f_int, 1, bias=True), nn.BatchNorm2d(f_int))
+        self.w_x = nn.Sequential(nn.Conv2d(f_l, f_int, 1, bias=True), nn.BatchNorm2d(f_int))
+        self.psi = nn.Sequential(nn.Conv2d(f_int, 1, 1, bias=True), nn.BatchNorm2d(1), nn.Sigmoid())
+        self.relu = nn.ReLU(inplace=True)
+
+    def forward(self, g, x):
+        a = self.relu(self.w_g(g) + self.w_x(x))
+        a = self.psi(a)
+        return x * a
+
+
+class AttentionUNet(nn.Module):
+    def __init__(self, in_channels: int = 3, num_classes: int = 2, base: int = 64):
+        super().__init__()
+        c = [base, base * 2, base * 4, base * 8, base * 16]
+        self.pool = nn.MaxPool2d(2, 2)
+        self.e1 = ConvBlock(in_channels, c[0])
+        self.e2 = ConvBlock(c[0], c[1])
+        self.e3 = ConvBlock(c[1], c[2])
+        self.e4 = ConvBlock(c[2], c[3])
+        self.e5 = ConvBlock(c[3], c[4])
+
+        self.u5 = UpConv(c[4], c[3]); self.a5 = AttentionGate(c[3], c[3], c[2]); self.d5 = ConvBlock(c[4], c[3])
+        self.u4 = UpConv(c[3], c[2]); self.a4 = AttentionGate(c[2], c[2], c[1]); self.d4 = ConvBlock(c[3], c[2])
+        self.u3 = UpConv(c[2], c[1]); self.a3 = AttentionGate(c[1], c[1], c[0]); self.d3 = ConvBlock(c[2], c[1])
+        self.u2 = UpConv(c[1], c[0]); self.a2 = AttentionGate(c[0], c[0], c[0] // 2); self.d2 = ConvBlock(c[1], c[0])
+        self.head = nn.Conv2d(c[0], num_classes, 1)
+
+    def forward(self, x):
+        x1 = self.e1(x)
+        x2 = self.e2(self.pool(x1))
+        x3 = self.e3(self.pool(x2))
+        x4 = self.e4(self.pool(x3))
+        x5 = self.e5(self.pool(x4))
+
+        d5 = self.u5(x5); d5 = self.d5(torch.cat([self.a5(d5, x4), d5], dim=1))
+        d4 = self.u4(d5); d4 = self.d4(torch.cat([self.a4(d4, x3), d4], dim=1))
+        d3 = self.u3(d4); d3 = self.d3(torch.cat([self.a3(d3, x2), d3], dim=1))
+        d2 = self.u2(d3); d2 = self.d2(torch.cat([self.a2(d2, x1), d2], dim=1))
+        return self.head(d2)
+
+
+def build_attention_unet(in_channels: int, num_classes: int, **_):
+    return AttentionUNet(in_channels=in_channels, num_classes=num_classes)

code/framework/models/registry.py

@@ -0,0 +1,46 @@
+"""Model registry: one entry point to build any in-framework segmenter.
+
+    build_model("unet", in_channels=3, num_classes=2, ...)
+
+Architectures:
+  * SMP zoo: unet, unetpp, manet, linknet, fpn, pspnet, deeplabv3, deeplabv3plus, pan
+  * attention_unet : reimplemented Oktay attention U-Net
+  * transunet      : sota/TransUNet model def (img_size divisible by 16; 224 canonical)
+  * swinunet       : sota/Swin-Unet model def (img_size must be 224)
+
+nnU-Net, U-Mamba (separate CLIs) and SAM-family (dropped) are intentionally NOT here.
+"""
+from __future__ import annotations
+
+import torch.nn as nn
+
+from .smp_models import is_smp_arch, build_smp
+from .attention_unet import build_attention_unet
+from .transunet_wrap import build_transunet
+from .swinunet_wrap import build_swinunet
+
+_FIXED_INPUT = {"swinunet": 224}   # archs that demand a specific img_size
+
+
+def required_img_size(arch: str):
+    return _FIXED_INPUT.get(arch.lower())
+
+
+def build_model(arch: str, in_channels: int, num_classes: int, img_size: int = 256,
+                encoder: str = "resnet34", encoder_weights: str = "imagenet",
+                pretrained_ckpt: str = "") -> nn.Module:
+    a = arch.lower()
+    if is_smp_arch(a):
+        return build_smp(a, in_channels, num_classes,
+                         encoder=encoder, encoder_weights=encoder_weights)
+    if a == "attention_unet":
+        return build_attention_unet(in_channels, num_classes)
+    if a == "transunet":
+        return build_transunet(in_channels, num_classes, img_size=img_size,
+                               encoder=encoder, pretrained_ckpt=pretrained_ckpt)
+    if a == "swinunet":
+        return build_swinunet(in_channels, num_classes, img_size=img_size,
+                              pretrained_ckpt=pretrained_ckpt)
+    raise ValueError(f"unknown arch '{arch}'. "
+                     f"SMP: unet/unetpp/manet/linknet/fpn/pspnet/deeplabv3/deeplabv3plus/pan; "
+                     f"plus attention_unet/transunet/swinunet.")

code/framework/models/smp_models.py

@@ -0,0 +1,39 @@
+"""segmentation_models.pytorch (SMP) backbones.
+
+SMP gives us a large architecture x encoder zoo behind one constructor and is a
+plain nn.Module (DDP- and bf16-friendly). Grayscale (in_channels=1) with ImageNet
+encoder weights is handled by SMP's built-in channel adaptation.
+"""
+from __future__ import annotations
+
+import segmentation_models_pytorch as smp
+
+# our arch name -> SMP architecture key
+_SMP_ARCH = {
+    "unet": "Unet",
+    "unetpp": "UnetPlusPlus",
+    "unetplusplus": "UnetPlusPlus",
+    "manet": "MAnet",
+    "linknet": "Linknet",
+    "fpn": "FPN",
+    "pspnet": "PSPNet",
+    "deeplabv3": "DeepLabV3",
+    "deeplabv3plus": "DeepLabV3Plus",
+    "pan": "PAN",
+}
+
+
+def is_smp_arch(arch: str) -> bool:
+    return arch.lower() in _SMP_ARCH
+
+
+def build_smp(arch: str, in_channels: int, num_classes: int,
+              encoder: str = "resnet34", encoder_weights: str = "imagenet", **_):
+    weights = None if encoder_weights in ("", "none", None) else encoder_weights
+    return smp.create_model(
+        arch=_SMP_ARCH[arch.lower()],
+        encoder_name=encoder,
+        encoder_weights=weights,
+        in_channels=in_channels,
+        classes=num_classes,
+    )

code/framework/models/swinunet_wrap.py

@@ -0,0 +1,75 @@
+"""Swin-Unet wrapper.
+
+Instantiates SwinTransformerSys from sota/Swin-Unet directly (bypassing the repo's
+yacs config + .npz/.h5 Synapse pipeline). Handles grayscale by repeating 1->3.
+
+Constraint: the windowed attention requires img_size = 224 (patches resolutions
+56/28/14/7 are each divisible by window_size=7). Other sizes will assert; we keep
+img_size=224 for this backbone.
+"""
+from __future__ import annotations
+
+import os
+import sys
+
+import torch
+import torch.nn as nn
+
+_REPO = os.path.join(os.path.dirname(__file__), "..", "..", "sota", "Swin-Unet")
+_REPO = os.path.abspath(_REPO)
+
+
+def _ensure_path():
+    if _REPO not in sys.path:
+        sys.path.insert(0, _REPO)
+
+
+class SwinUnetWrapper(nn.Module):
+    def __init__(self, in_channels: int, num_classes: int, img_size: int = 224,
+                 pretrained_ckpt: str = ""):
+        super().__init__()
+        _ensure_path()
+        from networks.swin_transformer_unet_skip_expand_decoder_sys import SwinTransformerSys
+        if img_size != 224:
+            raise ValueError("Swin-Unet backbone requires img_size=224.")
+        self.net = SwinTransformerSys(
+            img_size=img_size, patch_size=4, in_chans=3, num_classes=num_classes,
+            embed_dim=96, depths=[2, 2, 2, 2], num_heads=[3, 6, 12, 24],
+            window_size=7, mlp_ratio=4.0, qkv_bias=True, drop_path_rate=0.1,
+            ape=False, patch_norm=True, use_checkpoint=False,
+        )
+        if pretrained_ckpt and os.path.isfile(pretrained_ckpt):
+            self._load_pretrained(pretrained_ckpt)
+
+    def _load_pretrained(self, path):
+        """Port of Swin-Unet's load_from: load the ImageNet Swin-T encoder AND
+        mirror its `layers.X` weights into the decoder `layers_up.(3-X)` (the
+        scheme the paper uses to initialize the symmetric decoder)."""
+        import copy
+        ckpt = torch.load(path, map_location="cpu", weights_only=False)
+        if "model" not in ckpt:
+            self.net.load_state_dict(ckpt, strict=False)
+            return
+        pretrained = ckpt["model"]
+        model_dict = self.net.state_dict()
+        full = copy.deepcopy(pretrained)
+        for k, v in pretrained.items():
+            if "layers." in k:
+                n = 3 - int(k[7:8])
+                full["layers_up." + str(n) + k[8:]] = v
+        for k in list(full.keys()):
+            if k in model_dict and full[k].shape != model_dict[k].shape:
+                del full[k]
+        msg = self.net.load_state_dict(full, strict=False)
+        print(f"[swinunet] loaded pretrained {path}: "
+              f"missing={len(msg.missing_keys)} unexpected={len(msg.unexpected_keys)}")
+
+    def forward(self, x):
+        if x.size(1) == 1:
+            x = x.repeat(1, 3, 1, 1)
+        return self.net(x)
+
+
+def build_swinunet(in_channels: int, num_classes: int, img_size: int = 224,
+                   pretrained_ckpt: str = "", **_):
+    return SwinUnetWrapper(in_channels, num_classes, img_size, pretrained_ckpt)

code/framework/models/transunet_wrap.py

@@ -0,0 +1,42 @@
+"""TransUNet wrapper.
+
+Reuses ONLY the model definition from sota/TransUNet (networks/), not its
+.npz/.h5 Synapse data pipeline. The model's forward already repeats 1->3 channels
+for grayscale input, so it accepts our unified RGB or grayscale tensors.
+
+Notes:
+  * img_size must be divisible by 16 (ViT patch grid). 224 is the canonical value.
+  * pretrained_ckpt should be the R50+ViT-B_16 .npz (ImageNet-21k); optional.
+"""
+from __future__ import annotations
+
+import os
+import sys
+
+_REPO = os.path.join(os.path.dirname(__file__), "..", "..", "sota", "TransUNet")
+_REPO = os.path.abspath(_REPO)
+
+
+def _ensure_path():
+    if _REPO not in sys.path:
+        sys.path.insert(0, _REPO)
+
+
+def build_transunet(in_channels: int, num_classes: int, img_size: int = 224,
+                    encoder: str = "R50-ViT-B_16", pretrained_ckpt: str = "",
+                    vit_patches_size: int = 16, **_):
+    _ensure_path()
+    import numpy as np
+    from networks.vit_seg_modeling import VisionTransformer, CONFIGS
+
+    vit_name = encoder if encoder in CONFIGS else "R50-ViT-B_16"
+    config_vit = CONFIGS[vit_name]
+    config_vit.n_classes = num_classes
+    config_vit.n_skip = 3
+    if "R50" in vit_name:
+        config_vit.patches.grid = (img_size // vit_patches_size,
+                                   img_size // vit_patches_size)
+    model = VisionTransformer(config_vit, img_size=img_size, num_classes=num_classes)
+    if pretrained_ckpt and os.path.isfile(pretrained_ckpt):
+        model.load_from(weights=np.load(pretrained_ckpt))
+    return model

code/framework/nnunet_convert.py

@@ -0,0 +1,167 @@
+"""Convert processed_unified/<dataset>/<protocol> -> nnU-Net v2 raw format.
+
+Serves BOTH nnU-Net and U-Mamba (U-Mamba is nnU-Net v2 under the hood; same data
+format, only a different trainer/env). Honors our FIXED train/val/test split:
+  * train + val images go into imagesTr/labelsTr (nnU-Net needs them together)
+  * the exact train/val partition is emitted as splits_final.json (fold 0)
+  * test goes into imagesTs/labelsTs (excluded from training; for our evaluation)
+
+Key format rules (verified against sota/nnUNet):
+  * image file: <caseid>_0000.png  (ONE file; RGB read as 3 channels, gray as 1)
+  * mask file : <caseid>.png       (uint8, values 0..C-1, NEVER 0/255)
+  * channel_names: 3 entries (R,G,B) for RGB, 1 ("grayscale") for grayscale
+  * case ids are prefixed by split so train/val never collide inside imagesTr
+
+Usage:
+  python framework/nnunet_convert.py --data_root <processed_unified> \
+      --dataset cvc_clinicdb --protocol official --nnunet_raw <nnUNet_raw> --dataset_id 1
+"""
+from __future__ import annotations
+
+import os
+import json
+import argparse
+
+import numpy as np
+import cv2
+
+import sys
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+from framework.data.unified_dataset import (
+    _read_metadata, _pair_from_manifest, _pair_by_glob,
+    detect_in_channels, detect_num_classes,
+)
+
+
+def get_pairs(data_root, dataset, protocol, split):
+    split_dir = os.path.join(data_root, dataset, protocol, split)
+    if not os.path.isdir(split_dir):
+        return []
+    manifest = os.path.join(data_root, dataset, "manifest.jsonl")
+    return _pair_from_manifest(split_dir, manifest) or _pair_by_glob(split_dir)
+
+
+def _link(src, dst):
+    if os.path.lexists(dst):
+        os.remove(dst)
+    os.symlink(os.path.abspath(src), dst)
+
+
+def _emit_image(src, dst_png, in_ch):
+    """Place image as <...>_0000.png, ensuring a CONSISTENT channel count matching
+    in_ch (nnU-Net requires it). Grayscale datasets are re-encoded to true 1-channel
+    (some sources, e.g. kits19, store a mix of 1- and 3-channel PNGs); RGB .png are
+    symlinked (fast)."""
+    if os.path.lexists(dst_png):
+        os.remove(dst_png)  # never write THROUGH an existing symlink (would corrupt source)
+    if in_ch == 1:
+        im = cv2.imread(src, cv2.IMREAD_GRAYSCALE)
+        if im is None:
+            raise IOError(f"cannot read image {src}")
+        cv2.imwrite(dst_png, im)
+    elif os.path.splitext(src)[1].lower() == ".png":
+        _link(src, dst_png)
+    else:
+        im = cv2.imread(src, cv2.IMREAD_COLOR)
+        cv2.imwrite(dst_png, cv2.cvtColor(im, cv2.COLOR_BGR2RGB))
+
+
+def _emit_mask(src, dst_png, num_classes):
+    """Place mask as <...>.png with values 0..C-1 (remap 0/255 binary etc.)."""
+    m = cv2.imread(src, cv2.IMREAD_GRAYSCALE)
+    if m is None:
+        raise IOError(f"cannot read mask {src}")
+    uniq = set(int(v) for v in np.unique(m))
+    allowed = set(range(num_classes))
+    if uniq <= allowed and os.path.splitext(src)[1].lower() == ".png":
+        _link(src, dst_png)
+        return
+    if uniq <= {0, 255} and num_classes == 2:
+        m = (m > 0).astype(np.uint8)
+    elif len(uniq) <= num_classes:
+        remap = {v: i for i, v in enumerate(sorted(uniq))}
+        m = np.vectorize(remap.get)(m).astype(np.uint8)
+    else:
+        raise ValueError(f"mask {src} has values {sorted(uniq)} outside 0..{num_classes-1}")
+    cv2.imwrite(dst_png, m)
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--data_root", required=True)
+    ap.add_argument("--dataset", required=True)
+    ap.add_argument("--protocol", required=True)
+    ap.add_argument("--nnunet_raw", required=True, help="output nnUNet_raw root")
+    ap.add_argument("--dataset_id", type=int, required=True)
+    ap.add_argument("--name", default="", help="override dataset name suffix")
+    args = ap.parse_args()
+
+    meta = _read_metadata(args.data_root, args.dataset)
+    tr = get_pairs(args.data_root, args.dataset, args.protocol, "train")
+    va = get_pairs(args.data_root, args.dataset, args.protocol, "val")
+    ts = get_pairs(args.data_root, args.dataset, args.protocol, "test")
+    if not tr:
+        raise SystemExit(f"no train pairs for {args.dataset}/{args.protocol}")
+
+    in_ch = detect_in_channels(meta, tr[0][0])
+    num_classes = detect_num_classes(meta, [p[1] for p in tr + va + ts], args.dataset)
+
+    name = args.name or f"{args.dataset}_{args.protocol}"
+    dsname = f"Dataset{args.dataset_id:03d}_{name}"
+    root = os.path.join(args.nnunet_raw, dsname)
+    for d in ("imagesTr", "labelsTr", "imagesTs", "labelsTs"):
+        os.makedirs(os.path.join(root, d), exist_ok=True)
+
+    def emit(pairs, split, img_dir, lab_dir):
+        ids = []
+        for ip, mp in pairs:
+            stem = os.path.splitext(os.path.basename(ip))[0]
+            cid = f"{split}_{stem}"
+            _emit_image(ip, os.path.join(root, img_dir, f"{cid}_0000.png"), in_ch)
+            _emit_mask(mp, os.path.join(root, lab_dir, f"{cid}.png"), num_classes)
+            ids.append(cid)
+        return ids
+
+    train_ids = emit(tr, "train", "imagesTr", "labelsTr")
+    val_ids = emit(va, "val", "imagesTr", "labelsTr")   # val also in imagesTr
+    emit(ts, "test", "imagesTs", "labelsTs")
+
+    # dataset.json
+    if in_ch == 3:
+        channel_names = {"0": "R", "1": "G", "2": "B"}
+    else:
+        channel_names = {"0": "grayscale"}
+    labels = {"background": 0}
+    for c in range(1, num_classes):
+        labels[f"label{c}"] = c
+    dataset_json = {
+        "channel_names": channel_names,
+        "labels": labels,
+        "numTraining": len(train_ids) + len(val_ids),
+        "file_ending": ".png",
+        "name": dsname,
+        "description": f"converted from processed_unified/{args.dataset}/{args.protocol}",
+    }
+    with open(os.path.join(root, "dataset.json"), "w") as f:
+        json.dump(dataset_json, f, indent=2)
+
+    # staged splits_final.json (copy into nnUNet_preprocessed/<dsname>/ AFTER preprocessing).
+    # 3 IDENTICAL folds => train folds 0/1/2 = 3 runs of the SAME fixed split
+    # (run-to-run variance for mean±SD, matching the framework's 3-seed protocol).
+    splits = [{"train": train_ids, "val": val_ids} for _ in range(3)]
+    with open(os.path.join(root, "splits_final.json"), "w") as f:
+        json.dump(splits, f, indent=2)
+
+    print(f"[ok] {dsname}: in_ch={in_ch} num_classes={num_classes} "
+          f"train={len(train_ids)} val={len(val_ids)} test={len(ts)} -> {root}")
+    print("Next:")
+    print(f"  export nnUNet_raw=$(dirname {root})")
+    print("  export nnUNet_preprocessed=<fast_dir>  nnUNet_results=<dir>")
+    print(f"  nnUNetv2_plan_and_preprocess -d {args.dataset_id} -c 2d --verify_dataset_integrity")
+    print(f"  cp {root}/splits_final.json $nnUNet_preprocessed/{dsname}/splits_final.json")
+    print(f"  nnUNetv2_train {args.dataset_id} 2d 0                      # nnU-Net")
+    print(f"  nnUNetv2_train {args.dataset_id} 2d 0 -tr nnUNetTrainerUMambaBot   # U-Mamba")
+
+
+if __name__ == "__main__":
+    main()

code/framework/nnunet_eval.py

@@ -0,0 +1,94 @@
+"""Compute framework-format metrics.json from nnU-Net / U-Mamba test predictions.
+
+nnU-Net only reports validation Dice during training. To compare on the SAME
+held-out test set with the SAME 7 metrics as the framework, we: predict on
+imagesTs (done separately via nnUNetv2_predict), then run THIS script to score the
+predicted masks against labelsTs using framework/metrics.py, writing a metrics.json
+in the exact framework format so report/aggregate.py includes nnU-Net/U-Mamba rows.
+
+  python framework/nnunet_eval.py --data_root <processed_unified> --dataset <ds> \
+     --protocol <proto> --raw <nnUNet_raw> --dataset_id <ID> --fold <f> \
+     --pred_dir <predictions> --arch nnunet --exp_name baselines
+"""
+from __future__ import annotations
+
+import os
+import sys
+import json
+import glob
+import argparse
+
+import numpy as np
+import cv2
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+from framework.metrics.metrics import per_image_metrics, aggregate
+from framework.data.unified_dataset import _read_metadata, detect_num_classes
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--data_root", required=True)
+    ap.add_argument("--dataset", required=True)
+    ap.add_argument("--protocol", required=True)
+    ap.add_argument("--raw", required=True, help="nnUNet_raw root")
+    ap.add_argument("--dataset_id", type=int, required=True)
+    ap.add_argument("--name", default="", help="Dataset<DDD>_<name> suffix; default <dataset>_<protocol>")
+    ap.add_argument("--fold", type=int, required=True)
+    ap.add_argument("--pred_dir", required=True)
+    ap.add_argument("--arch", default="nnunet")
+    ap.add_argument("--exp_name", default="baselines")
+    ap.add_argument("--out_root", default="results")
+    ap.add_argument("--include_background", action="store_true")
+    ap.add_argument("--eval_size", type=int, default=0,
+                    help="resize pred+gt to R×R (nearest) before scoring; 0 = native GT resolution")
+    args = ap.parse_args()
+
+    name = args.name or f"{args.dataset}_{args.protocol}"
+    dsname = f"Dataset{args.dataset_id:03d}_{name}"
+    lab_dir = os.path.join(args.raw, dsname, "labelsTs")
+
+    meta = _read_metadata(args.data_root, args.dataset)
+    gt_masks = sorted(glob.glob(os.path.join(lab_dir, "*.png")))
+    num_classes = detect_num_classes(meta, gt_masks, args.dataset)
+
+    records = []
+    n_missing = 0
+    for pp in sorted(glob.glob(os.path.join(args.pred_dir, "*.png"))):
+        base = os.path.basename(pp)
+        gp = os.path.join(lab_dir, base)
+        if not os.path.isfile(gp):
+            n_missing += 1
+            continue
+        pred = cv2.imread(pp, cv2.IMREAD_GRAYSCALE)
+        gt = cv2.imread(gp, cv2.IMREAD_GRAYSCALE)
+        if pred is None or gt is None:
+            n_missing += 1
+            continue
+        if pred.shape != gt.shape:
+            pred = cv2.resize(pred, (gt.shape[1], gt.shape[0]), interpolation=cv2.INTER_NEAREST)
+        if args.eval_size > 0:  # resolution-fair common-R scoring
+            R = args.eval_size
+            pred = cv2.resize(pred, (R, R), interpolation=cv2.INTER_NEAREST)
+            gt = cv2.resize(gt, (R, R), interpolation=cv2.INTER_NEAREST)
+        records.append(per_image_metrics(pred.astype(np.int64), gt.astype(np.int64),
+                                         num_classes, include_background=args.include_background,
+                                         compute_hd95=True))
+    if not records:
+        raise SystemExit(f"no matched (pred,gt) pairs in {args.pred_dir} vs {lab_dir}")
+
+    agg = aggregate(records)
+    out = {"dataset": args.dataset, "protocol": args.protocol, "arch": args.arch,
+           "seed": args.fold, "num_classes": num_classes, "metrics": agg, "per_image": records}
+    out_dir = os.path.join(args.out_root, args.exp_name, f"{args.dataset}_{args.protocol}",
+                           args.arch, f"seed{args.fold}")
+    os.makedirs(out_dir, exist_ok=True)
+    with open(os.path.join(out_dir, "metrics.json"), "w") as f:
+        json.dump(out, f, indent=2)
+    print(f"[nnunet_eval] {dsname} fold{args.fold}: n={len(records)} (missing {n_missing}) "
+          f"dice={agg['dice_mean']:.4f} iou={agg['iou_mean']:.4f} hd95={agg['hd95_mean']:.2f} "
+          f"-> {out_dir}/metrics.json")
+
+
+if __name__ == "__main__":
+    main()

code/framework/report/aggregate.py

@@ -0,0 +1,310 @@
+"""Aggregate per-seed metrics.json files into mean +- SD tables (over seeds).
+
+Scans results/<exp_name>/**/seed*/metrics.json, groups by (dataset, protocol,
+arch), and reports mean +- SD across seeds for every metric. Emits CSV, Markdown,
+LaTeX. Overlap/clf metrics shown as percentages; HD95/ASSD as raw distances.
+
+  python framework/report/aggregate.py --exp_name baselines [--out_root results]
+"""
+from __future__ import annotations
+
+import os
+import json
+import glob
+import argparse
+from collections import defaultdict
+
+import numpy as np
+
+# (key, label, is_percent, higher_is_better)
+METRICS = [
+    ("dice", "Dice", True, True),
+    ("iou", "IoU", True, True),
+    ("hd95", "HD95", False, False),
+    ("assd", "ASSD", False, False),
+    ("sensitivity", "Sens", True, True),
+    ("specificity", "Spec", True, True),
+    ("precision", "Prec", True, True),
+]
+
+
+def load_runs(out_root, exp_name):
+    runs = []
+    for path in glob.glob(os.path.join(out_root, exp_name, "**", "seed*", "metrics.json"), recursive=True):
+        try:
+            with open(path) as f:
+                runs.append(json.load(f))
+        except Exception:
+            pass
+    return runs
+
+
+# Display labels for split protocols (replaces opaque "fold01" with the real
+# reporting protocol). Datasets with an official split keep "official"/"holdout";
+# fold-1-of-k single splits are labelled "single-split"; PanNuke is reported as
+# its official k-fold cross-validation (handled in summarize via fold-merge).
+_PROTO_LABEL = {
+    ("idridd_segmentation", "fold01"): "official",    # official 54/27; train re-split for val
+    ("busi", "fold01"): "single-split",
+    ("medsegdb_kits19", "fold01"): "single-split",
+    ("pannuke_semantic", "fold01"): "single-split",   # until the other folds are added
+}
+
+# Datasets reported as official k-fold cross-validation (mean±SD OVER folds, not
+# over seeds): folds are merged into one row when >1 fold is present.
+_CV_DATASETS = {"pannuke_semantic"}
+
+
+def _proto_label(dataset, protocol):
+    return _PROTO_LABEL.get((dataset, protocol), protocol)
+
+
+def _agg_over(items, key):
+    vals = np.array([it.get("metrics", {}).get(f"{key}_mean", np.nan) for it in items], np.float64)
+    vals = vals[~np.isnan(vals)]
+    return (float(vals.mean()), float(vals.std())) if vals.size else (float("nan"), float("nan"))
+
+
+def summarize(runs):
+    # (dataset, arch) -> protocol -> [runs]
+    by_da = defaultdict(lambda: defaultdict(list))
+    for d in runs:
+        by_da[(d.get("dataset"), d.get("arch"))][d.get("protocol")].append(d)
+    rows = []
+    for (dataset, arch), proto_map in sorted(by_da.items()):
+        protos = sorted(p for p in proto_map if p is not None)
+        row = {"dataset": dataset, "arch": arch}
+        if dataset in _CV_DATASETS and len(protos) > 1:
+            # k-fold CV: collapse seeds within each fold, then mean±SD OVER folds
+            row["protocol"] = f"{len(protos)}-fold"
+            row["n_seeds"] = len(protos)
+            for key, _, _, _ in METRICS:
+                fold_means = [m for m in (_agg_over(proto_map[p], key)[0] for p in protos)
+                              if not np.isnan(m)]
+                fm = np.array(fold_means, np.float64)
+                row[f"{key}_mean"] = float(fm.mean()) if fm.size else float("nan")
+                row[f"{key}_sd"] = float(fm.std()) if fm.size else float("nan")
+        else:
+            # single fixed split: mean±SD over seeds
+            proto = protos[0] if protos else None
+            items = proto_map.get(proto, [])
+            row["protocol"] = _proto_label(dataset, proto)
+            row["n_seeds"] = len(items)
+            for key, _, _, _ in METRICS:
+                row[f"{key}_mean"], row[f"{key}_sd"] = _agg_over(items, key)
+        rows.append(row)
+    return rows
+
+
+def _fmt(row, key, pct):
+    m, s = row[f"{key}_mean"], row[f"{key}_sd"]
+    if np.isnan(m):
+        return "—"
+    return f"{m*100:.2f}±{s*100:.2f}" if pct else f"{m:.2f}±{s:.2f}"
+
+
+def to_markdown(rows):
+    cols = ["dataset", "protocol", "arch", "seeds"] + [lbl for _, lbl, _, _ in METRICS]
+    out = "| " + " | ".join(cols) + " |\n" + "|" + "---|" * len(cols) + "\n"
+    for r in rows:
+        cells = [r["dataset"], r["protocol"], r["arch"], str(r["n_seeds"])]
+        cells += [_fmt(r, k, p) for k, _, p, _ in METRICS]
+        out += "| " + " | ".join(cells) + " |\n"
+    return out
+
+
+def to_csv(rows):
+    cols = ["dataset", "protocol", "arch", "n_seeds"]
+    for k, _, _, _ in METRICS:
+        cols += [f"{k}_mean", f"{k}_sd"]
+    out = ",".join(cols) + "\n"
+    for r in rows:
+        out += ",".join(str(r[c]) for c in cols) + "\n"
+    return out
+
+
+def to_latex(rows):
+    spec = "lll c" + "c" * len(METRICS)
+    out = "\\begin{tabular}{" + spec + "}\n\\toprule\n"
+    out += "Dataset & Protocol & Method & Seeds & " + " & ".join(lbl for _, lbl, _, _ in METRICS) + " \\\\\n\\midrule\n"
+    for r in rows:
+        cells = [str(r["dataset"]), str(r["protocol"]), str(r["arch"]), str(r["n_seeds"])]
+        cells += [_fmt(r, k, p).replace("±", "$\\pm$") for k, _, p, _ in METRICS]
+        out += " & ".join(cells) + " \\\\\n"
+    out += "\\bottomrule\n\\end{tabular}\n"
+    return out
+
+
+_ARCH_ORDER = ["unet", "unetpp", "deeplabv3plus", "attention_unet", "transunet", "swinunet",
+               "nnunet", "umamba"]
+
+
+# (#, dataset, modality, target, classes, channels, image size, reported protocol, train/val/test)
+_DATASETS_INFO = [
+    ("1", "CVC-ClinicDB", "Colonoscopy (endoscopy)", "Polyp", "2", "RGB", "384×288", "official", "490 / 61 / 61"),
+    ("2", "Kvasir-SEG", "GI endoscopy", "Polyp", "2", "RGB", "~622×529 (var)", "official", "800 / 100 / 100"),
+    ("3", "FIVES", "Retinal fundus", "Vessel", "2", "RGB", "2048×2048", "official", "480 / 120 / 200"),
+    ("4", "BUSI", "Breast ultrasound", "Tumor", "2", "grayscale¹", "variable", "single-split²", "545 / 78 / 157"),
+    ("5", "REFUGE2", "Retinal fundus", "Optic disc & cup", "3", "RGB", "~2124×2056", "official", "400 / 400 / 400"),
+    ("6", "ACDC", "Cardiac MRI (2D slices)", "RV / Myo / LV", "4", "grayscale", "~240×256 (var)", "official", "136 / 210 / 380"),
+    ("7", "IDRiD", "Retinal fundus", "DR lesions (4) + optic disc", "6", "RGB", "4288×2848", "official", "43 / 11 / 27"),
+    ("8", "PanNuke", "Histopathology (H&E)", "Nuclei (5 types)", "6", "RGB", "256×256", "official 3-fold CV", "~2.7k / 2.6k / 2.6k per fold"),
+    ("9", "ISIC2018", "Dermoscopy", "Skin lesion", "2", "RGB", "256×256", "holdout", "2582 / 369 / 737"),
+    ("10", "KiTS19", "Kidney CT (2D slices)", "Kidney (binary)", "2", "grayscale¹", "256×256", "single-split²", "2832 / 479 / 705"),
+]
+
+_METHODS_INFO = [
+    ("UNet", "CNN encoder–decoder", "SMP, ResNet-50 encoder (ImageNet-pretrained)"),
+    ("UNet++", "Nested UNet", "SMP, ResNet-50 (ImageNet)"),
+    ("DeepLabV3+", "Atrous CNN", "SMP, ResNet-50 (ImageNet)"),
+    ("Attention-UNet", "Attention-gated UNet", "Re-implemented, trained from scratch"),
+    ("TransUNet", "CNN–Transformer hybrid", "R50-ViT-B/16 (ImageNet-pretrained)"),
+    ("Swin-UNet", "Pure-Transformer UNet", "Swin-Tiny (ImageNet-pretrained)"),
+    ("nnU-Net (v2)", "Self-configuring CNN", "2D config, 250 epochs"),
+    ("U-Mamba", "State-space (Mamba) UNet", "U-Mamba_Bot, 100 epochs"),
+]
+
+# (metric, formula/definition, direction, unit, chinese role)
+_METRICS_INFO = [
+    ("Dice (DSC)", "2·|P∩G| / (|P|+|G|) = 2TP / (2TP+FP+FN)", "↑ higher", "%",
+     "区域重叠度(分割主指标);对类别不平衡较鲁棒,综合反映分割整体准确度。"),
+    ("IoU (Jaccard)", "|P∩G| / |P∪G| = TP / (TP+FP+FN)", "↑ higher", "%",
+     "交并比,同样衡量区域重叠但更严格(对错误惩罚更重),常与 Dice 并列报告。"),
+    ("Sensitivity / Recall", "TP / (TP+FN)", "↑ higher", "%",
+     "召回率/敏感度:真值前景中被正确分出的比例,反映“漏检/漏分割”程度(越高漏得越少)。"),
+    ("Specificity", "TN / (TN+FP)  (pixel-level, one-vs-rest)", "↑ higher", "%",
+     "特异度:真值背景中被正确判为背景的比例,反映对背景区域的误报控制能力。"),
+    ("Precision", "TP / (TP+FP)", "↑ higher", "%",
+     "精确率:预测为前景的像素中真正属于前景的比例,反映“过分割/误报”程度(越高误分越少)。"),
+    ("HD95", "95th-percentile Hausdorff distance between predicted & GT boundaries", "↓ lower", "pixels",
+     "边界距离:预测与真值轮廓间偏差的 95% 分位 Hausdorff 距离,衡量最大边界误差(取95%分位以抑制离群点),越小边界越贴合。"),
+    ("ASSD", "Average symmetric surface distance between boundaries", "↓ lower", "pixels",
+     "平均对称表面距离:两条轮廓间双向平均距离,衡量整体边界吻合度,越小边界越准。"),
+]
+
+
+def _intro_html():
+    h = ["<div class='ov'>"]
+    h.append("<h2>Project overview</h2>")
+    h.append("<p>A unified benchmark of <b>8 2D medical-image segmentation methods</b> across "
+             "<b>10 public datasets</b> spanning <b>7 imaging modalities</b> (endoscopy, retinal fundus, "
+             "ultrasound, cardiac MRI, dermoscopy, histopathology, abdominal CT). Every method is trained "
+             "and evaluated through one pipeline (bf16 AMP on A100, identical metrics). Reported as "
+             "<b>mean±SD</b> over <b>3 seeds</b> for fixed-split datasets and over <b>folds</b> for "
+             "cross-validation datasets (PanNuke: official 3-fold). Each (dataset,method) cell aggregates "
+             "~tens–thousands of test images; the suite totals ≈20k images. Per-method efficiency "
+             "(params / FLOPs / throughput) is in <code>efficiency.md</code>.</p>")
+
+    h.append("<h2>Datasets</h2>")
+    h.append("<table><tr><th>#</th><th>Dataset</th><th>Modality</th><th>Target</th><th>Classes</th>"
+             "<th>Channels</th><th>Image size</th><th>Reported protocol</th><th>Train / Val / Test</th></tr>")
+    for r in _DATASETS_INFO:
+        h.append("<tr><td>%s</td><td class='l'>%s</td><td class='l'>%s</td><td class='l'>%s</td>"
+                 "<td>%s</td><td>%s</td><td>%s</td><td>%s</td><td>%s</td></tr>" % r)
+    h.append("</table>")
+    h.append("<div class='cap'>¹ BUSI/KiTS19 are grayscale in content but stored as 3-channel PNG — read as "
+             "grayscale. ² BUSI/KiTS19 have no canonical split, so one fixed fold (of a 5-fold partition) is "
+             "reported with 3 seeds; datasets with an official split use it. Mask labels are 0…C-1 "
+             "(0=background); multi-class metrics are macro-averaged over foreground classes.</div>")
+
+    h.append("<h2>Segmentation methods</h2>")
+    h.append("<table><tr><th>Method</th><th>Family</th><th>Backbone / setup</th></tr>")
+    for m in _METHODS_INFO:
+        h.append("<tr><td class='l'>%s</td><td class='l'>%s</td><td class='l'>%s</td></tr>" % m)
+    h.append("</table>")
+
+    h.append("<h2>Evaluation metrics</h2>")
+    h.append("<table><tr><th>Metric</th><th>Definition</th><th>Direction</th><th>Unit</th>"
+             "<th>作用 / 含义（中文）</th></tr>")
+    for m in _METRICS_INFO:
+        h.append("<tr><td class='l'>%s</td><td class='l'>%s</td><td>%s</td><td>%s</td>"
+                 "<td class='l'>%s</td></tr>" % m)
+    h.append("</table>")
+    h.append("<p style='font-size:12.5px'>P = predicted foreground, G = ground-truth foreground; "
+             "TP/FP/FN/TN are pixel counts. Overlap metrics (Dice/IoU/Sens/Spec/Prec) are exact pixel "
+             "computations; boundary metrics (HD95/ASSD) are computed with MONAI on the mask edges.</p>")
+
+    h.append("<h2>Evaluation protocol</h2>")
+    h.append("<p>For each test image, every metric is computed <b>per foreground class</b> (background "
+             "excluded) and macro-averaged over the classes present; classes absent in <i>both</i> "
+             "prediction and ground truth are skipped so they do not dilute the score. Per-image values are "
+             "averaged to one score per (dataset, method, run); we then report <b>mean±SD across runs</b> — "
+             "3 random seeds for fixed-split datasets, or across folds for CV datasets. For surface metrics, "
+             "an empty–empty (no foreground in either) pair scores 0, while empty-vs-non-empty is undefined "
+             "(NaN) and excluded from the average. All methods share the same held-out test images per "
+             "dataset, so comparisons are paired.</p>")
+    h.append("<hr style='margin:18px 0;border:none;border-top:1px solid #ddd'>")
+    h.append("</div>")
+    return "\n".join(h)
+
+
+def to_html(rows, title="SegGen baselines"):
+    cell, dslist = {}, []
+    for r in rows:
+        ds = f"{r['dataset']}/{r['protocol']}"
+        if ds not in dslist:
+            dslist.append(ds)
+        cell[(ds, r["arch"])] = r
+    archs = [a for a in _ARCH_ORDER if any(r["arch"] == a for r in rows)] or \
+            sorted({r["arch"] for r in rows})
+    h = ["<!doctype html><html><head><meta charset='utf-8'><title>%s</title><style>" % title,
+         "body{font-family:Arial,Helvetica,sans-serif;margin:24px;color:#222}",
+         "h1{font-size:20px}h2{margin-top:28px;font-size:16px;color:#0a5}",
+         "table{border-collapse:collapse;margin:6px 0 18px}",
+         "th,td{border:1px solid #ccc;padding:5px 9px;text-align:center;font-size:12.5px}",
+         "th{background:#f3f3f3}td.ds{text-align:left;font-weight:bold;background:#fafafa}",
+         "td.l{text-align:left}.ov{font-size:12.5px;line-height:1.5;max-width:1100px}",
+         "b{color:#0a6}.cap{color:#777;font-size:12px;margin:2px 0 14px}</style></head><body>"]
+    nseeds = {f"{r['dataset']}/{r['protocol']}": r["n_seeds"] for r in rows}
+    h.append(f"<h1>{title} — mean±SD</h1>")
+    h.append("<div class='cap'>Higher is better for Dice/IoU/Sens/Spec/Prec (%); "
+             "lower for HD95/ASSD. <b>Bold</b>=best (highest/lowest mean) per dataset. "
+             "Variance is over seeds for single-split datasets and over folds for k-fold (CV) datasets. "
+             "'seeds' column = #seeds (single-split) or #folds (CV).</div>")
+    h.append(_intro_html())
+    h.append("<h1 style='font-size:18px'>Results</h1>")
+    for key, label, pct, hib in METRICS:
+        h.append(f"<h2>{label}{' (%)' if pct else ' (lower=better)'}</h2>")
+        h.append("<table><tr><th>Dataset</th>" + "".join(f"<th>{a}</th>" for a in archs)
+                 + "<th>seeds</th></tr>")
+        for ds in dslist:
+            vals = {a: cell[(ds, a)][f"{key}_mean"] for a in archs
+                    if (ds, a) in cell and cell[(ds, a)][f"{key}_mean"] == cell[(ds, a)][f"{key}_mean"]}
+            best = (max if hib else min)(vals, key=vals.get) if vals else None
+            bold = {best} if best is not None else set()
+            tds = [f"<td class='ds'>{ds}</td>"]
+            for a in archs:
+                if (ds, a) in cell:
+                    txt = _fmt(cell[(ds, a)], key, pct)
+                    tds.append(f"<td>{'<b>'+txt+'</b>' if a in bold else txt}</td>")
+                else:
+                    tds.append("<td>—</td>")
+            tds.append(f"<td>{nseeds.get(ds,'')}</td>")
+            h.append("<tr>" + "".join(tds) + "</tr>")
+        h.append("</table>")
+    h.append("</body></html>")
+    return "\n".join(h)
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--exp_name", required=True)
+    p.add_argument("--out_root", default="results")
+    args = p.parse_args()
+
+    runs = load_runs(args.out_root, args.exp_name)
+    if not runs:
+        print(f"no metrics.json under {args.out_root}/{args.exp_name}")
+        return
+    rows = summarize(runs)
+    base = os.path.join(args.out_root, args.exp_name)
+    open(os.path.join(base, "summary.csv"), "w").write(to_csv(rows))
+    open(os.path.join(base, "summary.md"), "w").write(to_markdown(rows))
+    open(os.path.join(base, "summary.tex"), "w").write(to_latex(rows))
+    open(os.path.join(base, "summary.html"), "w").write(to_html(rows, title=f"SegGen baselines ({args.exp_name})"))
+    print(to_markdown(rows))
+    print(f"{len(runs)} runs -> {len(rows)} (dataset,arch) cells; written {base}/summary.{{csv,md,tex,html}}")
+
+
+if __name__ == "__main__":
+    main()

code/framework/test.py

@@ -0,0 +1,44 @@
+"""Evaluation entrypoint (single process).
+
+  python framework/test.py --dataset cvc_clinicdb --arch unet --exp_name myrun --seed 0
+"""
+from __future__ import annotations
+
+import os
+import sys
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+
+import torch
+import cv2
+
+# Single-threaded OpenCV per process (parallelism via num_workers); avoids the
+# nproc-sized cv2 thread-pool oversubscription that starves the GPU at high res.
+cv2.setNumThreads(1)
+
+from framework.config import Config
+from framework.models.registry import build_model, required_img_size
+from framework.engine.evaluator import evaluate
+from framework.data.loaders import build_dataset
+
+
+def main():
+    cfg = Config.from_args()
+    req = required_img_size(cfg.arch)
+    if req and cfg.img_size != req:
+        cfg.img_size = req
+
+    if torch.cuda.is_available():
+        torch.cuda.set_device(0)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    probe = build_dataset(cfg, "test")
+    model = build_model(cfg.arch, in_channels=probe.in_channels, num_classes=probe.num_classes,
+                        img_size=cfg.img_size, encoder=cfg.encoder,
+                        encoder_weights="none",  # weights come from checkpoint
+                        pretrained_ckpt="")
+    evaluate(cfg, model, device, ckpt_path=cfg.resume)
+
+
+if __name__ == "__main__":
+    main()

code/framework/train.py

@@ -0,0 +1,60 @@
+"""Training entrypoint.
+
+Single GPU:   python framework/train.py --dataset cvc_clinicdb --arch unet ...
+Multi-GPU :   torchrun --nproc_per_node=4 framework/train.py --dataset ... --arch ...
+"""
+from __future__ import annotations
+
+import os
+import sys
+
+# allow `python framework/train.py` (add repo root to path)
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+
+import torch
+import cv2
+
+# Each DataLoader worker single-threaded for OpenCV; parallelism comes from num_workers.
+# Without this, cv2 spawns an nproc-sized (~384) thread pool per worker, whose per-op
+# dispatch overhead starves the GPU at high resolution (768) -> ~4x slower epochs.
+cv2.setNumThreads(1)
+
+from framework.config import Config
+from framework.engine.distributed import setup_distributed, cleanup_distributed, set_seed, print_main
+from framework.models.registry import build_model, required_img_size
+from framework.engine.trainer import Trainer
+
+
+def main():
+    cfg = Config.from_args()
+
+    # some backbones require a fixed input size
+    req = required_img_size(cfg.arch)
+    if req and cfg.img_size != req:
+        print_main(f"[info] arch '{cfg.arch}' requires img_size={req}; overriding {cfg.img_size}.")
+        cfg.img_size = req
+
+    local_rank = setup_distributed()
+    set_seed(cfg.seed, rank=local_rank)
+
+    # peek dataset to get in/out channels before building the model
+    from framework.data.loaders import build_dataset
+    probe = build_dataset(cfg, "train")
+    in_ch, n_cls = probe.in_channels, probe.num_classes
+    print_main(f"[data] {cfg.dataset}/{cfg.protocol}: in_channels={in_ch} num_classes={n_cls} "
+               f"train={len(probe)}")
+
+    model = build_model(cfg.arch, in_channels=in_ch, num_classes=n_cls,
+                        img_size=cfg.img_size, encoder=cfg.encoder,
+                        encoder_weights=cfg.encoder_weights,
+                        pretrained_ckpt=cfg.pretrained_ckpt)
+    print_main(f"[model] {cfg.arch} params={sum(p.numel() for p in model.parameters())/1e6:.1f}M "
+               f"amp={cfg.amp}")
+
+    trainer = Trainer(cfg, model, local_rank)
+    trainer.fit()
+    cleanup_distributed()
+
+
+if __name__ == "__main__":
+    main()

code/framework/visualize/overlay.py

@@ -0,0 +1,53 @@
+"""Save a side-by-side overlay: input | ground-truth | prediction.
+
+Used at test time to qualitatively inspect each method's output. Denormalizes the
+input tensor back to a viewable image and color-codes class masks.
+"""
+from __future__ import annotations
+
+import numpy as np
+import cv2
+import torch
+
+_IMAGENET_MEAN = np.array([0.485, 0.456, 0.406])
+_IMAGENET_STD = np.array([0.229, 0.224, 0.225])
+
+# distinct colors for up to 6 classes (0 = background -> transparent/black)
+_PALETTE = np.array([
+    [0, 0, 0], [255, 0, 0], [0, 255, 0], [0, 0, 255],
+    [255, 255, 0], [255, 0, 255], [0, 255, 255],
+], dtype=np.uint8)
+
+
+def _denorm(img: torch.Tensor) -> np.ndarray:
+    x = img.float().numpy()                       # C,H,W
+    c = x.shape[0]
+    x = np.transpose(x, (1, 2, 0))                # H,W,C
+    if c == 3:
+        x = x * _IMAGENET_STD + _IMAGENET_MEAN
+    else:
+        x = x * 0.5 + 0.5
+        x = np.repeat(x, 3, axis=2) if x.shape[2] == 1 else x
+    x = np.clip(x * 255.0, 0, 255).astype(np.uint8)
+    return x
+
+
+def _colorize(mask: np.ndarray, num_classes: int) -> np.ndarray:
+    h, w = mask.shape
+    out = np.zeros((h, w, 3), dtype=np.uint8)
+    for c in range(1, num_classes):
+        out[mask == c] = _PALETTE[c % len(_PALETTE)]
+    return out
+
+
+def save_overlay(img: torch.Tensor, gt: np.ndarray, pred: np.ndarray,
+                 num_classes: int, path: str, alpha: float = 0.5) -> None:
+    base = _denorm(img)
+    h, w = gt.shape
+    base = cv2.resize(base, (w, h), interpolation=cv2.INTER_LINEAR)
+    gt_c = _colorize(gt, num_classes)
+    pr_c = _colorize(pred, num_classes)
+    gt_o = cv2.addWeighted(base, 1 - alpha, gt_c, alpha, 0)
+    pr_o = cv2.addWeighted(base, 1 - alpha, pr_c, alpha, 0)
+    panel = np.concatenate([base, gt_o, pr_o], axis=1)
+    cv2.imwrite(path, cv2.cvtColor(panel, cv2.COLOR_RGB2BGR))