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 preprocess(self, x: torch.Tensor) -> torch.Tensor:
 379        """@private
 380        """
 381        return preprocess_vit_inputs(
 382            x,
 383            use_sam_stats=self.use_sam_stats,
 384            backbone=self.backbone,
 385            use_mae_stats=self.use_mae_stats,
 386            use_dino_stats=self.use_dino_stats,
 387            resize_input=self.resize_input,
 388            img_size=self.img_size,
 389            encoder_img_size=self.encoder.img_size,
 390            perform_range_checks=self.perform_range_checks,
 391        )
 392
 393    def postprocess_masks(
 394        self, masks: torch.Tensor, input_size: Tuple[int, ...], original_size: Tuple[int, ...],
 395    ) -> torch.Tensor:
 396        """@private
 397        """
 398        if masks.ndim == 4:  # i.e. 2d labels
 399            masks = F.interpolate(
 400                masks,
 401                (self.encoder.img_size, self.encoder.img_size),
 402                mode="bilinear",
 403                align_corners=False,
 404            )
 405            masks = masks[..., : input_size[0], : input_size[1]]
 406            masks = F.interpolate(masks, original_size, mode="bilinear", align_corners=False)
 407
 408        elif masks.ndim == 5:  # i.e. 3d volumetric labels
 409            masks = F.interpolate(
 410                masks,
 411                (input_size[0], self.img_size, self.img_size),
 412                mode="trilinear",
 413                align_corners=False,
 414            )
 415            masks = masks[..., :input_size[0], :input_size[1], :input_size[2]]
 416            masks = F.interpolate(masks, original_size, mode="trilinear", align_corners=False)
 417
 418        else:
 419            raise ValueError("Expected 4d or 5d labels, got", masks.shape)
 420
 421        return masks
 422
 423
 424def preprocess_vit_inputs(
 425    x: torch.Tensor,
 426    use_sam_stats: bool = False,
 427    backbone: str = "sam",
 428    use_mae_stats: bool = False,
 429    use_dino_stats: bool = False,
 430    resize_input: bool = True,
 431    img_size: int = 1024,
 432    encoder_img_size: int = 1024,
 433    perform_range_checks: bool = True,
 434) -> Tuple[torch.Tensor, Tuple]:
 435    """Preprocess inputs for ViT-backbones in UNETR models.
 436
 437    Handles normalization stat selection, input range validation, optional resizing to the longest side,
 438    and padding to `encoder_img_size`. Can be used as a standalone function without a model instance.
 439
 440    Args:
 441        x: Input tensor of shape (B, C, H, W) for 2D or (B, C, Z, H, W) for 3D.
 442        use_sam_stats: Whether to normalize with SAM/SAM2/SAM3 backbone statistics.
 443        backbone: The backbone name - controls which SAM stats are used when `use_sam_stats=True`.
 444        use_mae_stats: Whether to normalize with MAE statistics.
 445        use_dino_stats: Whether to normalize with DINOv2/DINOv3 statistics.
 446        resize_input: Whether to resize the input to the longest side before padding.
 447        img_size: The model image size, used for 3D resize.
 448        encoder_img_size: The encoder image size, used for 2D resize and padding.
 449        perform_range_checks: Whether to validate the expected input value range before normalization.
 450            You can disable the checks to avoid GPU sync overhead during training when inputs are known to be correct.
 451
 452    Returns:
 453        The preprocessed tensor and the spatial shape after resizing (before padding).
 454    """
 455    is_3d = (x.ndim == 5)
 456    device, dtype = x.device, x.dtype
 457    mean, std = (0.0, 0.0, 0.0), (1.0, 1.0, 1.0)
 458    expected_range = None
 459    unit_scale_max = None
 460
 461    if use_sam_stats:
 462        if backbone == "sam2":
 463            mean, std = (0.485, 0.456, 0.406), (0.229, 0.224, 0.225)
 464            expected_range = (0.0, 1.0)
 465        elif backbone == "sam3":
 466            mean, std = (0.5, 0.5, 0.5), (0.5, 0.5, 0.5)
 467            expected_range = (0.0, 1.0)
 468        else:  # sam1 / default
 469            mean, std = (123.675, 116.28, 103.53), (58.395, 57.12, 57.375)
 470            expected_range = (0.0, 255.0)
 471            unit_scale_max = 1.0
 472    elif use_mae_stats:  # TODO: add mean std from mae / scalemae experiments (or open up arguments for this)
 473        raise NotImplementedError
 474    elif use_dino_stats:
 475        mean, std = (0.485, 0.456, 0.406), (0.229, 0.224, 0.225)
 476        expected_range = (0.0, 1.0)
 477    else:
 478        mean, std = (0.0, 0.0, 0.0), (1.0, 1.0, 1.0)
 479        expected_range = None
 480
 481    if perform_range_checks:
 482        _check_input_normalization_range(x, expected_range, unit_scale_max)
 483    pixel_mean, pixel_std = _as_stats(mean, std, device=device, dtype=dtype, is_3d=is_3d)
 484
 485    if resize_input:
 486        if x.ndim == 4:
 487            target_size = UNETRBase.get_preprocess_shape(x.shape[2], x.shape[3], encoder_img_size)
 488            x = F.interpolate(x, target_size, mode="bilinear", align_corners=False, antialias=True)
 489        elif x.ndim == 5:
 490            B, C, Z, H, W = x.shape
 491            target_size = UNETRBase.get_preprocess_shape(H, W, img_size)
 492            x = F.interpolate(x, (Z, *target_size), mode="trilinear", align_corners=False)
 493
 494    input_shape = x.shape[-3:] if is_3d else x.shape[-2:]
 495
 496    x = (x - pixel_mean) / pixel_std
 497    h, w = x.shape[-2:]
 498    padh = encoder_img_size - h
 499    padw = encoder_img_size - w
 500
 501    if is_3d:
 502        x = F.pad(x, (0, padw, 0, padh, 0, 0))
 503    else:
 504        x = F.pad(x, (0, padw, 0, padh))
 505
 506    return x, input_shape
 507
 508
 509class UNETR(UNETRBase):
 510    """A (2d-only) UNet Transformer using a vision transformer as encoder and a convolutional decoder.
 511    """
 512    def __init__(
 513        self,
 514        img_size: int = 1024,
 515        backbone: Literal["sam", "sam2", "sam3", "cellpose_sam", "mae", "scalemae", "dinov2", "dinov3"] = "sam",
 516        encoder: Optional[Union[nn.Module, str]] = "vit_b",
 517        decoder: Optional[nn.Module] = None,
 518        out_channels: int = 1,
 519        use_sam_stats: bool = False,
 520        use_mae_stats: bool = False,
 521        use_dino_stats: bool = False,
 522        resize_input: bool = True,
 523        encoder_checkpoint: Optional[Union[str, OrderedDict]] = None,
 524        final_activation: Optional[Union[str, nn.Module]] = None,
 525        use_skip_connection: bool = True,
 526        embed_dim: Optional[int] = None,
 527        use_conv_transpose: bool = False,
 528        perform_range_checks: bool = True,
 529        **kwargs
 530    ) -> None:
 531
 532        super().__init__(
 533            img_size=img_size,
 534            backbone=backbone,
 535            encoder=encoder,
 536            decoder=decoder,
 537            out_channels=out_channels,
 538            use_sam_stats=use_sam_stats,
 539            use_mae_stats=use_mae_stats,
 540            use_dino_stats=use_dino_stats,
 541            resize_input=resize_input,
 542            encoder_checkpoint=encoder_checkpoint,
 543            final_activation=final_activation,
 544            use_skip_connection=use_skip_connection,
 545            embed_dim=embed_dim,
 546            use_conv_transpose=use_conv_transpose,
 547            perform_range_checks=perform_range_checks,
 548            **kwargs,
 549        )
 550
 551        encoder = self.encoder
 552
 553        if backbone == "sam2" and hasattr(encoder, "trunk"):
 554            in_chans = encoder.trunk.patch_embed.proj.in_channels
 555        elif hasattr(encoder, "in_chans"):
 556            in_chans = encoder.in_chans
 557        else:  # `nn.Module` ViT backbone.
 558            try:
 559                in_chans = encoder.patch_embed.proj.in_channels
 560            except AttributeError:  # for getting the input channels while using 'vit_t' from MobileSam
 561                in_chans = encoder.patch_embed.seq[0].c.in_channels
 562
 563        # parameters for the decoder network
 564        depth = 3
 565        gain = 2
 566        features_decoder = [self.initial_features * gain ** i for i in range(depth + 1)][::-1]
 567        scale_factors = depth * [2]
 568        self.out_channels = out_channels
 569
 570        # choice of upsampler - to use (bilinear interpolation + conv) or conv transpose
 571        _upsampler = SingleDeconv2DBlock if use_conv_transpose else Upsampler2d
 572
 573        self.decoder = decoder or Decoder(
 574            features=features_decoder,
 575            scale_factors=scale_factors[::-1],
 576            conv_block_impl=ConvBlock2d,
 577            sampler_impl=_upsampler,
 578        )
 579
 580        if use_skip_connection:
 581            self.deconv1 = Deconv2DBlock(
 582                in_channels=self.embed_dim,
 583                out_channels=features_decoder[0],
 584                use_conv_transpose=use_conv_transpose,
 585            )
 586            self.deconv2 = nn.Sequential(
 587                Deconv2DBlock(
 588                    in_channels=self.embed_dim,
 589                    out_channels=features_decoder[0],
 590                    use_conv_transpose=use_conv_transpose,
 591                ),
 592                Deconv2DBlock(
 593                    in_channels=features_decoder[0],
 594                    out_channels=features_decoder[1],
 595                    use_conv_transpose=use_conv_transpose,
 596                )
 597            )
 598            self.deconv3 = nn.Sequential(
 599                Deconv2DBlock(
 600                    in_channels=self.embed_dim,
 601                    out_channels=features_decoder[0],
 602                    use_conv_transpose=use_conv_transpose,
 603                ),
 604                Deconv2DBlock(
 605                    in_channels=features_decoder[0],
 606                    out_channels=features_decoder[1],
 607                    use_conv_transpose=use_conv_transpose,
 608                ),
 609                Deconv2DBlock(
 610                    in_channels=features_decoder[1],
 611                    out_channels=features_decoder[2],
 612                    use_conv_transpose=use_conv_transpose,
 613                )
 614            )
 615            self.deconv4 = ConvBlock2d(in_chans, features_decoder[-1])
 616        else:
 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 = Deconv2DBlock(
 623                in_channels=features_decoder[0],
 624                out_channels=features_decoder[1],
 625                use_conv_transpose=use_conv_transpose,
 626            )
 627            self.deconv3 = Deconv2DBlock(
 628                in_channels=features_decoder[1],
 629                out_channels=features_decoder[2],
 630                use_conv_transpose=use_conv_transpose,
 631            )
 632            self.deconv4 = Deconv2DBlock(
 633                in_channels=features_decoder[2],
 634                out_channels=features_decoder[3],
 635                use_conv_transpose=use_conv_transpose,
 636            )
 637
 638        self.base = ConvBlock2d(self.embed_dim, features_decoder[0])
 639        self.out_conv = nn.Conv2d(features_decoder[-1], out_channels, 1)
 640        self.deconv_out = _upsampler(
 641            scale_factor=2, in_channels=features_decoder[-1], out_channels=features_decoder[-1]
 642        )
 643        self.decoder_head = ConvBlock2d(2 * features_decoder[-1], features_decoder[-1])
 644
 645    def forward(self, x: torch.Tensor) -> torch.Tensor:
 646        """Apply the UNETR to the input data.
 647
 648        Args:
 649            x: The input tensor.
 650
 651        Returns:
 652            The UNETR output.
 653        """
 654        original_shape = x.shape[-2:]
 655
 656        # Reshape the inputs to the shape expected by the encoder
 657        # and normalize the inputs if normalization is part of the model.
 658        x, input_shape = self.preprocess(x)
 659
 660        encoder_outputs = self.encoder(x)
 661
 662        if isinstance(encoder_outputs[-1], list):
 663            # `encoder_outputs` can be arranged in only two forms:
 664            #   - either we only return the image embeddings
 665            #   - or, we return the image embeddings and the "list" of global attention layers
 666            z12, from_encoder = encoder_outputs
 667        else:
 668            z12 = encoder_outputs
 669
 670        if self.use_skip_connection:
 671            from_encoder = from_encoder[::-1]
 672            z9 = self.deconv1(from_encoder[0])
 673            z6 = self.deconv2(from_encoder[1])
 674            z3 = self.deconv3(from_encoder[2])
 675            z0 = self.deconv4(x)
 676
 677        else:
 678            z9 = self.deconv1(z12)
 679            z6 = self.deconv2(z9)
 680            z3 = self.deconv3(z6)
 681            z0 = self.deconv4(z3)
 682
 683        updated_from_encoder = [z9, z6, z3]
 684
 685        x = self.base(z12)
 686        x = self.decoder(x, encoder_inputs=updated_from_encoder)
 687        x = self.deconv_out(x)
 688
 689        x = torch.cat([x, z0], dim=1)
 690        x = self.decoder_head(x)
 691
 692        x = self.out_conv(x)
 693        if self.final_activation is not None:
 694            x = self.final_activation(x)
 695
 696        x = self.postprocess_masks(x, input_shape, original_shape)
 697        return x
 698
 699
 700class UNETR2D(UNETR):
 701    """A two-dimensional UNet Transformer using a vision transformer as encoder and a convolutional decoder.
 702    """
 703    pass
 704
 705
 706class UNETR3D(UNETRBase):
 707    """A three dimensional UNet Transformer using a vision transformer as encoder and a convolutional decoder.
 708    """
 709    def __init__(
 710        self,
 711        img_size: int = 1024,
 712        backbone: Literal["sam", "sam2", "sam3", "cellpose_sam", "mae", "scalemae", "dinov2", "dinov3"] = "sam",
 713        encoder: Optional[Union[nn.Module, str]] = "hvit_b",
 714        decoder: Optional[nn.Module] = None,
 715        out_channels: int = 1,
 716        use_sam_stats: bool = False,
 717        use_mae_stats: bool = False,
 718        use_dino_stats: bool = False,
 719        resize_input: bool = True,
 720        encoder_checkpoint: Optional[Union[str, OrderedDict]] = None,
 721        final_activation: Optional[Union[str, nn.Module]] = None,
 722        use_skip_connection: bool = False,
 723        embed_dim: Optional[int] = None,
 724        use_conv_transpose: bool = False,
 725        use_strip_pooling: bool = True,
 726        perform_range_checks: bool = True,
 727        **kwargs
 728    ):
 729        if use_skip_connection:
 730            raise NotImplementedError("The framework cannot handle skip connections atm.")
 731        if use_conv_transpose:
 732            raise NotImplementedError("It's not enabled to switch between interpolation and transposed convolutions.")
 733
 734        # Sort the `embed_dim` out
 735        embed_dim = 256 if embed_dim is None else embed_dim
 736
 737        super().__init__(
 738            img_size=img_size,
 739            backbone=backbone,
 740            encoder=encoder,
 741            decoder=decoder,
 742            out_channels=out_channels,
 743            use_sam_stats=use_sam_stats,
 744            use_mae_stats=use_mae_stats,
 745            use_dino_stats=use_dino_stats,
 746            resize_input=resize_input,
 747            encoder_checkpoint=encoder_checkpoint,
 748            final_activation=final_activation,
 749            use_skip_connection=use_skip_connection,
 750            embed_dim=embed_dim,
 751            use_conv_transpose=use_conv_transpose,
 752            perform_range_checks=perform_range_checks,
 753            **kwargs,
 754        )
 755
 756        # The 3d convolutional decoder.
 757        # First, get the important parameters for the decoder.
 758        depth = 3
 759        gain = 2
 760        features_decoder = [self.initial_features * gain ** i for i in range(depth + 1)][::-1]
 761        scale_factors = [1, 2, 2]
 762        self.out_channels = out_channels
 763
 764        # The mapping blocks.
 765        self.deconv1 = Deconv3DBlock(
 766            in_channels=embed_dim,
 767            out_channels=features_decoder[0],
 768            scale_factor=scale_factors,
 769            use_strip_pooling=use_strip_pooling,
 770        )
 771        self.deconv2 = Deconv3DBlock(
 772            in_channels=features_decoder[0],
 773            out_channels=features_decoder[1],
 774            scale_factor=scale_factors,
 775            use_strip_pooling=use_strip_pooling,
 776        )
 777        self.deconv3 = Deconv3DBlock(
 778            in_channels=features_decoder[1],
 779            out_channels=features_decoder[2],
 780            scale_factor=scale_factors,
 781            use_strip_pooling=use_strip_pooling,
 782        )
 783        self.deconv4 = Deconv3DBlock(
 784            in_channels=features_decoder[2],
 785            out_channels=features_decoder[3],
 786            scale_factor=scale_factors,
 787            use_strip_pooling=use_strip_pooling,
 788        )
 789
 790        # The core decoder block.
 791        self.decoder = decoder or Decoder(
 792            features=features_decoder,
 793            scale_factors=[scale_factors] * depth,
 794            conv_block_impl=partial(ConvBlock3dWithStrip, use_strip_pooling=use_strip_pooling),
 795            sampler_impl=Upsampler3d,
 796        )
 797
 798        # And the final upsampler to match the expected dimensions.
 799        self.deconv_out = Deconv3DBlock(  # NOTE: changed `end_up` to `deconv_out`
 800            in_channels=features_decoder[-1],
 801            out_channels=features_decoder[-1],
 802            scale_factor=scale_factors,
 803            use_strip_pooling=use_strip_pooling,
 804        )
 805
 806        # Additional conjunction blocks.
 807        self.base = ConvBlock3dWithStrip(
 808            in_channels=embed_dim,
 809            out_channels=features_decoder[0],
 810            use_strip_pooling=use_strip_pooling,
 811        )
 812
 813        # And the output layers.
 814        self.decoder_head = ConvBlock3dWithStrip(
 815            in_channels=2 * features_decoder[-1],
 816            out_channels=features_decoder[-1],
 817            use_strip_pooling=use_strip_pooling,
 818        )
 819        self.out_conv = nn.Conv3d(features_decoder[-1], out_channels, 1)
 820
 821    def forward(self, x: torch.Tensor):
 822        """Forward pass of the UNETR-3D model.
 823
 824        Args:
 825            x: Inputs of expected shape (B, C, Z, Y, X), where Z considers flexible inputs.
 826
 827        Returns:
 828            The UNETR output.
 829        """
 830        B, C, Z, H, W = x.shape
 831        original_shape = (Z, H, W)
 832
 833        # Preprocessing step
 834        x, input_shape = self.preprocess(x)
 835
 836        # Run the image encoder.
 837        curr_features = torch.stack([self.encoder(x[:, :, i])[0] for i in range(Z)], dim=2)
 838
 839        # Prepare the counterparts for the decoder.
 840        # NOTE: The section below is sequential, there's no skip connections atm.
 841        z9 = self.deconv1(curr_features)
 842        z6 = self.deconv2(z9)
 843        z3 = self.deconv3(z6)
 844        z0 = self.deconv4(z3)
 845
 846        updated_from_encoder = [z9, z6, z3]
 847
 848        # Align the features through the base block.
 849        x = self.base(curr_features)
 850        # Run the decoder
 851        x = self.decoder(x, encoder_inputs=updated_from_encoder)
 852        x = self.deconv_out(x)  # NOTE before `end_up`
 853
 854        # And the final output head.
 855        x = torch.cat([x, z0], dim=1)
 856        x = self.decoder_head(x)
 857        x = self.out_conv(x)
 858        if self.final_activation is not None:
 859            x = self.final_activation(x)
 860
 861        # Postprocess the output back to original size.
 862        x = self.postprocess_masks(x, input_shape, original_shape)
 863        return x
 864
 865#
 866#  ADDITIONAL FUNCTIONALITIES
 867#
 868
 869
 870def _strip_pooling_layers(enabled, channels) -> nn.Module:
 871    return DepthStripPooling(channels) if enabled else nn.Identity()
 872
 873
 874class DepthStripPooling(nn.Module):
 875    """@private
 876    """
 877    def __init__(self, channels: int, reduction: int = 4):
 878        """Block for strip pooling along the depth dimension (only).
 879
 880        eg. for 3D (Z > 1) - it aggregates global context across depth by adaptive avg pooling
 881        to Z=1, and then passes through a small 1x1x1 MLP, then broadcasts it back to Z to
 882        modulate the original features (using a gated residual).
 883
 884        For 2D (Z == 1): returns input unchanged (no-op).
 885
 886        Args:
 887            channels: The output channels.
 888            reduction: The reduction of the hidden layers.
 889        """
 890        super().__init__()
 891        hidden = max(1, channels // reduction)
 892        self.conv1 = nn.Conv3d(channels, hidden, kernel_size=1)
 893        self.bn1 = nn.BatchNorm3d(hidden)
 894        self.relu = nn.ReLU(inplace=True)
 895        self.conv2 = nn.Conv3d(hidden, channels, kernel_size=1)
 896
 897    def forward(self, x: torch.Tensor) -> torch.Tensor:
 898        if x.dim() != 5:
 899            raise ValueError(f"DepthStripPooling expects 5D tensors as input, got '{x.shape}'.")
 900
 901        B, C, Z, H, W = x.shape
 902        if Z == 1:  # i.e. always the case of all 2d.
 903            return x  # We simply do nothing there.
 904
 905        # We pool only along the depth dimension: i.e. target shape (B, C, 1, H, W)
 906        feat = F.adaptive_avg_pool3d(x, output_size=(1, H, W))
 907        feat = self.conv1(feat)
 908        feat = self.bn1(feat)
 909        feat = self.relu(feat)
 910        feat = self.conv2(feat)
 911        gate = torch.sigmoid(feat).expand(B, C, Z, H, W)  # Broadcast the collapsed depth context back to all slices
 912
 913        # Gated residual fusion
 914        return x * gate + x
 915
 916
 917class Deconv3DBlock(nn.Module):
 918    """@private
 919    """
 920    def __init__(
 921        self,
 922        scale_factor,
 923        in_channels,
 924        out_channels,
 925        kernel_size=3,
 926        anisotropic_kernel=True,
 927        use_strip_pooling=True,
 928    ):
 929        super().__init__()
 930        conv_block_kwargs = {
 931            "in_channels": out_channels,
 932            "out_channels": out_channels,
 933            "kernel_size": kernel_size,
 934            "padding": ((kernel_size - 1) // 2),
 935        }
 936        if anisotropic_kernel:
 937            conv_block_kwargs = _update_conv_kwargs(conv_block_kwargs, scale_factor)
 938
 939        self.block = nn.Sequential(
 940            Upsampler3d(scale_factor, in_channels, out_channels),
 941            nn.Conv3d(**conv_block_kwargs),
 942            nn.BatchNorm3d(out_channels),
 943            nn.ReLU(True),
 944            _strip_pooling_layers(enabled=use_strip_pooling, channels=out_channels),
 945        )
 946
 947    def forward(self, x):
 948        return self.block(x)
 949
 950
 951class ConvBlock3dWithStrip(nn.Module):
 952    """@private
 953    """
 954    def __init__(
 955        self, in_channels: int, out_channels: int, use_strip_pooling: bool = True, **kwargs
 956    ):
 957        super().__init__()
 958        self.block = nn.Sequential(
 959            ConvBlock3d(in_channels, out_channels, **kwargs),
 960            _strip_pooling_layers(enabled=use_strip_pooling, channels=out_channels),
 961        )
 962
 963    def forward(self, x):
 964        return self.block(x)
 965
 966
 967class SingleDeconv2DBlock(nn.Module):
 968    """@private
 969    """
 970    def __init__(self, scale_factor, in_channels, out_channels):
 971        super().__init__()
 972        self.block = nn.ConvTranspose2d(in_channels, out_channels, kernel_size=2, stride=2, padding=0, output_padding=0)
 973
 974    def forward(self, x):
 975        return self.block(x)
 976
 977
 978class SingleConv2DBlock(nn.Module):
 979    """@private
 980    """
 981    def __init__(self, in_channels, out_channels, kernel_size):
 982        super().__init__()
 983        self.block = nn.Conv2d(
 984            in_channels, out_channels, kernel_size=kernel_size, stride=1, padding=((kernel_size - 1) // 2)
 985        )
 986
 987    def forward(self, x):
 988        return self.block(x)
 989
 990
 991class Conv2DBlock(nn.Module):
 992    """@private
 993    """
 994    def __init__(self, in_channels, out_channels, kernel_size=3):
 995        super().__init__()
 996        self.block = nn.Sequential(
 997            SingleConv2DBlock(in_channels, out_channels, kernel_size),
 998            nn.BatchNorm2d(out_channels),
 999            nn.ReLU(True)
1000        )
1001
1002    def forward(self, x):
1003        return self.block(x)
1004
1005
1006class Deconv2DBlock(nn.Module):
1007    """@private
1008    """
1009    def __init__(self, in_channels, out_channels, kernel_size=3, use_conv_transpose=True):
1010        super().__init__()
1011        _upsampler = SingleDeconv2DBlock if use_conv_transpose else Upsampler2d
1012        self.block = nn.Sequential(
1013            _upsampler(scale_factor=2, in_channels=in_channels, out_channels=out_channels),
1014            SingleConv2DBlock(out_channels, out_channels, kernel_size),
1015            nn.BatchNorm2d(out_channels),
1016            nn.ReLU(True)
1017        )
1018
1019    def forward(self, x):
1020        return self.block(x)