o ) i(@s|ddlmZddlZddlmZgdZddd d ZGd d d ejZGdddeZ GdddeZ GdddejZ dS)) annotationsN) MedNeXtBlockMedNeXtDownBlockMedNeXtUpBlockMedNeXtOutBlockF spatial_dimint transposeboolcCs(|dkr |r tjStjS|rtjStjS)N)nnConvTranspose2dConv2dConvTranspose3dConv3drr re/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/networks/blocks/mednext_block.pyget_conv_layersrcs8eZdZdZ      ddfdd ZddZZS)ra MedNeXtBlock class for the MedNeXt model. Args: in_channels (int): Number of input channels. out_channels (int): Number of output channels. expansion_ratio (int): Expansion ratio for the block. Defaults to 4. kernel_size (int): Kernel size for convolutions. Defaults to 7. use_residual_connection (int): Whether to use residual connection. Defaults to True. norm_type (str): Type of normalization to use. Defaults to "group". dim (str): Dimension of the input. Can be "2d" or "3d". Defaults to "3d". global_resp_norm (bool): Whether to use global response normalization. Defaults to False. Tgroup3dF in_channelsr out_channelsexpansion_ratio kernel_sizeuse_residual_connection norm_typestrc s,t||_||_t|dkrdndd} d|dkrdnd} | |||d|d|d|_|dkr:tj||d |_n|d krQtj |g|g|dkrJdndd |_| |||ddd d |_ t |_ | |||ddd d |_ ||_|jrd||f| } tjt| dd|_tjt| dd|_dSdS)N2dr rr)r#rrrstridepaddinggroupsr) num_groups num_channelslayer)normalized_shaper)rrrr%r&T) requires_grad)super__init__do_resdimrconv1r GroupNormnorm LayerNormconv2GELUactconv3global_resp_norm Parametertorchzerosglobal_resp_betaglobal_resp_gamma) selfrrrrrrr0r9convZglobal_resp_norm_param_shape __class__rrr..s@   zMedNeXtBlock.__init__cCs|}||}||||}|jrD|jdkr$tj|dddd}n tj|dddd}||jdddd }|j|||j |}| |}|j rP||}|S) z Forward pass of the MedNeXtBlock. Args: x (torch.Tensor): Input tensor. Returns: torch.Tensor: Output tensor. r!r )T)pr0keepdim)rCrDr#)r0rFgư>) r1r7r5r3r9r0r;meanr>r=r8r/)r?xx1gxnxrrrforwardes   zMedNeXtBlock.forward)rrTrrF) rr rr rr rr rr rr __name__ __module__ __qualname____doc__r.rM __classcell__rrrArrs7rc<eZdZdZ      ddfdd ZfddZZS)ra MedNeXtDownBlock class for downsampling in the MedNeXt model. Args: in_channels (int): Number of input channels. out_channels (int): Number of output channels. expansion_ratio (int): Expansion ratio for the block. Defaults to 4. kernel_size (int): Kernel size for convolutions. Defaults to 7. use_residual_connection (bool): Whether to use residual connection. Defaults to False. norm_type (str): Type of normalization to use. Defaults to "group". dim (str): Dimension of the input. Can be "2d" or "3d". Defaults to "3d". global_resp_norm (bool): Whether to use global response normalization. Defaults to False. rrFrrrr rrrrr rr r0r9c sltj||||d|||dt|dkrdndd} ||_|r'| ||ddd|_| |||d|d|d |_dS) NFrrr0r9r!r rr"r#rrrr%r$)r-r.rresample_do_resres_convr1 r?rrrrrrr0r9r@rArrr.s,   zMedNeXtDownBlock.__init__cs(t|}|jr||}||}|S)z Forward pass of the MedNeXtDownBlock. Args: x (torch.Tensor): Input tensor. Returns: torch.Tensor: Output tensor. )r-rMrWrXr?rIrJresrArrrMs  zMedNeXtDownBlock.forwardrrFrrFrr rr rr rr rr rr r0r r9r rNrrrArrs%rcrT)ra MedNeXtUpBlock class for upsampling in the MedNeXt model. Args: in_channels (int): Number of input channels. out_channels (int): Number of output channels. expansion_ratio (int): Expansion ratio for the block. Defaults to 4. kernel_size (int): Kernel size for convolutions. Defaults to 7. use_residual_connection (bool): Whether to use residual connection. Defaults to False. norm_type (str): Type of normalization to use. Defaults to "group". dim (str): Dimension of the input. Can be "2d" or "3d". Defaults to "3d". global_resp_norm (bool): Whether to use global response normalization. Defaults to False. rrFrrrr rrrrr rr r0r9c sttj||||d|||d||_||_t|dkrdnddd} |r+| ||ddd |_| |||d|d|d |_dS) NFrUr!r rTrr#rVr$)r-r.rWr0rrXr1rYrArrr.s.   zMedNeXtUpBlock.__init__cst|}|jdkrtjj|d}ntjj|d}|jr>||}|jdkr2tjj|d}ntjj|d}||}|S)z Forward pass of the MedNeXtUpBlock. Args: x (torch.Tensor): Input tensor. Returns: torch.Tensor: Output tensor. r!)r#rr#r)r#rr#rr#r) r-rMr0r;r functionalpadrWrXrZrArrrMs   zMedNeXtUpBlock.forwardr\r]rNrrrArrs&rcs(eZdZdZfddZddZZS)rz MedNeXtOutBlock class for the output block in the MedNeXt model. Args: in_channels (int): Number of input channels. n_classes (int): Number of output classes. dim (str): Dimension of the input. Can be "2d" or "3d". cs6tt|dkr dnddd}|||dd|_dS)Nr!r rTrr#)r)r-r.rconv_out)r?r n_classesr0r@rArrr.%s zMedNeXtOutBlock.__init__cCs ||S)z Forward pass of the MedNeXtOutBlock. Args: x (torch.Tensor): Input tensor. Returns: torch.Tensor: Output tensor. )r`)r?rIrrrrM+s zMedNeXtOutBlock.forwardrNrrrArrs r)rF)rr r r ) __future__rr;torch.nnr allrModulerrrrrrrrs   cGR