torch_em.model.unetr

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