def __init__(self,
                 img_shape=(128, 128, 128),
                 in_channels=4,
                 embed_dim=768,
                 patch_size=16,
                 num_heads=12,
                 dropout=0.0,
                 ext_layers=[3, 6, 9, 12],
                 norm='instance',
                 base_filters=16,
                 dim_linear_block=3072,
                 num_classes=4):
        """
        The UNETR implementation based on PaddlePaddle.
        The original article refers to
        Ali Hatamizadeh, Yucheng Tang, Vishwesh Nath, Dong Yang, Andriy Myronenko, Bennett Landman, Holger Roth, Daguang Xu
        UNETR: Transformers for 3D Medical Image Segmentation
        (https://arxiv.org/abs/2103.10504)
        Implementation of this model is borrowed and modified(from torch to paddle) from here:
        https://github.com/tamasino52/UNETR/blob/main/unetr.py

        Args:
            img_shape: volume shape, provided as a tuple
            in_channels: input modalities/channels
            num_classes: number of classes
            embed_dim: transformer embed dim.
            patch_size: the non-overlapping patches to be created
            num_heads: for the transformer encoder
            dropout: percentage for dropout
            ext_layers: transformer layers to use their output
            version: 'light' saves some parameters in the decoding part
            norm: batch or instance norm for the conv blocks
        """
        super().__init__()
        self.num_layers = 12
        self.in_channels = in_channels
        self.num_classes = num_classes
        self.embed_dim = embed_dim
        self.img_shape = img_shape if type(img_shape) == 'tuple' else eval(
            img_shape)
        img_shape = self.img_shape
        self.patch_size = patch_size
        self.num_heads = num_heads
        self.dropout = dropout
        self.ext_layers = ext_layers
        self.patch_dim = [int(x / patch_size) for x in img_shape]

        self.norm = nn.BatchNorm3d if norm == 'batch' else nn.InstanceNorm3D

        self.embed = Embeddings3D(
            input_dim=in_channels,
            embed_dim=embed_dim,
            cube_size=img_shape,
            patch_size=patch_size,
            dropout=dropout)

        self.transformer = TransformerEncoder(
            embed_dim,
            num_heads,
            self.num_layers,
            dropout,
            ext_layers,
            dim_linear_block=dim_linear_block)

        self.init_conv = Conv3DBlock(
            in_channels, base_filters, double=True, norm=self.norm)

        # blue blocks in Fig.1
        self.z3_blue_conv = TranspConv3DConv3D(
            in_planes=embed_dim, out_planes=base_filters * 2, layers=3)

        self.z6_blue_conv = TranspConv3DConv3D(
            in_planes=embed_dim, out_planes=base_filters * 4, layers=2)

        self.z9_blue_conv = TranspConv3DConv3D(
            in_planes=embed_dim, out_planes=base_filters * 8, layers=1)

        # Green blocks in Fig.1
        self.z12_deconv = TranspConv3DBlock(embed_dim, base_filters * 8)

        self.z9_deconv = TranspConv3DBlock(base_filters * 8, base_filters * 4)
        self.z6_deconv = TranspConv3DBlock(base_filters * 4, base_filters * 2)
        self.z3_deconv = TranspConv3DBlock(base_filters * 2, base_filters)

        # Yellow blocks in Fig.1
        self.z9_conv = Conv3DBlock(
            base_filters * 8 * 2, base_filters * 8, double=True, norm=self.norm)
        self.z6_conv = Conv3DBlock(
            base_filters * 4 * 2, base_filters * 4, double=True, norm=self.norm)
        self.z3_conv = Conv3DBlock(
            base_filters * 2 * 2, base_filters * 2, double=True, norm=self.norm)
        # out convolutions
        self.out_conv = nn.Sequential(
            # last yellow conv block
            Conv3DBlock(
                base_filters * 2, base_filters, double=True, norm=self.norm),
            # grey block, final classification layer
            nn.Conv3D(
                base_filters, num_classes, kernel_size=1, stride=1))

        The UNETR implementation based on PaddlePaddle.
        The original article refers to
        Ali Hatamizadeh, Yucheng Tang, Vishwesh Nath, Dong Yang, Andriy Myronenko, Bennett Landman, Holger Roth, Daguang Xu
        UNETR: Transformers for 3D Medical Image Segmentation
        (https://arxiv.org/abs/2103.10504)
        Implementation of this model is borrowed and modified(from torch to paddle) from here:
        https://github.com/tamasino52/UNETR/blob/main/unetr.py

        Args:
            img_shape: volume shape, provided as a tuple
            in_channels: input modalities/channels
            num_classes: number of classes
            embed_dim: transformer embed dim.
            patch_size: the non-overlapping patches to be created
            num_heads: for the transformer encoder
            dropout: percentage for dropout
            ext_layers: transformer layers to use their output
            version: 'light' saves some parameters in the decoding part
            norm: batch or instance norm for the conv blocks
        
contrib/MedicalSeg/medicalseg/models/unetr.py