monai-weekly 1.5.dev2444__py3-none-any.whl → 1.5.dev2446__py3-none-any.whl

monai/networks/nets/mednext.py

@@ -0,0 +1,354 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Portions of this code are derived from the original repository at:
+# https://github.com/MIC-DKFZ/MedNeXt
+# and are used under the terms of the Apache License, Version 2.0.
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+import torch
+import torch.nn as nn
+
+from monai.networks.blocks.mednext_block import MedNeXtBlock, MedNeXtDownBlock, MedNeXtOutBlock, MedNeXtUpBlock
+
+__all__ = [
+    "MedNeXt",
+    "MedNeXtSmall",
+    "MedNeXtBase",
+    "MedNeXtMedium",
+    "MedNeXtLarge",
+    "MedNext",
+    "MedNextS",
+    "MedNeXtS",
+    "MedNextSmall",
+    "MedNextB",
+    "MedNeXtB",
+    "MedNextBase",
+    "MedNextM",
+    "MedNeXtM",
+    "MedNextMedium",
+    "MedNextL",
+    "MedNeXtL",
+    "MedNextLarge",
+]
+
+
+class MedNeXt(nn.Module):
+    """
+    MedNeXt model class from paper: https://arxiv.org/pdf/2303.09975
+
+    Args:
+        spatial_dims: spatial dimension of the input data. Defaults to 3.
+        init_filters: number of output channels for initial convolution layer. Defaults to 32.
+        in_channels: number of input channels for the network. Defaults to 1.
+        out_channels: number of output channels for the network. Defaults to 2.
+        encoder_expansion_ratio: expansion ratio for encoder blocks. Defaults to 2.
+        decoder_expansion_ratio: expansion ratio for decoder blocks. Defaults to 2.
+        bottleneck_expansion_ratio: expansion ratio for bottleneck blocks. Defaults to 2.
+        kernel_size: kernel size for convolutions. Defaults to 7.
+        deep_supervision: whether to use deep supervision. Defaults to False.
+        use_residual_connection: whether to use residual connections in standard, down and up blocks. Defaults to False.
+        blocks_down: number of blocks in each encoder stage. Defaults to [2, 2, 2, 2].
+        blocks_bottleneck: number of blocks in bottleneck stage. Defaults to 2.
+        blocks_up: number of blocks in each decoder stage. Defaults to [2, 2, 2, 2].
+        norm_type: type of normalization layer. Defaults to 'group'.
+        global_resp_norm: whether to use Global Response Normalization. Defaults to False. Refer: https://arxiv.org/abs/2301.00808
+    """
+
+    def __init__(
+        self,
+        spatial_dims: int = 3,
+        init_filters: int = 32,
+        in_channels: int = 1,
+        out_channels: int = 2,
+        encoder_expansion_ratio: Sequence[int] | int = 2,
+        decoder_expansion_ratio: Sequence[int] | int = 2,
+        bottleneck_expansion_ratio: int = 2,
+        kernel_size: int = 7,
+        deep_supervision: bool = False,
+        use_residual_connection: bool = False,
+        blocks_down: Sequence[int] = (2, 2, 2, 2),
+        blocks_bottleneck: int = 2,
+        blocks_up: Sequence[int] = (2, 2, 2, 2),
+        norm_type: str = "group",
+        global_resp_norm: bool = False,
+    ):
+        """
+        Initialize the MedNeXt model.
+
+        This method sets up the architecture of the model, including:
+        - Stem convolution
+        - Encoder stages and downsampling blocks
+        - Bottleneck blocks
+        - Decoder stages and upsampling blocks
+        - Output blocks for deep supervision (if enabled)
+        """
+        super().__init__()
+
+        self.do_ds = deep_supervision
+        assert spatial_dims in [2, 3], "`spatial_dims` can only be 2 or 3."
+        spatial_dims_str = f"{spatial_dims}d"
+        enc_kernel_size = dec_kernel_size = kernel_size
+
+        if isinstance(encoder_expansion_ratio, int):
+            encoder_expansion_ratio = [encoder_expansion_ratio] * len(blocks_down)
+
+        if isinstance(decoder_expansion_ratio, int):
+            decoder_expansion_ratio = [decoder_expansion_ratio] * len(blocks_up)
+
+        conv = nn.Conv2d if spatial_dims_str == "2d" else nn.Conv3d
+
+        self.stem = conv(in_channels, init_filters, kernel_size=1)
+
+        enc_stages = []
+        down_blocks = []
+
+        for i, num_blocks in enumerate(blocks_down):
+            enc_stages.append(
+                nn.Sequential(
+                    *[
+                        MedNeXtBlock(
+                            in_channels=init_filters * (2**i),
+                            out_channels=init_filters * (2**i),
+                            expansion_ratio=encoder_expansion_ratio[i],
+                            kernel_size=enc_kernel_size,
+                            use_residual_connection=use_residual_connection,
+                            norm_type=norm_type,
+                            dim=spatial_dims_str,
+                            global_resp_norm=global_resp_norm,
+                        )
+                        for _ in range(num_blocks)
+                    ]
+                )
+            )
+
+            down_blocks.append(
+                MedNeXtDownBlock(
+                    in_channels=init_filters * (2**i),
+                    out_channels=init_filters * (2 ** (i + 1)),
+                    expansion_ratio=encoder_expansion_ratio[i],
+                    kernel_size=enc_kernel_size,
+                    use_residual_connection=use_residual_connection,
+                    norm_type=norm_type,
+                    dim=spatial_dims_str,
+                )
+            )
+
+        self.enc_stages = nn.ModuleList(enc_stages)
+        self.down_blocks = nn.ModuleList(down_blocks)
+
+        self.bottleneck = nn.Sequential(
+            *[
+                MedNeXtBlock(
+                    in_channels=init_filters * (2 ** len(blocks_down)),
+                    out_channels=init_filters * (2 ** len(blocks_down)),
+                    expansion_ratio=bottleneck_expansion_ratio,
+                    kernel_size=dec_kernel_size,
+                    use_residual_connection=use_residual_connection,
+                    norm_type=norm_type,
+                    dim=spatial_dims_str,
+                    global_resp_norm=global_resp_norm,
+                )
+                for _ in range(blocks_bottleneck)
+            ]
+        )
+
+        up_blocks = []
+        dec_stages = []
+        for i, num_blocks in enumerate(blocks_up):
+            up_blocks.append(
+                MedNeXtUpBlock(
+                    in_channels=init_filters * (2 ** (len(blocks_up) - i)),
+                    out_channels=init_filters * (2 ** (len(blocks_up) - i - 1)),
+                    expansion_ratio=decoder_expansion_ratio[i],
+                    kernel_size=dec_kernel_size,
+                    use_residual_connection=use_residual_connection,
+                    norm_type=norm_type,
+                    dim=spatial_dims_str,
+                    global_resp_norm=global_resp_norm,
+                )
+            )
+
+            dec_stages.append(
+                nn.Sequential(
+                    *[
+                        MedNeXtBlock(
+                            in_channels=init_filters * (2 ** (len(blocks_up) - i - 1)),
+                            out_channels=init_filters * (2 ** (len(blocks_up) - i - 1)),
+                            expansion_ratio=decoder_expansion_ratio[i],
+                            kernel_size=dec_kernel_size,
+                            use_residual_connection=use_residual_connection,
+                            norm_type=norm_type,
+                            dim=spatial_dims_str,
+                            global_resp_norm=global_resp_norm,
+                        )
+                        for _ in range(num_blocks)
+                    ]
+                )
+            )
+
+        self.up_blocks = nn.ModuleList(up_blocks)
+        self.dec_stages = nn.ModuleList(dec_stages)
+
+        self.out_0 = MedNeXtOutBlock(in_channels=init_filters, n_classes=out_channels, dim=spatial_dims_str)
+
+        if deep_supervision:
+            out_blocks = [
+                MedNeXtOutBlock(in_channels=init_filters * (2**i), n_classes=out_channels, dim=spatial_dims_str)
+                for i in range(1, len(blocks_up) + 1)
+            ]
+
+            out_blocks.reverse()
+            self.out_blocks = nn.ModuleList(out_blocks)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor | Sequence[torch.Tensor]:
+        """
+        Forward pass of the MedNeXt model.
+
+        This method performs the forward pass through the model, including:
+        - Stem convolution
+        - Encoder stages and downsampling
+        - Bottleneck blocks
+        - Decoder stages and upsampling with skip connections
+        - Output blocks for deep supervision (if enabled)
+
+        Args:
+            x (torch.Tensor): Input tensor.
+
+        Returns:
+            torch.Tensor or Sequence[torch.Tensor]: Output tensor(s).
+        """
+        # Apply stem convolution
+        x = self.stem(x)
+
+        # Encoder forward pass
+        enc_outputs = []
+        for enc_stage, down_block in zip(self.enc_stages, self.down_blocks):
+            x = enc_stage(x)
+            enc_outputs.append(x)
+            x = down_block(x)
+
+        # Bottleneck forward pass
+        x = self.bottleneck(x)
+
+        # Initialize deep supervision outputs if enabled
+        if self.do_ds:
+            ds_outputs = []
+
+        # Decoder forward pass with skip connections
+        for i, (up_block, dec_stage) in enumerate(zip(self.up_blocks, self.dec_stages)):
+            if self.do_ds and i < len(self.out_blocks):
+                ds_outputs.append(self.out_blocks[i](x))
+
+            x = up_block(x)
+            x = x + enc_outputs[-(i + 1)]
+            x = dec_stage(x)
+
+        # Final output block
+        x = self.out_0(x)
+
+        # Return output(s)
+        if self.do_ds and self.training:
+            return (x, *ds_outputs[::-1])
+        else:
+            return x
+
+
+# Define the MedNeXt variants as reported in 10.48550/arXiv.2303.09975
+def create_mednext(
+    variant: str,
+    spatial_dims: int = 3,
+    in_channels: int = 1,
+    out_channels: int = 2,
+    kernel_size: int = 3,
+    deep_supervision: bool = False,
+) -> MedNeXt:
+    """
+    Factory method to create MedNeXt variants.
+
+    Args:
+        variant (str): The MedNeXt variant to create ('S', 'B', 'M', or 'L').
+        spatial_dims (int): Number of spatial dimensions. Defaults to 3.
+        in_channels (int): Number of input channels. Defaults to 1.
+        out_channels (int): Number of output channels. Defaults to 2.
+        kernel_size (int): Kernel size for convolutions. Defaults to 3.
+        deep_supervision (bool): Whether to use deep supervision. Defaults to False.
+
+    Returns:
+        MedNeXt: The specified MedNeXt variant.
+
+    Raises:
+        ValueError: If an invalid variant is specified.
+    """
+    common_args = {
+        "spatial_dims": spatial_dims,
+        "in_channels": in_channels,
+        "out_channels": out_channels,
+        "kernel_size": kernel_size,
+        "deep_supervision": deep_supervision,
+        "use_residual_connection": True,
+        "norm_type": "group",
+        "global_resp_norm": False,
+        "init_filters": 32,
+    }
+
+    if variant.upper() == "S":
+        return MedNeXt(
+            encoder_expansion_ratio=2,
+            decoder_expansion_ratio=2,
+            bottleneck_expansion_ratio=2,
+            blocks_down=(2, 2, 2, 2),
+            blocks_bottleneck=2,
+            blocks_up=(2, 2, 2, 2),
+            **common_args,  # type: ignore
+        )
+    elif variant.upper() == "B":
+        return MedNeXt(
+            encoder_expansion_ratio=(2, 3, 4, 4),
+            decoder_expansion_ratio=(4, 4, 3, 2),
+            bottleneck_expansion_ratio=4,
+            blocks_down=(2, 2, 2, 2),
+            blocks_bottleneck=2,
+            blocks_up=(2, 2, 2, 2),
+            **common_args,  # type: ignore
+        )
+    elif variant.upper() == "M":
+        return MedNeXt(
+            encoder_expansion_ratio=(2, 3, 4, 4),
+            decoder_expansion_ratio=(4, 4, 3, 2),
+            bottleneck_expansion_ratio=4,
+            blocks_down=(3, 4, 4, 4),
+            blocks_bottleneck=4,
+            blocks_up=(4, 4, 4, 3),
+            **common_args,  # type: ignore
+        )
+    elif variant.upper() == "L":
+        return MedNeXt(
+            encoder_expansion_ratio=(3, 4, 8, 8),
+            decoder_expansion_ratio=(8, 8, 4, 3),
+            bottleneck_expansion_ratio=8,
+            blocks_down=(3, 4, 8, 8),
+            blocks_bottleneck=8,
+            blocks_up=(8, 8, 4, 3),
+            **common_args,  # type: ignore
+        )
+    else:
+        raise ValueError(f"Invalid MedNeXt variant: {variant}")
+
+
+MedNext = MedNeXt
+MedNextS = MedNeXtS = MedNextSmall = MedNeXtSmall = lambda **kwargs: create_mednext("S", **kwargs)
+MedNextB = MedNeXtB = MedNextBase = MedNeXtBase = lambda **kwargs: create_mednext("B", **kwargs)
+MedNextM = MedNeXtM = MedNextMedium = MedNeXtMedium = lambda **kwargs: create_mednext("M", **kwargs)
+MedNextL = MedNeXtL = MedNextLarge = MedNeXtLarge = lambda **kwargs: create_mednext("L", **kwargs)

monai/networks/nets/vista3d.py

@@ -641,7 +641,6 @@ class ClassMappingClassify(nn.Module):
         # [b,1,feat] @ [1,feat,dim], batch dimension become class_embedding batch dimension.
         masks_embedding = class_embedding.squeeze() @ src.view(b, c, h * w * d)
         masks_embedding = masks_embedding.view(b, -1, h, w, d).transpose(0, 1)
-
         return masks_embedding, class_embedding