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