U PhH@sddlmZddlZddlmZddlmZddlmZm Z m Z m Z m Z ddl mZddd d gZGd ddejZGd ddejZGd d d eZGdddejZGdd d eZdS)) annotationsN)nn) functional)RegistrationDownSampleBlockRegistrationExtractionBlockRegistrationResidualConvBlockget_conv_blockget_deconv_block) meshgrid_ijRegUNet AffineHead GlobalNetLocalNetc seZdZdZd&dddddddd d d d d fd d ZddZddZddZddddZdddddZ ddZ dddddd Z dd!d"d#Z d$d%Z ZS)'r u Class that implements an adapted UNet. This class also serve as the parent class of LocalNet and GlobalNet Reference: O. Ronneberger, P. Fischer, and T. Brox, “U-net: Convolutional networks for biomedical image segmentation,”, Lecture Notes in Computer Science, 2015, vol. 9351, pp. 234–241. https://arxiv.org/abs/1505.04597 Adapted from: DeepReg (https://github.com/DeepRegNet/DeepReg) kaiming_uniformNTFint str | Noneztuple[int] | Noneboolint | list[int]) spatial_dims in_channelsnum_channel_initialdepthout_kernel_initializerout_activation out_channelsextract_levelspooling concat_skipencode_kernel_sizesc st|s|f}t||kr$t|_|_|_|_|_|_ |_ |_ | _ | _ t| trz| gjd} t| jdkrt| _fddtjdD_tj _dS)a, Args: spatial_dims: number of spatial dims in_channels: number of input channels num_channel_initial: number of initial channels depth: input is at level 0, bottom is at level depth. out_kernel_initializer: kernel initializer for the last layer out_activation: activation at the last layer out_channels: number of channels for the output extract_levels: list, which levels from net to extract. The maximum level must equal to ``depth`` pooling: for down-sampling, use non-parameterized pooling if true, otherwise use conv concat_skip: when up-sampling, concatenate skipped tensor if true, otherwise use addition encode_kernel_sizes: kernel size for down-sampling csg|]}jd|qS)r.0dselfP/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/networks/nets/regunet.py `sz$RegUNet.__init__..N)super__init__maxAssertionErrorrrrrrrrrrr isinstancerlenrrange num_channelsminmin_extract_level build_layers) r'rrrrrrrrrrr __class__r&r)r,,s:    zRegUNet.__init__cCs||dS)N)build_encode_layersbuild_decode_layersr&r(r(r)r5oszRegUNet.build_layerscs`tfddtjD_tfddtjD_jjdjdd_dS)Ncs@g|]8}j|dkrjn j|dj|j|dqS)rr rr kernel_size)build_conv_blockrr2rr#r&r(r)r*vs z/RegUNet.build_encode_layers..csg|]}jj|dqS)channels)build_down_sampling_blockr2r#r&r(r)r*srr) r ModuleListr1r encode_convs encode_poolsbuild_bottom_blockr2 bottom_blockr&r(r&r)r8ss  zRegUNet.build_encode_layersc Cs(tt|j|||dt|j|||dSN)rrrr;)r Sequentialrrrr'rrr;r(r(r)r<szRegUNet.build_conv_blockr=cCst|j||jdS)N)rr>r)rrr)r'r>r(r(r)r?sz!RegUNet.build_down_sampling_blockrBc Cs4|j|j}tt|j|||dt|j|||dSrH)rrrrIrrrrJr(r(r)rFs zRegUNet.build_bottom_blockcsjtfddtjdjddD_tfddtjdjddD__dS)Ncs*g|]"}jj|dj|dqS)r rB)build_up_sampling_blockr2r#r&r(r)r*sz/RegUNet.build_decode_layers..r rAcs<g|]4}jjr dj|nj|j|ddqS)r"rr:)r<rr2r#r&r(r)r*s ) rrCr1rr4decode_deconvs decode_convsbuild_output_block output_blockr&r(r&r)r9s   zRegUNet.build_decode_layers nn.ModulerrreturncCst|j||dSNrrr)r rr'rrr(r(r)rKszRegUNet.build_up_sampling_block)rRcCs t|j|j|j|j|j|jdS)N)rrr2rkernel_initializer activation)rrrr2rrrr&r(r(r)rNszRegUNet.build_output_blockcCs|jdd}g}|}t|j|jD]"\}}||}||}||q$||}|g} tt|j|jD]\\} \} } | |}|j rt j ||| dgdd}n||| d}| |}| |qj|j | |d} | S)z 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"Nr dim) image_size) shapeziprDrEappendrG enumeraterLrMrtorchcatrO)r'xrZskipsencodedZ encode_convZ encode_poolskipdecodedoutsiZ decode_deconvZ decode_convoutr(r(r)forwards$   zRegUNet.forward)rNrNTFr)__name__ __module__ __qualname____doc__r,r5r8r<r?rFr9rKrNri __classcell__r(r(r6r)r s$(C cs`eZdZdddddddfdd Zedd d d d Zd d ddZddd dddZZS)r Fr list[int]rrrZ decode_sizer save_thetac st||_|dkrN||d|d}d}tjddddddgtjd}n`|dkr||d|d|d}d}tjddddddddddddg tjd}ntd|tj||d |_ | ||_ |j j j |j jj |||_t|_d S) aR Args: spatial_dims: number of spatial dimensions image_size: output spatial size decode_size: input spatial size (two or three integers depending on ``spatial_dims``) in_channels: number of input channels save_theta: whether to save the theta matrix estimation r"rr dtyper z/only support 2D/3D operation, got spatial_dims=) in_features out_featuresN)r+r,rr_tensorfloat ValueErrorrLinearfcget_reference_gridgridweightdatazero_biascopy_rrTensortheta) r'rrZrqrrrrwrxZout_initr6r(r)r,s" * zAffineHead.__init__ztuple[int] | list[int] torch.Tensor)rZrRcCs.dd|D}tjt|dd}|jtjdS)NcSsg|]}td|qS)r)r_arange)r$rYr(r(r)r*sz1AffineHead.get_reference_grid..rrXrt)r_stackr torz)rZ mesh_pointsrr(r(r)r~szAffineHead.get_reference_grid)rc Cs|t|jt|jddg}|jdkrDtd||ddd}n4|jdkrhtd||ddd}ntd|j|S) Nr r"z qij,bpq->bpijrArzqijk,bpq->bpijkzdo not support spatial_dims=)r_r`r ones_likereinsumreshaper{)r'rZ grid_paddedZ grid_warpedr(r(r)affine_transforms   zAffineHead.affine_transformzlist[torch.Tensor])rarZrRcCsV|d}|jj|jd|_|||jdd}|jrB||_| ||j}|S)Nr)devicerA) rrrr}rr[rrdetachrr)r'rarZfrrhr(r(r)ri(s zAffineHead.forward)F) rjrkrlr, staticmethodr~rrirnr(r(r6r)r s ' c sDeZdZdZddddddd d d d d d d fd d ZddZZS)r z Build GlobalNet for image registration. Reference: Hu, Yipeng, et al. "Label-driven weakly-supervised learning for multimodal deformable image registration," https://arxiv.org/abs/1711.01666 rNTFrrorrrr) rZrrrrrrrrrrrc s||D]2} | ddkrtdddd|q||_fdd|D|_| |_tj|||||||| | d d S) a Args: image_size: output displacement field spatial size spatial_dims: number of spatial dims in_channels: number of input channels num_channel_initial: number of initial channels depth: input is at level 0, bottom is at level depth. out_kernel_initializer: kernel initializer for the last layer out_activation: activation at the last layer pooling: for down-sampling, use non-parameterized pooling if true, otherwise use conv concat_skip: when up-sampling, concatenate skipped tensor if true, otherwise use addition encode_kernel_sizes: kernel size for down-sampling save_theta: whether to save the theta matrix estimation r"rz given depth z3, all input spatial dimension must be divisible by z, got input of size csg|]}|dqSr!r(r$sizerr(r)r*asz&GlobalNet.__init__..) rrrrrrrrrrN)r{rZrqrrr+r,) r'rZrrrrrrrrrrrrr6rr)r,=s(zGlobalNet.__init__cCs t|j|j|j|jd|jdS)NrArp)r rrZrqr2rrr&r(r(r)rNpszGlobalNet.build_output_block)rNTFrF)rjrkrlrmr,rNrnr(r(r6r)r 2s(3cs<eZdZd ddddddfdd Zd d d d d ZZS)AdditiveUpSampleBlocknearestNrstr bool | Nonerrrmode align_cornerscs*tt|||d|_||_||_dSrS)r+r,r deconvrr)r'rrrrrr6r(r)r,|s zAdditiveUpSampleBlock.__init__r)rarRcCspdd|jddD}||}tj|||j|jd}tjtj|j |jdddddddd}||}|S) NcSsg|] }|dqSr!r(rr(r(r)r*sz1AdditiveUpSampleBlock.forward..r")rrr ) split_sizerYrArX) r[rF interpolaterrr_sumrsplit)r'ra output_sizeZdeconvedresizedrhr(r(r)ris  ,zAdditiveUpSampleBlock.forward)rN)rjrkrlr,rirnr(r(r6r)rzs rcs`eZdZdZddddd d d dd d d d d d fdd ZdddddZddddddZZS)ra Reimplementation of 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) rNrTFrrz tuple[int]rrrr) rrrrrrrruse_additive_samplingrrrc sL| |_| |_| |_tj||||t|||||| dgdgt|d dS)a Args: spatial_dims: number of spatial dims in_channels: number of input channels num_channel_initial: number of initial channels out_kernel_initializer: kernel initializer for the last layer out_activation: activation at the last layer out_channels: number of channels for the output extract_levels: list, which levels from net to extract. The maximum level must equal to ``depth`` pooling: for down-sampling, use non-parameterized pooling if true, otherwise use conv3d use_additive_sampling: whether use additive up-sampling layer for decoding. concat_skip: when up-sampling, concatenate skipped tensor if true, otherwise use addition mode: mode for interpolation when use_additive_sampling, default is "nearest". align_corners: align_corners for interpolation when use_additive_sampling, default is None. r) rrrrrrrrrrrN)use_additive_upsamplingrrr+r,r-) r'rrrrrrrrrrrrr6r(r)r,s zLocalNet.__init__rBcCs|j|j}t|j|||dSrH)rrrrrJr(r(r)rFs zLocalNet.build_bottom_blockrPrQcCs.|jrt|j|||j|jdSt|j||dS)NrrT)rrrrrr rUr(r(r)rKsz LocalNet.build_up_sampling_block)rNrTTFrN)rjrkrlrmr,rFrKrnr(r(r6r)rs*/) __future__rr_rtorch.nnrrZ#monai.networks.blocks.regunet_blockrrrrr monai.networks.utilsr __all__Moduler r r rrr(r(r(r) s     OFH