mednext.py

# 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 torch.nn.functional import interpolate

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``.
            If ``True``, in training mode, the forward function will output not only the final feature map
            (from the `out_0` block), but also the feature maps that come from the intermediate up sample layers.
            In order to unify the return type, all intermediate feature maps are interpolated into the same size
            as the final feature map and stacked together (with a new dimension in the first axis) into one single tensor.
            For instance, if there are two intermediate feature maps with shapes: (1, 2, 16, 12) and
            (1, 2, 8, 6), and the final feature map has the shape (1, 2, 32, 24), then all intermediate feature maps
            will be interpolated into (1, 2, 32, 24), and the stacked tensor will has the shape (1, 3, 2, 32, 24).
            When calculating the loss, you can use torch.unbind to get all feature maps can compute the loss
            one by one with the ground truth, then do a weighted average for all losses to achieve the final loss.
        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:
            out_all = [x]
            for feature_map in ds_outputs[::-1]:
                out_all.append(interpolate(feature_map, x.shape[2:]))
            return torch.stack(out_all, dim=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)