o ( iy"@sddlmZddlmZddlZddlmZddlmZddl m Z ddl m Z m Z mZmZ   d)d*ddZ d+d,dd ZGd!d"d"ejZGd#d$d$ejZd-d%d&ZGd'd(d(ejZdS).) annotations)SequenceN)nn) functional Convolution)ConvNormPool same_paddingRELUBATCHkaiming_uniform spatial_dimsint in_channels out_channels kernel_sizeSequence[int] | intstridespaddingtuple[int, ...] | int | Noneacttuple | str | Nonenorm initializer str | Nonereturn nn.Modulec Cs|durt|}t|||||||dd|d } ttj|f} | D]-} t| | rN|dkr7tjt | j q!|dkrFtj t | j q!t d|dq!| 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_typemr1e/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/networks/blocks/regunet_block.pyget_conv_blocks4     r3c Cs"t|}t||||dd|d}|S)NFT)rr!r"r)r r)rrrrrmodr1r1r2get_conv_layer?s r5cs0eZdZdZ ddfd d ZdddZZS)RegistrationResidualConvBlockz A block with skip links and layer - norm - activation. Only changes the number of channels, the spatial size is kept same. r rrrr 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|dkr ndqS)r)rrrr)r5).0irrrrr1r2 ^sz:RegistrationResidualConvBlock.__init__..csg|] }ttjfqSr1)r rr9_)rrr1r2r<hscSsg|]}tqSr1)rReLUr=r1r1r2r<isN) super__init__r8r ModuleListrangelayersnormsacts)selfrrrr8r __class__r;r2rAOs  z&RegistrationResidualConvBlock.__init__x torch.TensorrcCs\|}tt|j|j|jD]\}\}}}||}||}||jdkr'||}||}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 ) enumerateziprDrErFr8)rGrJskipr:convrrr1r1r2forwardks $ z%RegistrationResidualConvBlock.forward)r7r ) rrrrrrr8rrrrJrKrrK__name__ __module__ __qualname____doc__rArP __classcell__r1r1rHr2r6Is r6cs,eZdZdZdfd d Zdd dZZS)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) rrchannelspoolingboolrNonecsBt|rttj|fdd|_dSt|||dddd|_dS)z Args: spatial_dims: number of spatial dimensions. channels: channels pooling: use MaxPool if True, strided conv if False r7)rr)rrrrrrN)r@rAr MAXlayerr3)rGrrYrZrHr1r2rAs  z$RegistrationDownSampleBlock.__init__rJrKcCs>|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. r7Nrz7expecting x spatial dimensions be even, got x of shape )shaper-r^)rGrJr:outr1r1r2rPs   z#RegistrationDownSampleBlock.forward)rrrYrrZr[rr\rQrRr1r1rHr2rXsrXc Cs t|||dddddddd }|S)Nr7rrFTr ) rrrrrrr! is_transposedroutput_paddingr)rrrr4r1r1r2get_deconv_blocks rccs6eZdZdZ    ddfdd ZdddZZS) RegistrationExtractionBlockzx The Extraction Block used in RegUNet. Extracts feature from each ``extract_levels`` and takes the average. rNnearestrrextract_levels tuple[int] num_channelstuple[int] | list[int]rkernel_initializerr activationmodestr align_corners bool | Nonec 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)r3)r9drkrjrhrrr1r2r<s z8RegistrationExtractionBlock.__init__..N) r@rArfmax max_levelrrBrDrlrn) rGrrfrhrrjrkrlrnrHrqr2rAs    z$RegistrationExtractionBlock.__init__rJlist[torch.Tensor] image_size list[int]rrKcs<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))sizerlrn)F interpolatersrlrn)r9r^levelrurGrJr1r2r<s z7RegistrationExtractionBlock.forward..r)dim)rMrDrfr)meanstack)rGrJruZ feature_listr`r1r{r2rPs  z#RegistrationExtractionBlock.forward)rNreN)rrrfrgrhrirrrjrrkrrlrmrnro)rJrtrurvrrKrRr1r1rHr2rds +rd)r r Nrrr)rrrrrrrrrrrrrrrrrrrr )r ) rrrrrrrrrr )rrrrrrrr ) __future__rcollections.abcrr)rtorch.nnrrxZmonai.networks.blocksrmonai.networks.layersrr r r r3r5Moduler6rXrcrdr1r1r1r2s(     (  7 /