# Source code for train_tools.models.MEDIARFormer

    import torch
    import torch.nn as nn

    from segmentation_models_pytorch import MAnet
    from segmentation_models_pytorch.base.modules import Activation

    __all__ = ["MEDIARFormer"]

    class MEDIARFormer(MAnet):
        """MEDIAR-Former Model"""

        def __init__(
            self,
            encoder_name="mit_b5",  # Default encoder
            encoder_weights="imagenet",  # Pre-trained weights
            decoder_channels=(1024, 512, 256, 128, 64),  # Decoder configuration
            decoder_pab_channels=256,  # Decoder Pyramid Attention Block channels
            in_channels=3,  # Number of input channels
            classes=3,  # Number of output classes
        ):
            # Initialize the MAnet model with provided parameters
            super().__init__(
                encoder_name=encoder_name,
                encoder_weights=encoder_weights,
                decoder_channels=decoder_channels,
                decoder_pab_channels=decoder_pab_channels,
                in_channels=in_channels,
                classes=classes,
            )

            # Remove the default segmentation head as it's not used in this architecture

            self.segmentation_head = None

            # Modify all activation functions in the encoder and decoder from ReLU to Mish
            _convert_activations(self.encoder, nn.ReLU, nn.Mish(inplace=True))
            _convert_activations(self.decoder, nn.ReLU, nn.Mish(inplace=True))

            # Add custom segmentation heads for different segmentation tasks

            self.cellprob_head = DeepSegmentationHead(
                in_channels=decoder_channels[-1], out_channels=1
            )

            self.gradflow_head = DeepSegmentationHead(
                in_channels=decoder_channels[-1], out_channels=2
            )

        def forward(self, x):
            """Forward pass through the network"""
            # Ensure the input shape is correct
            self.check_input_shape(x)

            # Encode the input and then decode it
            features = self.encoder(x)
            decoder_output = self.decoder(*features)

            # Generate masks for cell probability and gradient flows
            cellprob_mask = self.cellprob_head(decoder_output)
            gradflow_mask = self.gradflow_head(decoder_output)

            # Concatenate the masks for output
            masks = torch.cat([gradflow_mask, cellprob_mask], dim=1)

            return masks

    class DeepSegmentationHead(nn.Sequential):
        """Custom segmentation head for generating specific masks"""

        def __init__(
            self, in_channels, out_channels, kernel_size=3, activation=None, upsampling=1
        ):
            # Define a sequence of layers for the segmentation head
            layers = [
                nn.Conv2d(
                    in_channels,
                    in_channels // 2,
                    kernel_size=kernel_size,
                    padding=kernel_size // 2,
                ),
                nn.Mish(inplace=True),
                nn.BatchNorm2d(in_channels // 2),
                nn.Conv2d(
                    in_channels // 2,
                    out_channels,
                    kernel_size=kernel_size,
                    padding=kernel_size // 2,
                ),
                nn.UpsamplingBilinear2d(scale_factor=upsampling)
                if upsampling > 1
                else nn.Identity(),
                Activation(activation) if activation else nn.Identity(),
            ]
            super().__init__(*layers)

    def _convert_activations(module, from_activation, to_activation):
        """Recursively convert activation functions in a module"""
        for name, child in module.named_children():
            if isinstance(child, from_activation):
                setattr(module, name, to_activation)
            else:
                _convert_activations(child, from_activation, to_activation)