o , i3@sddlmZddlmZddlZddlmZddlmZm Z m Z m Z gdZ Gdddej Z   ddddZeZddZZZZddZZZZddZZZZddZZZZ dS)) annotations)SequenceN) MedNeXtBlockMedNeXtDownBlockMedNeXtOutBlockMedNeXtUpBlock)MedNeXt MedNeXtSmall MedNeXtBase MedNeXtMedium MedNeXtLargeMedNextMedNextSMedNeXtS MedNextSmallMedNextBMedNeXtB MedNextBaseMedNextMMedNeXtM MedNextMediumMedNextLMedNeXtL MedNextLargecsLeZdZdZ              d&d'fdd Zd(d$d%ZZS))rak 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  Frrrrgroup spatial_dimsint init_filters in_channels out_channelsencoder_expansion_ratioSequence[int] | intdecoder_expansion_ratiobottleneck_expansion_ratio kernel_sizedeep_supervisionbooluse_residual_connection blocks_down Sequence[int]blocks_bottleneck blocks_up norm_typestrglobal_resp_normc sHt| |_|dvsJd|d |ttr%gtttr1gt dkr8tjntj}|| dd|_ g}g}t D]6\}| tj  fddt |D| t d  d d d qKt||_t||_tj  fd dt | D|_g}g}t D]@\}| t d t d td d | tj  f d dt |Dqt||_t||_t d|_| r" fddt dtdD}|t||_dSdS)aG 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) )rrz"`spatial_dims` can only be 2 or 3.d2dr)r*c s8g|]}tdddqSrr$r%expansion_ratior*r-r2dimr4)r.0_)enc_kernel_sizer&r4ir#r2spatial_dims_strr-]/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/networks/nets/mednext.py xs   z$MedNeXt.__init__..r)r$r%r9r*r-r2r:c s<g|]}tdtdtdqSr7rlenr;)r.r)dec_kernel_sizer4r#r2r@r-rArBrCs r8c sPg|]$}tdtddtddqS)rrr8rDr;) r1rFr(r4r?r#r2r@r-rArBrCs r$ n_classesr:cs"g|] }td|dqS)rrG)r)r<r?)r#r%r@rArBrCsN)super__init__do_ds isinstancer"rEnnConv2dConv3dstem enumerateappend Sequentialranger ModuleList enc_stages down_blocks bottleneckr up_blocks dec_stagesrout_0reverse out_blocks)selfr!r#r$r%r&r(r)r*r+r-r.r0r1r2r4convrVrW num_blocksrYrZr] __class__)r.r1r)rFr(r>r&r4r?r#r2r%r@r-rBrJEs            zMedNeXt.__init__x torch.Tensorreturn%torch.Tensor | Sequence[torch.Tensor]c Cs||}g}t|j|jD]\}}||}||||}q||}|jr*g}tt|j|j D]+\}\}}|jrM|t |j krM||j ||||}|||d }||}q3| |}|jru|j ru|g|dddRS|S)a 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). rN)rPziprVrWrRrXrKrQrYrZrEr]r[training) r^rcZ enc_outputsZ enc_stage down_blockZ ds_outputsr?up_blockZ dec_stagerArArBforwards&       zMedNeXt.forward)rrrrrrrrFFrrrr F)r!r"r#r"r$r"r%r"r&r'r(r'r)r"r*r"r+r,r-r,r.r/r0r"r1r/r2r3r4r,)rcrdrerf)__name__ __module__ __qualname____doc__rJrl __classcell__rArArarBr/s(rrrrFvariantr3r!r"r$r%r*r+r,rec Cs|||||ddddd }|dkr tdddddddd |S|d kr4tdd d d dddd |S|dkrHtdd d d dd dd |S|dkr\tdddddddd |Std|)a 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. Tr Fr) r!r$r%r*r+r-r2r4r#Srr)r&r(r)r.r0r1B)rrru)rururrruM)rrururu)rurururL)rrurx)rxrxrurrxzInvalid MedNeXt variant: NrA)upperr ValueError)rrr!r$r%r*r+ common_argsrArArBcreate_mednext sn     r|cKtdi|S)Nrs)rsr|kwargsrArArB_rcKr})Nrt)rtr~rrArArBr`rcKr})Nrv)rvr~rrArArBrarcKr})Nrw)rwr~rrArArBrbr)rrrrF)rrr3r!r"r$r"r%r"r*r"r+r,rer)! __future__rcollections.abcrtorchtorch.nnrMZ#monai.networks.blocks.mednext_blockrrrr__all__Modulerr|r rrrr rrrr rrrr rrrr rArArArBs&   a Q