U PhY#@szddlmZddlmZddlZddlmZddlmZm Z m Z Gdddej Z Gdddej Z Gd d d ej ZdS) ) annotations)SequenceN)UnetBasicBlock UnetResBlockget_conv_layerc s>eZdZdZd dddddddddfd d Zd d ZZS) UnetrUpBlockz An upsampling module that can be used for UNETR: "Hatamizadeh et al., UNETR: Transformers for 3D Medical Image Segmentation " FintSequence[int] | int tuple | strboolNone) spatial_dims in_channels out_channels kernel_sizeupsample_kernel_size norm_name res_blockreturnc sdt|}t|||||ddd|_|rFt|||||d|d|_nt|||||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. Trstride conv_only is_transposed)rrrN)super__init__r transp_convr conv_blockr) selfr rrrrrrupsample_stride __class__V/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/networks/blocks/unetr_block.pyrs8   zUnetrUpBlock.__init__cCs*||}tj||fdd}||}|S)Nr)dim)rtorchcatr)rinpskipoutr"r"r#forwardQs  zUnetrUpBlock.forward)F__name__ __module__ __qualname____doc__rr* __classcell__r"r"r r#rs "5rc sDeZdZdZd dddddddddddd fd d Zd d ZZS)UnetrPrUpBlockz A projection upsampling module that can be used for UNETR: "Hatamizadeh et al., UNETR: Transformers for 3D Medical Image Segmentation " Frr r r r ) r rr num_layerrrrrrrrc stt|ddd|_| r| rZtfddt|D|_qtfddt|D|_n$tfddt|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 s:g|]2}ttdddtdqSTrr rrrrr)nn Sequentialrr.0irrrr rrrr"r# s( z+UnetrPrUpBlock.__init__..c s:g|]2}ttdddtdqSr3)r5r6rrr7r:r"r#r;s( c s"g|]}tdddqS)Tr)rr7)rr rrr"r#r;s N)rrrtransp_conv_initr5 ModuleListrangeblocks) rr rrr2rrrrrrr r:r#r_s8   zUnetrPrUpBlock.__init__cCs"||}|jD] }||}q|SN)r<r?)rxblkr"r"r#r*s   zUnetrPrUpBlock.forward)FFr+r"r"r r#r1Ys (kr1c s>eZdZdZd dddddddddfd d Zd d ZZS)UnetrBasicBlockz A CNN module that can be used for UNETR, based on: "Hatamizadeh et al., UNETR: Transformers for 3D Medical Image Segmentation " Frr r r r )r rrrrrrrcs@t|r&t||||||d|_nt||||||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. r4N)rrrlayerr)rr rrrrrrr r"r#rs$  zUnetrBasicBlock.__init__cCs ||Sr@)rD)rr'r"r"r#r*szUnetrBasicBlock.forward)Fr+r"r"r r#rCs "+rC) __future__rcollections.abcrr%torch.nnr5Z#monai.networks.blocks.dynunet_blockrrrModulerr1rCr"r"r"r# s   Cx