o , iP@sddlmZddlmZddlZddlmZddlmmZ ddl m Z ddl m Z ddlmZmZddlmZgdZGd d d ejZeZGd d d ejZeZdS) ) annotations)SequenceN) Convolution)UpSample)DVF2DDFWarp)SkipConnection)VoxelMorphUNetvoxelmorphunet VoxelMorph voxelmorphcsheZdZdZ         d4d5fd d! Zd6d&d'Zd7d*d+Zd8d,d-Zd7d.d/Zd9d2d3Z Z S):r aY The backbone network used in VoxelMorph. See :py:class:`monai.networks.nets.VoxelMorph` for more details. A concatenated pair of images (moving and fixed) is first passed through a UNet. The output of the UNet is then passed through a series of convolution blocks to produce the final prediction of the displacement field (DDF) or the stationary velocity field (DVF). In the original implementation, downsample is achieved through maxpooling, here one has the option to use either maxpooling or strided convolution for downsampling. The default is to use maxpooling as it is consistent with the original implementation. Note that for upsampling, the authors of VoxelMorph used nearest neighbor interpolation instead of transposed convolution. In this implementation, only nearest neighbor interpolation is supported in order to be consistent with the original implementation. An instance of this class can be used as a backbone network for constructing a VoxelMorph network. See the documentation of :py:class:`monai.networks.nets.VoxelMorph` for more details and an example on how to construct a VoxelMorph network. Args: spatial_dims: number of spatial dimensions. in_channels: number of channels in the input volume after concatenation of moving and fixed images. unet_out_channels: number of channels in the output of the UNet. channels: number of channels in each layer of the UNet. See the following example for more details. final_conv_channels: number of channels in each layer of the final convolution block. final_conv_act: activation type for the final convolution block. Defaults to LeakyReLU. Since VoxelMorph was originally implemented in tensorflow where the default negative slope for LeakyReLU was 0.2, we use the same default value here. kernel_size: kernel size for all convolution layers in the UNet. Defaults to 3. up_kernel_size: kernel size for all convolution layers in the upsampling path of the UNet. Defaults to 3. act: activation type for all convolution layers in the UNet. Defaults to LeakyReLU with negative slope 0.2. norm: feature normalization type and arguments for all convolution layers in the UNet. Defaults to None. dropout: dropout ratio for all convolution layers in the UNet. Defaults to 0.0 (no dropout). bias: whether to use bias in all convolution layers in the UNet. Defaults to True. use_maxpool: whether to use maxpooling in the downsampling path of the UNet. Defaults to True. Using maxpooling is the consistent with the original implementation of VoxelMorph. But one can optionally use strided convolution instead (i.e. set `use_maxpool` to False). adn_ordering: ordering of activation, dropout, and normalization. Defaults to "NDA". LEAKYRELUNTNDA spatial_dimsint in_channelsunet_out_channelschannels Sequence[int]final_conv_channelsfinal_conv_acttuple | str | None kernel_sizeSequence[int] | intup_kernel_sizeact tuple | strnormdropoutfloatbiasbool use_maxpool adn_orderingstrreturnNonecs|t|dvr td|ddkrtdt|dkr!tdt|ddkr-tdt|tr Builds the UNet structure recursively. Args: inc: number of input channels. outc: number of output channels. channels: sequence of channels for each pair of down and up layers. is_top: True if this is the top block. rr)NFr/)len_get_bottom_layer_get_down_layer _get_up_layer_get_connection_block) r-r.rr/Z next_c_inZ next_c_outupcsubblockdownup _create_blockself`/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/networks/nets/voxelmorph.pyr<}s   z.VoxelMorphUNet.__init__.._create_blockcst}t|D]!\}}|d|tj||jjjj j j d |}q|dtj||jdjj j j d |S)a Builds the final convolution blocks. Args: inc: number of input channels, should be the same as `unet_out_channels`. outc: number of output channels, should be the same as `spatial_dims`. channels: sequence of channels for each convolution layer. Note: there is no activation after the last convolution layer as per the original implementation. Z final_conv_)rrrr r"r%Zfinal_conv_outN) nn Sequential enumerate add_moduler dimensionsrrrr r"r%)r-r.rmodic)r=r>r?_create_final_convs@ z3VoxelMorphUNet.__init__.._create_final_convr1) r-rr.rrrr/r#r'r0)r-rr.rrrr'r0)super__init__ ValueErrorr2 isinstancerrDrrrrrr&upperrrr r"r$r%rrr@rAnet)r=rrrrrrrrrrr r"r$r%rH __class__r;r?rJCsV     0 zVoxelMorphUNet.__init__ down_pathr0up_pathr8cCst|t||S)a Returns the block object defining a layer of the UNet structure including the implementation of the skip between encoding (down) and decoding (up) sides of the network. Args: down_path: encoding half of the layer up_path: decoding half of the layer subblock: block defining the next layer in the network. Returns: block for this layer: `nn.Sequential(down_path, SkipConnection(subblock), up_path)` )r@rAr)r=rQrRr8r>r>r?r6s z$VoxelMorphUNet._get_connection_block out_channelsr/c Cs|js|rdnd}t|j||||j|j|j|j|j|jd }|jr>|s>|jdkr3t t j ddd|n t t j ddd|}|S)a In each down layer, the input is first downsampled using maxpooling, then passed through a convolution block, unless this is the top layer in which case the input is passed through a convolution block only without maxpooling first. Args: in_channels: number of input channels. out_channels: number of output channels. is_top: True if this is the top block. r))stridesrrrr r"r%r)rstride) r$rrDrrrr r"r%r@rA MaxPool3d MaxPool2dr=rrSr/rUrEr>r>r?r4s& zVoxelMorphUNet._get_down_layerc Cs4|j||dd}t|j||ddddd}t||S)z Bottom layer (bottleneck) in voxelmorph consists of a typical down layer followed by an upsample layer. Args: in_channels: number of input channels. out_channels: number of output channels. Fr1r) nontrainablenearestN scale_factormode interp_mode align_corners)r4rrDr@rA)r=rrSrEupsampler>r>r?r3s  z VoxelMorphUNet._get_bottom_layerc CsXd}t|j||||j|j|j|j|jd|jd }|s*t |t |j||ddddd}|S) a In each up layer, the input is passed through a convolution block before upsampled, unless this is the top layer in which case the input is passed through a convolution block only without upsampling. Args: in_channels: number of input channels. out_channels: number of output channels. is_top: True if this is the top block. rTF)rUrrrr r" is_transposedr%r)rZr[Nr\) rrDrrrr r"r%r@rArrYr>r>r?r5s8 zVoxelMorphUNet._get_up_layerconcatenated_pairs torch.TensorcCs||}|S)N)rN)r=rcxr>r>r?forwardLs zVoxelMorphUNet.forward) r rrr NrTTr)rrrrrrrrrrrrrrrrrrrrr r!r"r#r$r#r%r&r'r()rQr0rRr0r8r0r'r0)rrrSrr/r#r'r0)rrrSrr'r0)rcrdr'rd) __name__ __module__ __qualname____doc__rJr6r4r3r5rf __classcell__r>r>rOr?r s$-  ' /r cs6eZdZdZ    ddfdd ZdddZZS)r a A re-implementation of VoxelMorph framework for medical image registration as described in https://arxiv.org/pdf/1809.05231.pdf. For more details, please refer to VoxelMorph: A Learning Framework for Deformable Medical Image Registration, Guha Balakrishnan, Amy Zhao, Mert R. Sabuncu, John Guttag, Adrian V. Dalca IEEE TMI: Transactions on Medical Imaging. 2019. eprint arXiv:1809.05231. This class is intended to be a general framework, based on which a deformable image registration network can be built. Given a user-specified backbone network (e.g., UNet in the original VoxelMorph paper), this class serves as a wrapper that concatenates the input pair of moving and fixed images, passes through the backbone network, integrate the predicted stationary velocity field (DVF) from the backbone network to obtain the displacement field (DDF), and, finally, warp the moving image using the DDF. To construct a VoxelMorph network, one need to first construct a backbone network (e.g., a :py:class:`monai.networks.nets.VoxelMorphUNet`) and pass it to the constructor of :py:class:`monai.networks.nets.VoxelMorph`. The backbone network should be able to take a pair of moving and fixed images as input and produce a DVF (or DDF, details to be discussed later) as output. When `forward` is called, the input moving and fixed images are first concatenated along the channel dimension and passed through the specified backbone network to produce the prediction of the displacement field (DDF) in the non-diffeomorphic variant (i.e. when `integration_steps` is set to 0) or the stationary velocity field (DVF) in the diffeomorphic variant (i.e. when `integration_steps` is set to a positive integer). The DVF is then integrated using a scaling-and-squaring approach via a :py:class:`monai.networks.blocks.warp.DVF2DDF` module to produce the DDF. Finally, the DDF is used to warp the moving image to the fixed image using a :py:class:`monai.networks.blocks.warp.Warp` module. Optionally, the integration from DVF to DDF can be performed on reduced resolution by specifying `half_res` to be True, in which case the output DVF from the backbone network is first linearly interpolated to half resolution before integration. The output DDF is then linearly interpolated again back to full resolution before being used to warp the moving image. Args: backbone: a backbone network. integration_steps: number of integration steps used for obtaining DDF from DVF via scaling-and-squaring. Defaults to 7. If set to 0, the network will be non-diffeomorphic. half_res: whether to perform integration on half resolution. Defaults to False. spatial_dims: number of spatial dimensions, defaults to 3. Example:: from monai.networks.nets import VoxelMorphUNet, VoxelMorph # The following example construct an instance of VoxelMorph that matches the original VoxelMorph paper # https://arxiv.org/pdf/1809.05231.pdf # First, a backbone network is constructed. In this case, we use a VoxelMorphUNet as the backbone network. backbone = VoxelMorphUNet( spatial_dims=3, in_channels=2, unet_out_channels=32, channels=(16, 32, 32, 32, 32, 32), # this indicates the down block at the top takes 16 channels as # input, the corresponding up block at the top produces 32 # channels as output, the second down block takes 32 channels as # input, and the corresponding up block at the same level # produces 32 channels as output, etc. final_conv_channels=(16, 16) ) # Then, a full VoxelMorph network is constructed using the specified backbone network. net = VoxelMorph( backbone=backbone, integration_steps=7, half_res=False ) # A forward pass through the network would look something like this moving = torch.randn(1, 1, 160, 192, 224) fixed = torch.randn(1, 1, 160, 192, 224) warped, ddf = net(moving, fixed) NFrbackbone!VoxelMorphUNet | nn.Module | Noneintegration_stepsrhalf_resr#rr'r(cszt|dur |nt|ddddd|_||_||_||_|jdkr%dnd|_|jr4t|jd d d |_ t d d d |_ dS) Nr) )rqrqrqrqrq)rrrr)rrrrrrTFbilinearzeros) num_stepsr^ padding_mode)r^rv) rIrJr rmrrpro diffeomorphicrdvf2ddfrwarp)r=rmrorprrOr>r?rJs$  zVoxelMorph.__init__movingrdfixed!tuple[torch.Tensor, torch.Tensor]cCs |j|jkrtd|jd|jd|tj||gdd}|jd|jkr6td|jd|jdd|jd d|jd dkrYtd |jd dd |jd dd |jrgtj|d dddd}|j ro| |}|jr}tj|d dddd}| |||fS)NzfThe spatial shape of the moving image should be the same as the spatial shape of the fixed image. Got z and z instead.rT)dimzqThe number of channels in the output of the backbone network should be equal to the number of spatial dimensions z. Got z channels instead.r)zvThe spatial shape of the output of the backbone network should be equal to the spatial shape of the input images. Got z instead of .g? trilinearT)r]r^r`g@) shaperKrmtorchcatrrpF interpolaterwrxry)r=rzr{rer>r>r?rfs@    zVoxelMorph.forward)NrlFr) rmrnrorrpr#rrr'r()rzrdr{rdr'r|)rgrhrirjrJrfrkr>r>rOr?r TsG!r ) __future__rcollections.abcrrtorch.nnr@torch.nn.functional functionalr"monai.networks.blocks.convolutionsrmonai.networks.blocks.upsamplerZmonai.networks.blocks.warprr"monai.networks.layers.simplelayersr__all__Moduler r r r r>r>r>r?s"     7