torch_em.model.unetr

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