U Phy" @sddlmZddlmZddlZddlmZddlmZddl m Z ddl m Z m Z mZmZd"d d d dd dddddd ddZd#d d d dddddZGdddejZGdddejZd d d ddddZGd d!d!ejZdS)$) annotations)SequenceN)nn) functional Convolution)ConvNormPool same_paddingRELUBATCHkaiming_uniformintzSequence[int] | intztuple[int, ...] | int | Noneztuple | str | None str | Nonez nn.Module) spatial_dims in_channels out_channels kernel_sizestridespaddingactnorm initializerreturnc Cs|dkrt|}t|||||||dd|d } ttj|f} | D]Z} t| | rB|dkrntjt | j qB|dkrtj t | j qBt d|dqB| S)NF)rrrrbias conv_onlyrrzerosz initializer zA is not supported, currently supporting kaiming_uniform and zeros)r rrCONVmodules isinstancerinitkaiming_normal_torch as_tensorweightzeros_ ValueError) rrrrrrrrr conv_block conv_typemr-X/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/networks/blocks/regunet_block.pyget_conv_blocks2     r/)rrrrrc Cs"t|}t||||dd|d}|S)NFT)rrrr)r r)rrrrrmodr-r-r.get_conv_layer?sr1cs@eZdZdZd ddddddfdd Zddd d d ZZS) RegistrationResidualConvBlockz A block with skip links and layer - norm - activation. Only changes the number of channels, the spatial size is kept same. r r)rrr num_layersrcsrt||_tfddt|D|_tfddt|D|_tddt|D|_dS)a Args: spatial_dims: number of spatial dimensions in_channels: number of input channels out_channels: number of output channels num_layers: number of layers inside the block kernel_size: kernel_size cs(g|] }t|dkrndqS)r)rrrr)r1).0irrrrr-r. ^sz:RegistrationResidualConvBlock.__init__..csg|]}ttjfqSr-)r rr5_)rrr-r.r8hscSsg|] }tqSr-)rReLUr9r-r-r.r8isN) super__init__r4r ModuleListrangelayersnormsacts)selfrrrr4r __class__r7r.r=Os  z&RegistrationResidualConvBlock.__init__ torch.TensorxrcCs\|}tt|j|j|jD]<\}\}}}||}||}||jdkrN||}||}q|S)a  Args: x: Tensor in shape (batch, ``in_channels``, insize_1, insize_2, [insize_3]) Returns: Tensor in shape (batch, ``out_channels``, insize_1, insize_2, [insize_3]), with the same spatial size as ``x`` r ) enumeratezipr@rArBr4)rCrHskipr6convrrr-r-r.forwardks $ z%RegistrationResidualConvBlock.forward)r3r __name__ __module__ __qualname____doc__r=rM __classcell__r-r-rDr.r2Is r2cs<eZdZdZdddddfdd Zddd d d ZZS) RegistrationDownSampleBlockz A down-sample module used in RegUNet to half the spatial size. The number of channels is kept same. Adapted from: DeepReg (https://github.com/DeepRegNet/DeepReg) rboolNone)rchannelspoolingrcs@t|r&ttj|fdd|_nt|||dddd|_dS)z Args: spatial_dims: number of spatial dimensions. channels: channels pooling: use MaxPool if True, strided conv if False r3)rr)rrrrrrN)r<r=r MAXlayerr/)rCrrWrXrDr-r.r=s z$RegistrationDownSampleBlock.__init__rFrGcCs>|jddD] }|ddkrtd|jq||}|S)a_ Halves the spatial dimensions and keeps the same channel. output in shape (batch, ``channels``, insize_1 / 2, insize_2 / 2, [insize_3 / 2]), Args: x: Tensor in shape (batch, ``channels``, insize_1, insize_2, [insize_3]) Raises: ValueError: when input spatial dimensions are not even. r3Nrz7expecting x spatial dimensions be even, got x of shape )shaper)rZ)rCrHr6outr-r-r.rMs   z#RegistrationDownSampleBlock.forwardrNr-r-rDr.rTsrT)rrrrc Cs t|||dddddddd }|S)Nr3rrFTr ) rrrrrrr is_transposedroutput_paddingr)rrrr0r-r-r.get_deconv_blocks r_c sHeZdZdZdddddddd d d fd d ZddddddZZS)RegistrationExtractionBlockzx The Extraction Block used in RegUNet. Extracts feature from each ``extract_levels`` and takes the average. rNnearestrz tuple[int]ztuple[int] | list[int]rstrz bool | None)rextract_levels num_channelsrkernel_initializer activationmode align_cornersc sLt||_t||_tfdd|D|_||_||_ dS)an Args: spatial_dims: number of spatial dimensions extract_levels: spatial levels to extract feature from, 0 refers to the input scale num_channels: number of channels at each scale level, List or Tuple of length equals to `depth` of the RegNet out_channels: number of output channels kernel_initializer: kernel initializer activation: kernel activation function mode: feature map interpolation mode, default to "nearest". align_corners: whether to align corners for feature map interpolation. c s$g|]}t|ddqS)N)rrrrrr)r/)r5drfrerdrrr-r.r8s z8RegistrationExtractionBlock.__init__..N) r<r=rcmax max_levelrr>r@rgrh) rCrrcrdrrerfrgrhrDrjr.r=s    z$RegistrationExtractionBlock.__init__zlist[torch.Tensor]z list[int]rF)rH image_sizercs<fddtjjD}tjtj|dddd}|S)a# Args: x: Decoded feature at different spatial levels, sorted from deep to shallow image_size: output image size Returns: Tensor of shape (batch, `out_channels`, size1, size2, size3), where (size1, size2, size3) = ``image_size`` cs4g|],\}}tj|j|jjdqS))sizergrh)F interpolaterlrgrh)r5rZlevelrmrCrHr-r.r8sz7RegistrationExtractionBlock.forward..r)dim)rJr@rcr%meanstack)rCrHrmZ feature_listr\r-rrr.rMs  z#RegistrationExtractionBlock.forward)rNraNrNr-r-rDr.r`s "+r`)r r Nrrr)r ) __future__rcollections.abcrr%rtorch.nnrroZmonai.networks.blocksrmonai.networks.layersrr r r r/r1Moduler2rTr_r`r-r-r-r. s&     "( 7/