U Ph, @s ddlmZddlmZddlZddlmZddlmZddl m Z ddl m Z ddl mZmZmZd#d d d d ddddddZd$d d d d ddddZd d d ddddZGdddejZGdddejZGdddejZGdd d ejZGd!d"d"ejZdS)%) annotations)SequenceN)nn) functional Convolution) same_padding)ConvNormPoolRELUBATCHintSequence[int] | inttuple | str | Nonez nn.Module) spatial_dims in_channels out_channels kernel_sizeactnormreturnc Cs&t|}t||||||dd|d }|S)NF)rrrbias conv_onlypaddingrr)rrrrrrrmodrY/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/networks/blocks/localnet_block.pyget_conv_blocks r rrrrrc Cs"t|}t||||dd|d}|S)NFT)rrrrr)rrrrrrrrrget_conv_layer0sr"rrrrc Cs t|||dddddddd }|S)Nr rFT) rrrstridesrrr is_transposedroutput_paddingr)rrrrrrrget_deconv_block:s r)cs8eZdZddddddfdd Zddd d ZZS) ResidualBlockrrNoner!cslt||kr&td|d|t||||d|_t||||d|_ttj|f||_ t |_ dS)N7expecting in_channels == out_channels, got in_channels=, out_channels=rrrr) super__init__ ValueErrorr conv_blockr"convr rrrReLUreluselfrrrr __class__rrr0Ls& zResidualBlock.__init__ torch.Tensorrc Cs$||||||}|SN)r5rr3r2r7xoutrrrforward]s zResidualBlock.forward__name__ __module__ __qualname__r0r@ __classcell__rrr8rr*Jsr*cs6eZdZdddddfdd Zdddd ZZS) LocalNetResidualBlockrr+r#csXt||kr&td|d|t|||d|_ttj|f||_t |_ dS)Nr,r-rrr) r/r0r1r" conv_layerr rrrr4r5)r7rrrr8rrr0ds zLocalNetResidualBlock.__init__r:r;cCs|||||}|Sr<)r5rrHr7r>midr?rrrr@nszLocalNetResidualBlock.forwardrArrr8rrFbs rFcs<eZdZdZddddddfdd Zdd d d ZZS) LocalNetDownSampleBlocka A down-sample module that can be used for LocalNet, based on: `Weakly-supervised convolutional neural networks for multimodal image registration `_. `Label-driven weakly-supervised learning for multimodal deformable image registration `_. Adapted from: DeepReg (https://github.com/DeepRegNet/DeepReg) rrr+r!csHtt||||d|_t||||d|_ttj|fdd|_dS)a7 Args: spatial_dims: number of spatial dimensions. in_channels: number of input channels. out_channels: number of output channels. kernel_size: convolution kernel size. Raises: NotImplementedError: when ``kernel_size`` is even r.r$)rN) r/r0r r2r*residual_blockr MAXmax_poolr6r8rrr0s z LocalNetDownSampleBlock.__init__z!tuple[torch.Tensor, torch.Tensor]r;cCsV|jddD] }|ddkrtd|jq||}||}||}||fS)a Halves the spatial dimensions. A tuple of (x, mid) is returned: - x is the downsample result, in shape (batch, ``out_channels``, insize_1 / 2, insize_2 / 2, [insize_3 / 2]), - mid is the mid-level feature, in shape (batch, ``out_channels``, insize_1, insize_2, [insize_3]) Args: x: Tensor in shape (batch, ``in_channels``, insize_1, insize_2, [insize_3]) Raises: ValueError: when input spatial dimensions are not even. r$Nrz7expecting x spatial dimensions be even, got x of shape )shaper1r2rLrN)r7r>irJrrrr@s    zLocalNetDownSampleBlock.forwardrBrCrD__doc__r0r@rErrr8rrKss rKcsNeZdZdZddddddddfd d Zd d d dZd d ddZZS)LocalNetUpSampleBlocka An up-sample module that can be used for LocalNet, based on: `Weakly-supervised convolutional neural networks for multimodal image registration `_. `Label-driven weakly-supervised learning for multimodal deformable image registration `_. Adapted from: DeepReg (https://github.com/DeepRegNet/DeepReg) nearestNrstrz bool | Noner+)rrrmode align_cornersrcsptt|||d|_t|||d|_t|||d|_||dkrZtd|d|||_ ||_ ||_ dS)a Args: spatial_dims: number of spatial dimensions. in_channels: number of input channels. out_channels: number of output channels. mode: interpolation mode of the additive upsampling, default to 'nearest'. align_corners: whether to align corners for the additive upsampling, default to None. Raises: ValueError: when ``in_channels != 2 * out_channels`` rGr$z;expecting in_channels == 2 * out_channels, got in_channels=r-N) r/r0r) deconv_blockr r2rFrLr1rrVrW)r7rrrrVrWr8rrr0s&  zLocalNetUpSampleBlock.__init__r:r;cCsPtj||jdd|j|jd}|jt|jdd}tj tj |dddd}|S)Nr$)rVrWr%) split_sizedim)rZ) F interpolaterOrVrWsplitrrtorchsumstackrIrrradditive_upsamplings z)LocalNetUpSampleBlock.additive_upsamplingc Cst|jdd|jddD],\}}|d|krtd|jd|jq|||||}||}||}|||}|S)a Halves the channel and doubles the spatial dimensions. Args: x: feature to be up-sampled, in shape (batch, ``in_channels``, insize_1, insize_2, [insize_3]) mid: mid-level feature saved during down-sampling, in shape (batch, ``out_channels``, midsize_1, midsize_2, [midsize_3]) Raises: ValueError: when ``midsize != insize * 2`` r$Nz_expecting mid spatial dimensions be exactly the double of x spatial dimensions, got x of shape z, mid of shape )ziprOr1rXrbr2rL) r7r>rJrPjh0r1r2r?rrrr@s &   zLocalNetUpSampleBlock.forward)rTN)rBrCrDrRr0rbr@rErrr8rrSs #rScs@eZdZdZddddddddfd d Zd d d dZZS)LocalNetFeatureExtractorBlocka A feature-extraction module that can be used for LocalNet, based on: `Weakly-supervised convolutional neural networks for multimodal image registration `_. `Label-driven weakly-supervised learning for multimodal deformable image registration `_. Adapted from: DeepReg (https://github.com/DeepRegNet/DeepReg) r kaiming_uniformrrrUr+)rrrr initializerrcstt||||dd|_ttj|f}|jD]Z}t||r6|dkrbtj t |j q6|dkrtj t |j q6td|dq6dS)a+ Args: spatial_dims: number of spatial dimensions. in_channels: number of input channels. out_channels: number of output channels. act: activation type and arguments. Defaults to ReLU. kernel_initializer: kernel initializer. Defaults to None. N)rrrrrrizerosz initializer zA is not supported, currently supporting kaiming_uniform and zeros)r/r0r r2r CONVmodules isinstancerinitkaiming_normal_r_ as_tensorweightzeros_r1)r7rrrrrj conv_typemr8rrr0s$   z&LocalNetFeatureExtractorBlock.__init__r:r;cCs||}|S)zo Args: x: Tensor in shape (batch, ``in_channels``, insize_1, insize_2, [insize_3]) )r2r=rrrr@'s z%LocalNetFeatureExtractorBlock.forward)r rirQrrr8rrhs  rh)r r r)r ) __future__rcollections.abcrr_rtorch.nnrr\Zmonai.networks.blocksrmonai.networks.layersrmonai.networks.layers.factoriesr r r r r"r)Moduler*rFrKrSrhrrrr s&       8P