o ) i< @srddlmZddlmZddlmZddlmZddlm Z m Z m Z ddl m Z ddlmZGdd d ejZdS) ) annotations)SequenceN) UnetOutBlock)UnetrBasicBlockUnetrPrUpBlock UnetrUpBlock)ViT)ensure_tuple_repcsLeZdZdZ         d)d*fd#d$ Zd%d&Zd'd(ZZS)+UNETRz UNETR based on: "Hatamizadeh et al., UNETR: Transformers for 3D Medical Image Segmentation "  convinstanceTF in_channelsint out_channelsimg_sizeSequence[int] | int feature_size hidden_sizemlp_dim num_heads proj_typestr norm_name tuple | str conv_blockbool res_block dropout_ratefloat spatial_dimsqkv_bias save_attnreturnNonecstd| krdkstdtd||dkr!tdd|_t|| }td| |_tddt||jD|_||_ d |_ t |||j|||j|||j | | ||d |_ t | ||d d| | d |_t| ||d d d dd | | | d |_t| ||ddd dd | | | d |_t| ||ddd dd | | | d |_t| ||dd d | | d|_t| |d|dd d | | d|_t| |d|d d d | | d|_t| |d |d d | | d|_t| ||d|_d| dftddt| D|_t|j|j g|_dS)aT Args: in_channels: dimension of input channels. out_channels: dimension of output channels. img_size: dimension of input image. feature_size: dimension of network feature size. Defaults to 16. hidden_size: dimension of hidden layer. Defaults to 768. mlp_dim: dimension of feedforward layer. Defaults to 3072. num_heads: number of attention heads. Defaults to 12. proj_type: patch embedding layer type. Defaults to "conv". norm_name: feature normalization type and arguments. Defaults to "instance". conv_block: if convolutional block is used. Defaults to True. res_block: if residual block is used. Defaults to True. dropout_rate: fraction of the input units to drop. Defaults to 0.0. spatial_dims: number of spatial dims. Defaults to 3. qkv_bias: apply the bias term for the qkv linear layer in self attention block. Defaults to False. save_attn: to make accessible the attention in self attention block. Defaults to False. Examples:: # for single channel input 4-channel output with image size of (96,96,96), feature size of 32 and batch norm >>> net = UNETR(in_channels=1, out_channels=4, img_size=(96,96,96), feature_size=32, norm_name='batch') # for single channel input 4-channel output with image size of (96,96), feature size of 32 and batch norm >>> net = UNETR(in_channels=1, out_channels=4, img_size=96, feature_size=32, norm_name='batch', spatial_dims=2) # for 4-channel input 3-channel output with image size of (128,128,128), conv position embedding and instance norm >>> net = UNETR(in_channels=4, out_channels=3, img_size=(128,128,128), proj_type='conv', norm_name='instance') rz'dropout_rate should be between 0 and 1.z-hidden_size should be divisible by num_heads.rr css|] \}}||VqdS)N).0Zimg_dp_dr+r+[/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/networks/nets/unetr.py Zsz!UNETR.__init__..F) rr patch_sizerr num_layersrrclassificationr#r%r&r'r)r%rr kernel_sizestriderr") r%rr num_layerr3r4upsample_kernel_sizerr r")r%rrr3r7rr")r%rrcss|]}|dVqdS)r*Nr+)r,dr+r+r.r/sN)super__init__ ValueErrorr1r r0tuplezipZ feat_sizerr2rvitrencoder1rencoder2encoder3encoder4rdecoder5decoder4decoder3decoder2routrange proj_axeslistproj_view_shape)selfrrrrrrrrrr r"r#r%r&r' __class__r+r.r<s 1           $zUNETR.__init__cCs0|dg|j}||}||j}|S)Nr)sizerMviewpermuterK contiguous)rNxnew_viewr+r+r. proj_feats zUNETR.proj_featcCs||\}}||}|d}|||}|d}|||}|d} ||| } ||} || | } || |} || |}| ||}| |S)Nr ) r@rArBrWrCrDrErFrGrHrI)rNx_inrUhidden_states_outenc1x2enc2x3enc3x4Zenc4dec4dec3dec2dec1rIr+r+r.forwards       z UNETR.forward) r r r rrrTTrrFF) rrrrrrrrrrrrrrrrrrr r!r"r!r#r$r%rr&r!r'r!r(r))__name__ __module__ __qualname____doc__r<rWrf __classcell__r+r+rOr.r s$ $r ) __future__rcollections.abcrtorch.nnnnZ#monai.networks.blocks.dynunet_blockrZ!monai.networks.blocks.unetr_blockrrrZmonai.networks.nets.vitr monai.utilsr Moduler r+r+r+r.s