o ( i,@sddlmZddlmZddlZddlmZddlmZddl m Z ddl m Z ddl mZmZmZ d'd(ddZ d)d*ddZd+ddZGdddejZGdd d ejZGd!d"d"ejZGd#d$d$ejZGd%d&d&ejZdS),) annotations)SequenceN)nn) functional Convolution) same_padding)ConvNormPoolRELUBATCH spatial_dimsint in_channels out_channels kernel_sizeSequence[int] | intacttuple | str | Nonenormreturn nn.Modulec Cs&t|}t||||||dd|d }|S)NF)rrrbias conv_onlypaddingrr)rrrrrrrmodrf/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/networks/blocks/localnet_block.pyget_conv_blocks r!c Cs"t|}t||||dd|d}|S)NFT)rrrrr)rrrrrrrrr get_conv_layer0s r"c Cs t|||dddddddd }|S)Nr rFT) rrrstridesrrr is_transposedroutput_paddingr)rrrrrrr get_deconv_block:s r(cs(eZdZdfd d Zdd d ZZS) ResidualBlockrrrrrrrNonecslt||krtd|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__rr r/Ls zResidualBlock.__init__ torch.Tensorc Cs$||||||}|SN)r4rr2r1r6xoutrrr forward]s zResidualBlock.forward rrrrrrrrrr*rr9__name__ __module__ __qualname__r/r> __classcell__rrr7r r)Jsr)cs(eZdZd fdd Zd d d ZZS)LocalNetResidualBlockrrrrrr*csXt||krtd|d|t|||d|_ttj|f||_t |_ dS)Nr+r,rrr) r.r/r0r" conv_layerr rrrr3r4)r6rrrr7rr r/ds zLocalNetResidualBlock.__init__r9cCs|||||}|Sr:)r4rrHr6r<midr=rrr r>nszLocalNetResidualBlock.forward)rrrrrrrr*r@rArrr7r rFbs rFcs,eZdZdZdfd d Zdd dZZS)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) rrrrrrrr*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.r/r!r1r)residual_blockr MAXmax_poolr5r7rr r/s z LocalNetDownSampleBlock.__init__!tuple[torch.Tensor, torch.Tensor]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 )shaper0r1rLrN)r6r<irJrrr r>s    zLocalNetDownSampleBlock.forwardr?)rrOrBrCrD__doc__r/r>rErrr7r rKss rKcs<eZdZdZ  ddfdd ZdddZdddZZS)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) nearestNrrrrmodestr align_corners bool | Nonerr*csptt|||d|_t|||d|_t|||d|_||dkr-td|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.r/r( deconv_blockr!r1rFrLr0rrVrX)r6rrrrVrXr7rr r/s&   zLocalNetUpSampleBlock.__init__r9cCsPtj||jdd|j|jd}|jt|jdd}tj tj |dddd}|S)Nr#)rVrXr$) split_sizedim)r\) F interpolaterPrVrXsplitrrtorchsumstackrIrrr additive_upsamplings z)LocalNetUpSampleBlock.additive_upsamplingc Cst|jdd|jddD]\}}|d|kr%td|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 )ziprPr0rZrdr1rL) r6r<rJrQjh0r1r2r=rrr r>s&   zLocalNetUpSampleBlock.forward)rUN) rrrrrrrVrWrXrYrr*r@)rBrCrDrSr/rdr>rErrr7r rTs #rTcs2eZdZdZ  ddfdd ZdddZZS)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_uniformrrrrrr initializerrWrr*cstt||||dd|_ttj|f}|jD]-}t||rH|dkr1tj t |j q|dkr@tj t |j qtd|dqdS)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)rrrrrrkzerosz initializer zA is not supported, currently supporting kaiming_uniform and zeros)r.r/r!r1r CONVmodules isinstancerinitkaiming_normal_ra as_tensorweightzeros_r0)r6rrrrrl conv_typemr7rr r/s    z&LocalNetFeatureExtractorBlock.__init__r9cCs||}|S)zo Args: x: Tensor in shape (batch, ``in_channels``, insize_1, insize_2, [insize_3]) )r1r;rrr r>'s z%LocalNetFeatureExtractorBlock.forward)r rk) rrrrrrrrrlrWrr*r@rRrrr7r rjs  rj)r r r)rrrrrrrrrrrrrr)r ) rrrrrrrrrr)rrrrrrrr) __future__rcollections.abcrrartorch.nnrr^Zmonai.networks.blocksrmonai.networks.layersrmonai.networks.layers.factoriesr r r r!r"r(Moduler)rFrKrTrjrrrr s(       8P