o ) iY#@szddlmZddlmZddlZddlmZddlmZm Z m Z Gdddej Z Gdddej Z Gd d d ej ZdS) ) annotations)SequenceN)UnetBasicBlock UnetResBlockget_conv_layerc.eZdZdZ ddfdd ZddZZS) UnetrUpBlockz An upsampling module that can be used for UNETR: "Hatamizadeh et al., UNETR: Transformers for 3D Medical Image Segmentation " F spatial_dimsint in_channels out_channels kernel_sizeSequence[int] | intupsample_kernel_size norm_name tuple | str res_blockboolreturnNonec sft|}t|||||ddd|_|r$t|||||d|d|_dSt|||||d|d|_dS)a Args: spatial_dims: number of spatial dimensions. in_channels: number of input channels. out_channels: number of output channels. kernel_size: convolution kernel size. upsample_kernel_size: convolution kernel size for transposed convolution layers. norm_name: feature normalization type and arguments. res_block: bool argument to determine if residual block is used. Tr stride conv_only is_transposed)r rrN)super__init__r transp_convr conv_blockr) selfr r r r rrrupsample_stride __class__c/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/networks/blocks/unetr_block.pyrs8     zUnetrUpBlock.__init__cCs*||}tj||fdd}||}|S)Nr)dim)rtorchcatr)rinpskipoutr#r#r$forwardQs  zUnetrUpBlock.forwardF)r r r r r r r rrrrrrrrr__name__ __module__ __qualname____doc__rr+ __classcell__r#r#r!r$rs  5rcs0eZdZdZ  ddfdd ZddZZS)UnetrPrUpBlockz A projection upsampling module that can be used for UNETR: "Hatamizadeh et al., UNETR: Transformers for 3D Medical Image Segmentation " Fr r r r num_layerr rrrrrrrrrrc stt|ddd|_| rE| r.tfddt|D|_dStfddt|D|_dStfddt|D|_dS)a Args: spatial_dims: number of spatial dimensions. in_channels: number of input channels. out_channels: number of output channels. num_layer: number of upsampling blocks. kernel_size: convolution kernel size. stride: convolution stride. upsample_kernel_size: convolution kernel size for transposed convolution layers. norm_name: feature normalization type and arguments. conv_block: bool argument to determine if convolutional block is used. res_block: bool argument to determine if residual block is used. Trc :g|]}ttdddtdqSTrr r r r rr)nn Sequentialrr.0ir rr r rrr r#r$ , z+UnetrPrUpBlock.__init__..c r5r6)r8r9rrr:r=r#r$r>r?c s"g|] }tdddqS)Tr)rr:)r r rr r#r$r>s N)rrrtransp_conv_initr8 ModuleListrangeblocks) rr r r r4r rrrrrr!r=r$r_s8      zUnetrPrUpBlock.__init__cCs"||}|jD]}||}q|SN)r@rC)rxblkr#r#r$r+s   zUnetrPrUpBlock.forward)FF)r r r r r r r4r r rrrrrrrrrrrrrr-r#r#r!r$r3Ys kr3cr)UnetrBasicBlockz A CNN module that can be used for UNETR, based on: "Hatamizadeh et al., UNETR: Transformers for 3D Medical Image Segmentation " Fr r r r r rrrrrrrrcsBt|rt||||||d|_dSt||||||d|_dS)a Args: spatial_dims: number of spatial dimensions. in_channels: number of input channels. out_channels: number of output channels. kernel_size: convolution kernel size. stride: convolution stride. norm_name: feature normalization type and arguments. res_block: bool argument to determine if residual block is used. r7N)rrrlayerr)rr r r r rrrr!r#r$rs$    zUnetrBasicBlock.__init__cCs ||SrD)rH)rr(r#r#r$r+s zUnetrBasicBlock.forwardr,)r r r r r r r rrrrrrrrrr-r#r#r!r$rGs  +rG) __future__rcollections.abcrr&torch.nnr8Z#monai.networks.blocks.dynunet_blockrrrModulerr3rGr#r#r#r$s  Cx