torch_em.model.unetr

   1from os import PathLike
   2from functools import partial
   3from collections.abc import Mapping
   4from typing import Callable, Optional, Tuple, Union
   5
   6import torch
   7import torch.nn as nn
   8import torch.nn.functional as F
   9
  10from .util import VIT_CONFIGS
  11from .vit import get_vision_transformer
  12from .unet import Decoder, ConvBlock2d, ConvBlock3d, Upsampler2d, Upsampler3d, _update_conv_kwargs
  13
  14
  15#
  16# UNETR WITH PRETRAINED ENCODERS
  17#
  18
  19
  20def get_unetr_model(
  21    ndim: int,
  22    backbone: str,
  23    model_type: Union[str, nn.Module],
  24    encoder_checkpoint: Optional[Union[str, PathLike, Mapping]] = None,
  25    init_decoder: Optional[Callable[[nn.Module], nn.Module]] = None,
  26    **init_kwargs,
  27) -> nn.Module:
  28    """Create a UNETR with a ViT encoder for 2D or 3D segmentation.
  29
  30    Pass `encoder_checkpoint` to load encoder weights.
  31    Torchvision loads its default pretrained weights without a checkpoint.
  32    Set `pretrained=False` to disable this torchvision download.
  33    3D models apply the encoder to each slice and use a 3D convolutional decoder.
  34
  35    Args:
  36        ndim: The number of spatial dimensions. Use 2 or 3.
  37        backbone: The ViT family supported by `get_vision_transformer`: "sam", "sam2", "sam3", "cellpose_sam",
  38            "mae", "scalemae", "dinov2", "dinov3", or "torchvision".
  39            Use a custom name for an injected encoder from another family.
  40        model_type: The encoder name or a module from `get_vision_transformer`. Examples: "hvit_t" for SAM2,
  41            "vit_b_reg4" for DINOv2, or "vit_b_16" for torchvision. SAM weight names can include training suffixes.
  42            Custom modules must expose `img_size`, `embed_dim`, and `in_chans`.
  43            They return feature maps or a tuple of feature maps and intermediate features.
  44        encoder_checkpoint: The encoder checkpoint, as a file path or a loaded state dictionary.
  45        init_decoder: The function that initializes the decoder with pretrained weights. Use this only for 2D models.
  46        init_kwargs: The keyword arguments that override the defaults for UNETR2D or UNETR3D.
  47            The factory selects the image size and normalization for the backbone.
  48            It disables skip connections when the encoder features do not match the decoder.
  49
  50    Returns:
  51        The UNETR model.
  52    """
  53    if ndim not in (2, 3):
  54        raise ValueError(f"Invalid ndim={ndim}. Use 2 or 3.")
  55    if init_decoder is not None and ndim != 2:
  56        raise ValueError(f"`init_decoder` is only supported for 2D models, got ndim={ndim}.")
  57
  58    config = VIT_CONFIGS.get(backbone)
  59    if config is None and not isinstance(model_type, nn.Module):
  60        raise ValueError(f"Unsupported backbone '{backbone}'.")
  61
  62    # Get the model class.
  63    model_class = UNETR2D if ndim == 2 else UNETR3D
  64    if isinstance(model_type, str):
  65        encoder = "_".join(model_type.split("_")[:2]) if config.weight_suffixes else model_type
  66        img_size = config.img_size
  67        encoder_name = encoder
  68    elif isinstance(model_type, nn.Module):
  69        encoder = model_type
  70        img_size = encoder.img_size
  71        encoder_name = getattr(encoder, "model_type", None)
  72    else:
  73        raise TypeError("The model_type must be an encoder name or a torch module.")
  74
  75    model_kwargs = {"img_size": img_size}
  76    if config is not None and config.stats is not None:
  77        model_kwargs[f"use_{config.stats}_stats"] = True
  78    compatible_skips = config is not None and encoder_name in config.skip_models
  79    if ndim == 2 and (not compatible_skips or getattr(encoder, "apply_neck", False)):
  80        model_kwargs["use_skip_connection"] = False
  81    if backbone == "torchvision" and encoder_checkpoint is not None:
  82        model_kwargs["pretrained"] = False
  83    model_kwargs.update(init_kwargs)
  84
  85    model = model_class(backbone=backbone, encoder=encoder, encoder_checkpoint=encoder_checkpoint, **model_kwargs)
  86    if init_decoder is not None:
  87        model = init_decoder(model)
  88
  89    return model
  90
  91
  92#
  93# UNETR IMPLEMENTATION [Vision Transformer (ViT from SAM / CellposeSAM / SAM2 / SAM3 /
  94# DINOv2 / DINOv3 / MAE / ScaleMAE / torchvision) + UNet Decoder from `torch_em`]
  95#
  96
  97
  98def _check_input_normalization_range(
  99    x: torch.Tensor,
 100    expected_range: Optional[Tuple[float, float]],
 101    unit_scale_max: Optional[float] = None,
 102) -> None:
 103    """Check whether raw inputs match the value range expected by the model normalizer.
 104
 105    Args:
 106        x: The input tensor to validate.
 107        expected_range: The (min, max) value range the input must lie within. If None, all checks are skipped.
 108        unit_scale_max: If set, raises an error when the input's maximum value is at or below this threshold,
 109            catching inputs that are likely in the wrong scale (e.g. [0, 1] instead of [0, 255] for SAM1).
 110    """
 111    if expected_range is None:
 112        return
 113
 114    if not torch.all(torch.isfinite(x)):
 115        raise ValueError("The input contains NaN or infinite values before normalization.")
 116
 117    min_value, max_value = expected_range
 118    if torch.any((x < min_value) | (x > max_value)):
 119        actual_min, actual_max = torch.aminmax(x.detach())
 120        raise ValueError(
 121            "The input is outside the expected scale before normalization: "
 122            f"expected values in [{min_value}, {max_value}], got [{actual_min.item()}, {actual_max.item()}]. "
 123            "Please check whether the raw inputs should be scaled to [0, 1] or kept in [0, 255] "
 124            "before applying the pretrained normalization statistics."
 125        )
 126
 127    if unit_scale_max is not None:
 128        actual_max = x.detach().max().item()
 129        if actual_max <= unit_scale_max:
 130            raise ValueError(
 131                f"The input maximum value ({actual_max:.4f}) suggests the input is in the wrong scale: "
 132                f"expected inputs with values in [{min_value}, {max_value}], "
 133                f"but the maximum is only {actual_max:.4f}. "
 134                "Please check whether the raw inputs should be scaled to [0, 255] instead of [0, 1]."
 135            )
 136
 137
 138def _as_stats(mean, std, device, dtype, is_3d: bool):
 139    view_shape = (1, -1, 1, 1, 1) if is_3d else (1, -1, 1, 1)
 140    pixel_mean = torch.tensor(mean, device=device, dtype=dtype).view(*view_shape)
 141    pixel_std = torch.tensor(std, device=device, dtype=dtype).view(*view_shape)
 142    return pixel_mean, pixel_std
 143
 144
 145class UNETRBase(nn.Module):
 146    """Base class for implementing a UNETR.
 147
 148    Args:
 149        img_size: The size of the input for the image encoder. Input images will be resized to match this size.
 150        backbone: The name of the vision transformer implementation.
 151            One of "sam", "sam2", "sam3", "cellpose_sam", "mae", "scalemae", "dinov2", "dinov3"
 152            (see all combinations below)
 153        encoder: The vision transformer. Can either be a name, such as "vit_b"
 154            (see all combinations for this below) or a torch module.
 155        decoder: The convolutional decoder.
 156        out_channels: The number of output channels of the UNETR.
 157        use_sam_stats: Whether to normalize the input data with the statistics of the
 158            pretrained SAM / SAM2 / SAM3 model.
 159        use_dino_stats: Whether to normalize the input data with the statistics of the
 160            pretrained DINOv2 / DINOv3 model.
 161        use_imagenet_stats: Whether to normalize with standard ImageNet statistics, i.e.
 162            mean - (0.485, 0.456, 0.406) and std - (0.229, 0.224, 0.225), raw inputs between range [0, 1].
 163            Use this with the 'torchvision' backbone when loading pretrained weights.
 164        use_mae_stats: Whether to normalize the input data with the statistics of the pretrained MAE model.
 165        resize_input: Whether to resize the input images to match `img_size`.
 166            By default, it resizes the inputs to match the `img_size`.
 167        encoder_checkpoint: Checkpoint for initializing the vision transformer.
 168            Can either be a filepath or an already loaded checkpoint.
 169        final_activation: The activation to apply to the UNETR output.
 170        use_skip_connection: Whether to use skip connections. By default, it uses skip connections.
 171        embed_dim: The embedding dimensionality, corresponding to the output dimension of the vision transformer.
 172        use_conv_transpose: Whether to use transposed convolutions instead of resampling for upsampling.
 173            By default, it uses resampling for upsampling.
 174        perform_range_checks: Whether to validate the input value range before normalization on each forward pass.
 175            You can disable the checks to avoid GPU sync overhead during training when inputs are known to be correct.
 176        initial_features: The number of features of the finest decoder level. The features per level are
 177            'initial_features * gain ** i', so this scales the decoder parameters quadratically.
 178
 179        NOTE: The currently supported combinations of 'backbone' x 'encoder' are the following:
 180
 181        SAM_family_models:
 182            - 'sam' x 'vit_b'
 183            - 'sam' x 'vit_l'
 184            - 'sam' x 'vit_h'
 185            - 'sam2' x 'hvit_t'
 186            - 'sam2' x 'hvit_s'
 187            - 'sam2' x 'hvit_b'
 188            - 'sam2' x 'hvit_l'
 189            - 'sam3' x 'vit_pe'
 190            - 'cellpose_sam' x 'vit_l'
 191
 192        DINO_family_models:
 193            - 'dinov2' x 'vit_s'
 194            - 'dinov2' x 'vit_b'
 195            - 'dinov2' x 'vit_l'
 196            - 'dinov2' x 'vit_g'
 197            - 'dinov2' x 'vit_s_reg4'
 198            - 'dinov2' x 'vit_b_reg4'
 199            - 'dinov2' x 'vit_l_reg4'
 200            - 'dinov2' x 'vit_g_reg4'
 201            - 'dinov3' x 'vit_s'
 202            - 'dinov3' x 'vit_s+'
 203            - 'dinov3' x 'vit_b'
 204            - 'dinov3' x 'vit_l'
 205            - 'dinov3' x 'vit_l+'
 206            - 'dinov3' x 'vit_h+'
 207            - 'dinov3' x 'vit_7b'
 208
 209        MAE_family_models:
 210            - 'mae' x 'vit_b'
 211            - 'mae' x 'vit_l'
 212            - 'mae' x 'vit_h'
 213            - 'scalemae' x 'vit_b'
 214            - 'scalemae' x 'vit_l'
 215            - 'scalemae' x 'vit_h'
 216
 217        torchvision_models:
 218            - 'torchvision' x 'vit_b_16'
 219            - 'torchvision' x 'vit_b_32'
 220            - 'torchvision' x 'vit_l_16'
 221            - 'torchvision' x 'vit_l_32'
 222            - 'torchvision' x 'vit_h_14'
 223    """
 224    def __init__(
 225        self,
 226        img_size: int = 1024,
 227        backbone: str = "sam",
 228        encoder: Optional[Union[nn.Module, str]] = "vit_b",
 229        decoder: Optional[nn.Module] = None,
 230        out_channels: int = 1,
 231        use_sam_stats: bool = False,
 232        use_mae_stats: bool = False,
 233        use_dino_stats: bool = False,
 234        use_imagenet_stats: bool = False,
 235        resize_input: bool = True,
 236        encoder_checkpoint: Optional[Union[str, PathLike, Mapping]] = None,
 237        final_activation: Optional[Union[str, nn.Module]] = None,
 238        use_skip_connection: bool = True,
 239        embed_dim: Optional[int] = None,
 240        use_conv_transpose: bool = False,
 241        perform_range_checks: bool = True,
 242        initial_features: int = 64,
 243        **kwargs
 244    ) -> None:
 245        super().__init__()
 246
 247        self.img_size = img_size
 248        self.use_sam_stats = use_sam_stats
 249        self.use_mae_stats = use_mae_stats
 250        self.use_dino_stats = use_dino_stats
 251        self.use_imagenet_stats = use_imagenet_stats
 252        self.use_skip_connection = use_skip_connection
 253        self.resize_input = resize_input
 254        self.perform_range_checks = perform_range_checks
 255        self.use_conv_transpose = use_conv_transpose
 256        self.initial_features = initial_features
 257        self.backbone = backbone
 258
 259        if isinstance(encoder, str):  # e.g. "vit_b" / "hvit_b" / "vit_pe"
 260            print(f"Using {encoder} from {backbone.upper()}")
 261            self.encoder = get_vision_transformer(img_size=img_size, backbone=backbone, model=encoder, **kwargs)
 262        else:  # `nn.Module` ViT backbone
 263            self.encoder = encoder
 264
 265        if encoder_checkpoint is not None:
 266            self._load_encoder_from_checkpoint(backbone, encoder_checkpoint)
 267
 268        # Sort the `embed_dim` out using the encoder output width.
 269        if embed_dim is None:
 270            have_neck = hasattr(self.encoder, "neck") and not hasattr(self.encoder, "embed_dim")
 271            # For SAM1 encoder, if 'apply_neck' is applied, the embedding dimension must change.
 272            if getattr(self.encoder, "apply_neck", have_neck):
 273                embed_dim = self.encoder.neck[2].out_channels  # 256 for SAM1.
 274            elif hasattr(self.encoder, "embed_dim"):
 275                embed_dim = self.encoder.embed_dim
 276            else:
 277                embed_dim = self.encoder.patch_embed.proj.out_channels
 278
 279        self.embed_dim = embed_dim
 280        self.final_activation = self._get_activation(final_activation)
 281
 282    def _load_encoder_from_checkpoint(self, backbone, checkpoint):
 283        """Load encoder weights from an encoder state or a full model checkpoint."""
 284        if isinstance(checkpoint, (str, PathLike)):
 285            # Load the checkpoint directly; extract encoder weights from full model checkpoints below.
 286            encoder_state = torch.load(checkpoint, map_location="cpu", weights_only=False)
 287        else:
 288            encoder_state = checkpoint
 289
 290        for state_key in ("model_state", "model", "state_dict"):
 291            if state_key in encoder_state:
 292                encoder_state = encoder_state[state_key]
 293        # Handle DataParallel/DistributedDataParallel prefix.
 294        encoder_state = {key.removeprefix("module."): value for key, value in encoder_state.items()}
 295
 296        if backbone in ("sam", "sam2"):
 297            # Extract the image encoder state from the full SAM or SAM2 checkpoint.
 298            prefixes = ("image_encoder.", "encoder.")
 299        elif backbone == "sam3":
 300            # Align the trunk names and leave out the neck's 'convs' and 'sam2_convs' upsampling blocks.
 301            prefixes = (
 302                "detector.backbone.vision_backbone.trunk.", "backbone.vision_backbone.trunk.",
 303                "vision_backbone.trunk.", "trunk.", "encoder."
 304            )
 305        elif backbone == "torchvision":
 306            prefixes = ("encoder.",) if "encoder.conv_proj.weight" in encoder_state else ()
 307        else:
 308            # CellposeSAM matches the pretrained rel_pos sizes, so its weights need no interpolation.
 309            # Extract encoder weights from CellposeSAM checkpoint format (strip 'encoder.' prefix).
 310            prefixes = ("encoder.",)
 311        for prefix in prefixes:
 312            if any(key.startswith(prefix) for key in encoder_state):
 313                encoder_state = {
 314                    key[len(prefix):]: value for key, value in encoder_state.items() if key.startswith(prefix)
 315                }
 316                break
 317
 318        if backbone in ("mae", "scalemae"):
 319            # ViT initialization hints from:
 320            # https://github.com/facebookresearch/mae/blob/main/main_finetune.py#L233-L242
 321            excluded = ("mask_token", "decoder", "head.")
 322            if backbone == "scalemae":
 323                # NOTE: ScaleMAE computes positional embeddings from the input resolution.
 324                excluded += ("fcn", "fpn", "pos_embed")
 325                if hasattr(self.encoder, "pos_embed"):
 326                    del self.encoder.pos_embed
 327            encoder_state = {key: value for key, value in encoder_state.items() if not key.startswith(excluded)}
 328            # Remove the classification head, which the MAE pretrained encoder weights do not need.
 329            if hasattr(self.encoder, "head"):
 330                del self.encoder.head
 331        elif backbone == "torchvision":
 332            encoder_state = {k: v for k, v in encoder_state.items() if not k.startswith("heads.")}
 333
 334        self.encoder.load_state_dict(encoder_state)
 335
 336    def _get_activation(self, activation):
 337        return_activation = None
 338        if activation is None:
 339            return None
 340        if isinstance(activation, nn.Module):
 341            return activation
 342        if isinstance(activation, str):
 343            return_activation = getattr(nn, activation, None)
 344        if return_activation is None:
 345            raise ValueError(f"Invalid activation: {activation}")
 346
 347        return return_activation()
 348
 349    @staticmethod
 350    def get_preprocess_shape(oldh: int, oldw: int, long_side_length: int) -> Tuple[int, int]:
 351        """Compute the output size given input size and target long side length.
 352
 353        Args:
 354            oldh: The input image height.
 355            oldw: The input image width.
 356            long_side_length: The longest side length for resizing.
 357
 358        Returns:
 359            The new image height.
 360            The new image width.
 361        """
 362        scale = long_side_length * 1.0 / max(oldh, oldw)
 363        newh, neww = oldh * scale, oldw * scale
 364        neww = int(neww + 0.5)
 365        newh = int(newh + 0.5)
 366        return (newh, neww)
 367
 368    def resize_longest_side(self, image: torch.Tensor) -> torch.Tensor:
 369        """Resize the image so that the longest side has the correct length.
 370
 371        Expects batched images with shape BxCxHxW OR BxCxDxHxW and float format.
 372
 373        Args:
 374            image: The input image.
 375
 376        Returns:
 377            The resized image.
 378        """
 379        if image.ndim == 4:  # i.e. 2d image
 380            target_size = self.get_preprocess_shape(image.shape[2], image.shape[3], self.encoder.img_size)
 381            return F.interpolate(image, target_size, mode="bilinear", align_corners=False, antialias=True)
 382        elif image.ndim == 5:  # i.e. 3d volume
 383            B, C, Z, H, W = image.shape
 384            target_size = self.get_preprocess_shape(H, W, self.img_size)
 385            return F.interpolate(image, (Z, *target_size), mode="trilinear", align_corners=False)
 386        else:
 387            raise ValueError("Expected 4d or 5d inputs, got", image.shape)
 388
 389    def _as_stats(self, mean, std, device, dtype, is_3d: bool):
 390        """@private
 391        """
 392        return _as_stats(mean, std, device, dtype, is_3d)
 393
 394    def _check_input_normalization_range(self, x: torch.Tensor, expected_range: Optional[Tuple[float, float]]) -> None:
 395        """@private
 396        """
 397        _check_input_normalization_range(x, expected_range)
 398
 399    def encode(self, x: torch.Tensor):
 400        """Preprocess the input and run the image encoder.
 401
 402        Args:
 403            x: The input tensor.
 404
 405        Returns:
 406            The encoder features to pass to `decode` and the spatial shape after preprocessing.
 407        """
 408        raise NotImplementedError
 409
 410    def decode(self, features, input_shape: Tuple[int, ...], original_shape: Tuple[int, ...]) -> torch.Tensor:
 411        """Run the convolutional decoder on the encoder features.
 412
 413        Args:
 414            features: The encoder features returned by `encode`.
 415            input_shape: The spatial shape after preprocessing, returned by `encode`.
 416            original_shape: The spatial shape of the original input.
 417
 418        Returns:
 419            The UNETR output, resized to `original_shape`.
 420        """
 421        raise NotImplementedError
 422
 423    def forward(self, x: torch.Tensor) -> torch.Tensor:
 424        """Apply the UNETR to the input data.
 425
 426        Args:
 427            x: The input tensor.
 428
 429        Returns:
 430            The UNETR output.
 431        """
 432        features, input_shape = self.encode(x)
 433        return self.decode(features, input_shape, tuple(x.shape[2:]))
 434
 435    def _encode_image(self, x):
 436        # Encoders can return either:
 437        # - only the image embeddings
 438        # - image embeddings and intermediate attention features
 439        features = self.encoder(x)
 440        return (features, None) if isinstance(features, torch.Tensor) else features
 441
 442    def preprocess(self, x: torch.Tensor) -> torch.Tensor:
 443        """@private
 444        """
 445        return preprocess_vit_inputs(
 446            x,
 447            use_sam_stats=self.use_sam_stats,
 448            backbone=self.backbone,
 449            use_mae_stats=self.use_mae_stats,
 450            use_dino_stats=self.use_dino_stats,
 451            use_imagenet_stats=self.use_imagenet_stats,
 452            resize_input=self.resize_input,
 453            img_size=self.img_size,
 454            encoder_img_size=self.encoder.img_size,
 455            perform_range_checks=self.perform_range_checks,
 456        )
 457
 458    def postprocess_masks(
 459        self, masks: torch.Tensor, input_size: Tuple[int, ...], original_size: Tuple[int, ...],
 460    ) -> torch.Tensor:
 461        """@private
 462        """
 463        if masks.ndim == 4:  # i.e. 2d labels
 464            masks = F.interpolate(
 465                masks,
 466                (self.encoder.img_size, self.encoder.img_size),
 467                mode="bilinear",
 468                align_corners=False,
 469            )
 470            masks = masks[..., : input_size[0], : input_size[1]]
 471            masks = F.interpolate(masks, original_size, mode="bilinear", align_corners=False)
 472
 473        elif masks.ndim == 5:  # i.e. 3d volumetric labels
 474            masks = F.interpolate(
 475                masks,
 476                (input_size[0], self.img_size, self.img_size),
 477                mode="trilinear",
 478                align_corners=False,
 479            )
 480            masks = masks[..., :input_size[0], :input_size[1], :input_size[2]]
 481            masks = F.interpolate(masks, original_size, mode="trilinear", align_corners=False)
 482
 483        else:
 484            raise ValueError("Expected 4d or 5d labels, got", masks.shape)
 485
 486        return masks
 487
 488
 489def preprocess_vit_inputs(
 490    x: torch.Tensor,
 491    use_sam_stats: bool = False,
 492    backbone: str = "sam",
 493    use_mae_stats: bool = False,
 494    use_dino_stats: bool = False,
 495    use_imagenet_stats: bool = False,
 496    resize_input: bool = True,
 497    img_size: int = 1024,
 498    encoder_img_size: int = 1024,
 499    perform_range_checks: bool = True,
 500) -> Tuple[torch.Tensor, Tuple]:
 501    """Preprocess inputs for ViT-backbones in UNETR models.
 502
 503    Handles normalization stat selection, input range validation, optional resizing to the longest side,
 504    and padding to `encoder_img_size`. Can be used as a standalone function without a model instance.
 505
 506    Args:
 507        x: Input tensor of shape (B, C, H, W) for 2D or (B, C, Z, H, W) for 3D.
 508        use_sam_stats: Whether to normalize with SAM/SAM2/SAM3 backbone statistics.
 509        backbone: The backbone name - controls which SAM stats are used when `use_sam_stats=True`.
 510        use_mae_stats: Whether to normalize with MAE statistics.
 511        use_dino_stats: Whether to normalize with DINOv2/DINOv3 statistics.
 512        use_imagenet_stats: Whether to normalize with standard ImageNet statistics
 513            (mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), range [0, 1]).
 514            Use this for torchvision pretrained backbones.
 515        resize_input: Whether to resize the input to the longest side before padding.
 516        img_size: The model image size, used for 3D resize.
 517        encoder_img_size: The encoder image size, used for 2D resize and padding.
 518        perform_range_checks: Whether to validate the expected input value range before normalization.
 519            You can disable the checks to avoid GPU sync overhead during training when inputs are known to be correct.
 520
 521    Returns:
 522        The preprocessed tensor and the spatial shape after resizing (before padding).
 523    """
 524    is_3d = (x.ndim == 5)
 525    device, dtype = x.device, x.dtype
 526    mean, std = (0.0, 0.0, 0.0), (1.0, 1.0, 1.0)
 527    expected_range = None
 528    unit_scale_max = None
 529
 530    if use_sam_stats:
 531        if backbone == "sam2":
 532            mean, std = (0.485, 0.456, 0.406), (0.229, 0.224, 0.225)
 533            expected_range = (0.0, 1.0)
 534        elif backbone == "sam3":
 535            mean, std = (0.5, 0.5, 0.5), (0.5, 0.5, 0.5)
 536            expected_range = (0.0, 1.0)
 537        else:  # sam1 / default
 538            mean, std = (123.675, 116.28, 103.53), (58.395, 57.12, 57.375)
 539            expected_range = (0.0, 255.0)
 540            unit_scale_max = 1.0
 541    elif use_mae_stats or use_dino_stats or use_imagenet_stats:
 542        # MAE, ScaleMAE, and DINO use the same ImageNet mean and standard deviation.
 543        mean, std = (0.485, 0.456, 0.406), (0.229, 0.224, 0.225)
 544        expected_range = (0.0, 1.0)
 545    else:
 546        mean, std = (0.0, 0.0, 0.0), (1.0, 1.0, 1.0)
 547        expected_range = None
 548
 549    if perform_range_checks:
 550        _check_input_normalization_range(x, expected_range, unit_scale_max)
 551    pixel_mean, pixel_std = _as_stats(mean, std, device=device, dtype=dtype, is_3d=is_3d)
 552
 553    if resize_input:
 554        if x.ndim == 4:
 555            target_size = UNETRBase.get_preprocess_shape(x.shape[2], x.shape[3], encoder_img_size)
 556            x = F.interpolate(x, target_size, mode="bilinear", align_corners=False, antialias=True)
 557        elif x.ndim == 5:
 558            B, C, Z, H, W = x.shape
 559            target_size = UNETRBase.get_preprocess_shape(H, W, img_size)
 560            x = F.interpolate(x, (Z, *target_size), mode="trilinear", align_corners=False)
 561
 562    input_shape = x.shape[-3:] if is_3d else x.shape[-2:]
 563
 564    x = (x - pixel_mean) / pixel_std
 565    h, w = x.shape[-2:]
 566    padh = encoder_img_size - h
 567    padw = encoder_img_size - w
 568
 569    if is_3d:
 570        x = F.pad(x, (0, padw, 0, padh, 0, 0))
 571    else:
 572        x = F.pad(x, (0, padw, 0, padh))
 573
 574    return x, input_shape
 575
 576
 577class UNETR(UNETRBase):
 578    """A (2d-only) UNet Transformer using a vision transformer as encoder and a convolutional decoder.
 579    """
 580    def __init__(
 581        self,
 582        img_size: int = 1024,
 583        backbone: str = "sam",
 584        encoder: Optional[Union[nn.Module, str]] = "vit_b",
 585        decoder: Optional[nn.Module] = None,
 586        out_channels: int = 1,
 587        use_sam_stats: bool = False,
 588        use_mae_stats: bool = False,
 589        use_dino_stats: bool = False,
 590        use_imagenet_stats: bool = False,
 591        resize_input: bool = True,
 592        encoder_checkpoint: Optional[Union[str, PathLike, Mapping]] = None,
 593        final_activation: Optional[Union[str, nn.Module]] = None,
 594        use_skip_connection: bool = True,
 595        embed_dim: Optional[int] = None,
 596        use_conv_transpose: bool = False,
 597        perform_range_checks: bool = True,
 598        **kwargs
 599    ) -> None:
 600
 601        super().__init__(
 602            img_size=img_size,
 603            backbone=backbone,
 604            encoder=encoder,
 605            decoder=decoder,
 606            out_channels=out_channels,
 607            use_sam_stats=use_sam_stats,
 608            use_mae_stats=use_mae_stats,
 609            use_dino_stats=use_dino_stats,
 610            use_imagenet_stats=use_imagenet_stats,
 611            resize_input=resize_input,
 612            encoder_checkpoint=encoder_checkpoint,
 613            final_activation=final_activation,
 614            use_skip_connection=use_skip_connection,
 615            embed_dim=embed_dim,
 616            use_conv_transpose=use_conv_transpose,
 617            perform_range_checks=perform_range_checks,
 618            **kwargs,
 619        )
 620        self.init_kwargs = {
 621            "img_size": img_size, "backbone": backbone, "encoder": encoder, "decoder": decoder,
 622            "out_channels": out_channels, "use_sam_stats": use_sam_stats,
 623            "use_mae_stats": use_mae_stats, "use_dino_stats": use_dino_stats,
 624            "use_imagenet_stats": use_imagenet_stats, "resize_input": resize_input,
 625            "final_activation": final_activation,
 626            "use_skip_connection": use_skip_connection, "embed_dim": embed_dim,
 627            "use_conv_transpose": use_conv_transpose, "perform_range_checks": perform_range_checks,
 628            **kwargs,
 629        }
 630
 631        encoder = self.encoder
 632
 633        if backbone == "sam2" and hasattr(encoder, "trunk"):
 634            in_chans = encoder.trunk.patch_embed.proj.in_channels
 635        elif hasattr(encoder, "in_chans"):
 636            in_chans = encoder.in_chans
 637        else:  # `nn.Module` ViT backbone.
 638            try:
 639                in_chans = encoder.patch_embed.proj.in_channels
 640            except AttributeError:  # for getting the input channels while using 'vit_t' from MobileSam
 641                in_chans = encoder.patch_embed.seq[0].c.in_channels
 642
 643        # parameters for the decoder network
 644        depth = 3
 645        gain = 2
 646        features_decoder = [self.initial_features * gain ** i for i in range(depth + 1)][::-1]
 647        scale_factors = depth * [2]
 648        self.out_channels = out_channels
 649
 650        # choice of upsampler - to use (bilinear interpolation + conv) or conv transpose
 651        _upsampler = SingleDeconv2DBlock if use_conv_transpose else Upsampler2d
 652
 653        self.decoder = decoder or Decoder(
 654            features=features_decoder,
 655            scale_factors=scale_factors[::-1],
 656            conv_block_impl=ConvBlock2d,
 657            sampler_impl=_upsampler,
 658        )
 659
 660        if use_skip_connection:
 661            self.deconv1 = Deconv2DBlock(
 662                in_channels=self.embed_dim,
 663                out_channels=features_decoder[0],
 664                use_conv_transpose=use_conv_transpose,
 665            )
 666            self.deconv2 = nn.Sequential(
 667                Deconv2DBlock(
 668                    in_channels=self.embed_dim,
 669                    out_channels=features_decoder[0],
 670                    use_conv_transpose=use_conv_transpose,
 671                ),
 672                Deconv2DBlock(
 673                    in_channels=features_decoder[0],
 674                    out_channels=features_decoder[1],
 675                    use_conv_transpose=use_conv_transpose,
 676                )
 677            )
 678            self.deconv3 = nn.Sequential(
 679                Deconv2DBlock(
 680                    in_channels=self.embed_dim,
 681                    out_channels=features_decoder[0],
 682                    use_conv_transpose=use_conv_transpose,
 683                ),
 684                Deconv2DBlock(
 685                    in_channels=features_decoder[0],
 686                    out_channels=features_decoder[1],
 687                    use_conv_transpose=use_conv_transpose,
 688                ),
 689                Deconv2DBlock(
 690                    in_channels=features_decoder[1],
 691                    out_channels=features_decoder[2],
 692                    use_conv_transpose=use_conv_transpose,
 693                )
 694            )
 695            self.deconv4 = ConvBlock2d(in_chans, features_decoder[-1])
 696        else:
 697            self.deconv1 = Deconv2DBlock(
 698                in_channels=self.embed_dim,
 699                out_channels=features_decoder[0],
 700                use_conv_transpose=use_conv_transpose,
 701            )
 702            self.deconv2 = Deconv2DBlock(
 703                in_channels=features_decoder[0],
 704                out_channels=features_decoder[1],
 705                use_conv_transpose=use_conv_transpose,
 706            )
 707            self.deconv3 = Deconv2DBlock(
 708                in_channels=features_decoder[1],
 709                out_channels=features_decoder[2],
 710                use_conv_transpose=use_conv_transpose,
 711            )
 712            self.deconv4 = Deconv2DBlock(
 713                in_channels=features_decoder[2],
 714                out_channels=features_decoder[3],
 715                use_conv_transpose=use_conv_transpose,
 716            )
 717
 718        self.base = ConvBlock2d(self.embed_dim, features_decoder[0])
 719        self.out_conv = nn.Conv2d(features_decoder[-1], out_channels, 1)
 720        self.deconv_out = _upsampler(
 721            scale_factor=2, in_channels=features_decoder[-1], out_channels=features_decoder[-1]
 722        )
 723        self.decoder_head = ConvBlock2d(2 * features_decoder[-1], features_decoder[-1])
 724
 725    def encode(self, x: torch.Tensor):
 726        """Preprocess the input and run the image encoder.
 727
 728        Args:
 729            x: The input tensor of shape (B, C, Y, X).
 730
 731        Returns:
 732            The features as a tuple of the image embeddings, the list of intermediate encoder outputs
 733            (None if the encoder returns only the embeddings) and the preprocessed input, which the
 734            skip connections consume, and the spatial shape after preprocessing.
 735        """
 736        # Reshape the inputs to the shape expected by the encoder
 737        # and normalize the inputs if normalization is part of the model.
 738        x, input_shape = self.preprocess(x)
 739        z12, from_encoder = self._encode_image(x)
 740        return (z12, from_encoder, x), input_shape
 741
 742    def decode(self, features, input_shape: Tuple[int, ...], original_shape: Tuple[int, ...]) -> torch.Tensor:
 743        """Run the convolutional decoder on the encoder features.
 744
 745        Args:
 746            features: The tuple returned by `encode`.
 747            input_shape: The spatial shape (Y, X) after preprocessing.
 748            original_shape: The spatial shape (Y, X) of the original input.
 749
 750        Returns:
 751            The UNETR output, resized to `original_shape`.
 752        """
 753        z12, from_encoder, x = features
 754
 755        if self.use_skip_connection:
 756            from_encoder = from_encoder[::-1]
 757            z9 = self.deconv1(from_encoder[0])
 758            z6 = self.deconv2(from_encoder[1])
 759            z3 = self.deconv3(from_encoder[2])
 760            z0 = self.deconv4(x)
 761
 762        else:
 763            z9 = self.deconv1(z12)
 764            z6 = self.deconv2(z9)
 765            z3 = self.deconv3(z6)
 766            z0 = self.deconv4(z3)
 767
 768        updated_from_encoder = [z9, z6, z3]
 769
 770        x = self.base(z12)
 771        x = self.decoder(x, encoder_inputs=updated_from_encoder)
 772        x = self.deconv_out(x)
 773
 774        x = torch.cat([x, z0], dim=1)
 775        x = self.decoder_head(x)
 776
 777        x = self.out_conv(x)
 778        if self.final_activation is not None:
 779            x = self.final_activation(x)
 780
 781        return self.postprocess_masks(x, input_shape, original_shape)
 782
 783
 784class UNETR2D(UNETR):
 785    """A two-dimensional UNet Transformer using a vision transformer as encoder and a convolutional decoder.
 786    """
 787    pass
 788
 789
 790class UNETR3D(UNETRBase):
 791    """A 3D UNet Transformer with a vision transformer encoder and a convolutional decoder."""
 792    def __init__(
 793        self,
 794        img_size: int = 1024,
 795        backbone: str = "sam",
 796        encoder: Optional[Union[nn.Module, str]] = "vit_b",
 797        decoder: Optional[nn.Module] = None,
 798        out_channels: int = 1,
 799        use_sam_stats: bool = False,
 800        use_mae_stats: bool = False,
 801        use_dino_stats: bool = False,
 802        use_imagenet_stats: bool = False,
 803        resize_input: bool = True,
 804        encoder_checkpoint: Optional[Union[str, PathLike, Mapping]] = None,
 805        final_activation: Optional[Union[str, nn.Module]] = None,
 806        use_skip_connection: bool = False,
 807        embed_dim: Optional[int] = None,
 808        use_conv_transpose: bool = False,
 809        use_strip_pooling: bool = True,
 810        perform_range_checks: bool = True,
 811        **kwargs
 812    ):
 813        if use_skip_connection:
 814            raise NotImplementedError("The framework cannot handle skip connections atm.")
 815        if use_conv_transpose:
 816            raise NotImplementedError("It's not enabled to switch between interpolation and transposed convolutions.")
 817
 818        super().__init__(
 819            img_size=img_size,
 820            backbone=backbone,
 821            encoder=encoder,
 822            decoder=decoder,
 823            out_channels=out_channels,
 824            use_sam_stats=use_sam_stats,
 825            use_mae_stats=use_mae_stats,
 826            use_dino_stats=use_dino_stats,
 827            use_imagenet_stats=use_imagenet_stats,
 828            resize_input=resize_input,
 829            encoder_checkpoint=encoder_checkpoint,
 830            final_activation=final_activation,
 831            use_skip_connection=use_skip_connection,
 832            embed_dim=embed_dim,
 833            use_conv_transpose=use_conv_transpose,
 834            perform_range_checks=perform_range_checks,
 835            **kwargs,
 836        )
 837        self.init_kwargs = {
 838            "img_size": img_size, "backbone": backbone, "encoder": encoder, "decoder": decoder,
 839            "out_channels": out_channels, "use_sam_stats": use_sam_stats,
 840            "use_mae_stats": use_mae_stats, "use_dino_stats": use_dino_stats,
 841            "use_imagenet_stats": use_imagenet_stats, "resize_input": resize_input,
 842            "final_activation": final_activation,
 843            "use_skip_connection": use_skip_connection, "embed_dim": embed_dim,
 844            "use_conv_transpose": use_conv_transpose, "use_strip_pooling": use_strip_pooling,
 845            "perform_range_checks": perform_range_checks, **kwargs,
 846        }
 847
 848        # The 3d convolutional decoder.
 849        # First, get the important parameters for the decoder.
 850        depth = 3
 851        gain = 2
 852        features_decoder = [self.initial_features * gain ** i for i in range(depth + 1)][::-1]
 853        scale_factors = [1, 2, 2]
 854        self.out_channels = out_channels
 855
 856        # The mapping blocks.
 857        self.deconv1 = Deconv3DBlock(
 858            in_channels=self.embed_dim,
 859            out_channels=features_decoder[0],
 860            scale_factor=scale_factors,
 861            use_strip_pooling=use_strip_pooling,
 862        )
 863        self.deconv2 = Deconv3DBlock(
 864            in_channels=features_decoder[0],
 865            out_channels=features_decoder[1],
 866            scale_factor=scale_factors,
 867            use_strip_pooling=use_strip_pooling,
 868        )
 869        self.deconv3 = Deconv3DBlock(
 870            in_channels=features_decoder[1],
 871            out_channels=features_decoder[2],
 872            scale_factor=scale_factors,
 873            use_strip_pooling=use_strip_pooling,
 874        )
 875        self.deconv4 = Deconv3DBlock(
 876            in_channels=features_decoder[2],
 877            out_channels=features_decoder[3],
 878            scale_factor=scale_factors,
 879            use_strip_pooling=use_strip_pooling,
 880        )
 881
 882        # The core decoder block.
 883        self.decoder = decoder or Decoder(
 884            features=features_decoder,
 885            scale_factors=[scale_factors] * depth,
 886            conv_block_impl=partial(ConvBlock3dWithStrip, use_strip_pooling=use_strip_pooling),
 887            sampler_impl=Upsampler3d,
 888        )
 889
 890        # And the final upsampler to match the expected dimensions.
 891        self.deconv_out = Deconv3DBlock(  # NOTE: changed `end_up` to `deconv_out`
 892            in_channels=features_decoder[-1],
 893            out_channels=features_decoder[-1],
 894            scale_factor=scale_factors,
 895            use_strip_pooling=use_strip_pooling,
 896        )
 897
 898        # Additional conjunction blocks.
 899        self.base = ConvBlock3dWithStrip(
 900            in_channels=self.embed_dim, out_channels=features_decoder[0], use_strip_pooling=use_strip_pooling
 901        )
 902
 903        # And the output layers.
 904        self.decoder_head = ConvBlock3dWithStrip(
 905            in_channels=2 * features_decoder[-1],
 906            out_channels=features_decoder[-1],
 907            use_strip_pooling=use_strip_pooling
 908        )
 909        self.out_conv = nn.Conv3d(features_decoder[-1], out_channels, 1)
 910
 911    def encode(self, x: torch.Tensor):
 912        """Preprocess the input and run the image encoder on every z-slice.
 913
 914        Args:
 915            x: The input tensor with shape (B, C, Z, Y, X). The depth Z can vary.
 916
 917        Returns:
 918            The encoder features of shape (B, D, Z, Y', X') and the spatial shape after preprocessing.
 919        """
 920        Z = x.shape[2]
 921        x, input_shape = self.preprocess(x)
 922        features = torch.stack([self._encode_image(x[:, :, i])[0] for i in range(Z)], dim=2)
 923        return features, input_shape
 924
 925    def decode(self, features, input_shape: Tuple[int, ...], original_shape: Tuple[int, ...]) -> torch.Tensor:
 926        """Run the convolutional decoder on the encoder features.
 927
 928        Args:
 929            features: The encoder features with shape (B, D, Z, Y', X'), returned by `encode`.
 930            input_shape: The spatial shape (Z, Y, X) after preprocessing.
 931            original_shape: The spatial shape (Z, Y, X) of the original input.
 932
 933        Returns:
 934            The UNETR output, resized to `original_shape`.
 935        """
 936        # Prepare the counterparts for the decoder.
 937        # NOTE: These mappings are sequential; they do not use intermediate encoder skip connections.
 938        z9 = self.deconv1(features)
 939        z6 = self.deconv2(z9)
 940        z3 = self.deconv3(z6)
 941        z0 = self.deconv4(z3)
 942
 943        updated_from_encoder = [z9, z6, z3]
 944
 945        # Align the features through the base block.
 946        x = self.base(features)
 947        # Run the decoder
 948        x = self.decoder(x, encoder_inputs=updated_from_encoder)
 949        x = self.deconv_out(x)  # NOTE: formerly `end_up`.
 950
 951        # And the final output head.
 952        x = torch.cat([x, z0], dim=1)
 953        x = self.decoder_head(x)
 954        x = self.out_conv(x)
 955        if self.final_activation is not None:
 956            x = self.final_activation(x)
 957
 958        # Postprocess the output back to original size.
 959        return self.postprocess_masks(x, input_shape, original_shape)
 960
 961
 962#
 963# ADDITIONAL FUNCTIONALITIES
 964#
 965
 966
 967def _strip_pooling_layers(enabled, channels) -> nn.Module:
 968    return DepthStripPooling(channels) if enabled else nn.Identity()
 969
 970
 971class DepthStripPooling(nn.Module):
 972    """@private
 973    """
 974    def __init__(self, channels: int, reduction: int = 4):
 975        """Block for strip pooling along the depth dimension (only).
 976
 977        eg. for 3D (Z > 1) - it aggregates global context across depth by adaptive avg pooling
 978        to Z=1, and then passes through a small 1x1x1 MLP, then broadcasts it back to Z to
 979        modulate the original features (using a gated residual).
 980
 981        For 2D (Z == 1): returns input unchanged (no-op).
 982
 983        Args:
 984            channels: The output channels.
 985            reduction: The reduction of the hidden layers.
 986        """
 987        super().__init__()
 988        hidden = max(1, channels // reduction)
 989        self.conv1 = nn.Conv3d(channels, hidden, kernel_size=1)
 990        self.bn1 = nn.BatchNorm3d(hidden)
 991        self.relu = nn.ReLU(inplace=True)
 992        self.conv2 = nn.Conv3d(hidden, channels, kernel_size=1)
 993
 994    def forward(self, x: torch.Tensor) -> torch.Tensor:
 995        if x.dim() != 5:
 996            raise ValueError(f"DepthStripPooling expects 5D tensors as input, got '{x.shape}'.")
 997
 998        B, C, Z, H, W = x.shape
 999        if Z == 1:  # i.e. always the case of all 2d.
1000            return x  # We simply do nothing there.
1001
1002        # We pool only along the depth dimension: i.e. target shape (B, C, 1, H, W).
1003        # A plain mean over Z is the same operation as adaptive_avg_pool3d to (1, H, W), but its
1004        # reduction kernel is several times faster at full resolution.
1005        feat = x.mean(dim=2, keepdim=True)
1006        feat = self.conv1(feat)
1007        feat = self.bn1(feat)
1008        feat = self.relu(feat)
1009        feat = self.conv2(feat)
1010        gate = torch.sigmoid(feat).expand(B, C, Z, H, W)  # Broadcast the collapsed depth context back to all slices
1011
1012        # Gated residual fusion
1013        return x * gate + x
1014
1015
1016class Deconv3DBlock(nn.Module):
1017    """@private
1018    """
1019    def __init__(
1020        self,
1021        scale_factor,
1022        in_channels,
1023        out_channels,
1024        kernel_size=3,
1025        anisotropic_kernel=True,
1026        use_strip_pooling=True,
1027    ):
1028        super().__init__()
1029        conv_block_kwargs = {
1030            "in_channels": out_channels,
1031            "out_channels": out_channels,
1032            "kernel_size": kernel_size,
1033            "padding": ((kernel_size - 1) // 2),
1034        }
1035        if anisotropic_kernel:
1036            conv_block_kwargs = _update_conv_kwargs(conv_block_kwargs, scale_factor)
1037
1038        self.block = nn.Sequential(
1039            Upsampler3d(scale_factor, in_channels, out_channels),
1040            nn.Conv3d(**conv_block_kwargs),
1041            nn.BatchNorm3d(out_channels),
1042            nn.ReLU(True),
1043            _strip_pooling_layers(enabled=use_strip_pooling, channels=out_channels),
1044        )
1045
1046    def forward(self, x):
1047        return self.block(x)
1048
1049
1050class ConvBlock3dWithStrip(nn.Module):
1051    """@private
1052    """
1053    def __init__(
1054        self, in_channels: int, out_channels: int, use_strip_pooling: bool = True, **kwargs
1055    ):
1056        super().__init__()
1057        self.block = nn.Sequential(
1058            ConvBlock3d(in_channels, out_channels, **kwargs),
1059            _strip_pooling_layers(enabled=use_strip_pooling, channels=out_channels),
1060        )
1061
1062    def forward(self, x):
1063        return self.block(x)
1064
1065
1066class SingleDeconv2DBlock(nn.Module):
1067    """@private
1068    """
1069    def __init__(self, scale_factor, in_channels, out_channels):
1070        super().__init__()
1071        self.block = nn.ConvTranspose2d(in_channels, out_channels, kernel_size=2, stride=2, padding=0, output_padding=0)
1072
1073    def forward(self, x):
1074        return self.block(x)
1075
1076
1077class SingleConv2DBlock(nn.Module):
1078    """@private
1079    """
1080    def __init__(self, in_channels, out_channels, kernel_size):
1081        super().__init__()
1082        self.block = nn.Conv2d(
1083            in_channels, out_channels, kernel_size=kernel_size, stride=1, padding=((kernel_size - 1) // 2)
1084        )
1085
1086    def forward(self, x):
1087        return self.block(x)
1088
1089
1090class Conv2DBlock(nn.Module):
1091    """@private
1092    """
1093    def __init__(self, in_channels, out_channels, kernel_size=3):
1094        super().__init__()
1095        self.block = nn.Sequential(
1096            SingleConv2DBlock(in_channels, out_channels, kernel_size),
1097            nn.BatchNorm2d(out_channels),
1098            nn.ReLU(True)
1099        )
1100
1101    def forward(self, x):
1102        return self.block(x)
1103
1104
1105class Deconv2DBlock(nn.Module):
1106    """@private
1107    """
1108    def __init__(self, in_channels, out_channels, kernel_size=3, use_conv_transpose=True):
1109        super().__init__()
1110        _upsampler = SingleDeconv2DBlock if use_conv_transpose else Upsampler2d
1111        self.block = nn.Sequential(
1112            _upsampler(scale_factor=2, in_channels=in_channels, out_channels=out_channels),
1113            SingleConv2DBlock(out_channels, out_channels, kernel_size),
1114            nn.BatchNorm2d(out_channels),
1115            nn.ReLU(True)
1116        )
1117
1118    def forward(self, x):
1119        return self.block(x)