o  i@srddlmZddlZddlmZddlmZddlmZddl m Z ddl m Z ddl mZGdd d ejZdS) ) annotationsN)Tensor) complex_abs_t)root_sum_of_squares_t) VarNetBlock)ifftn_centered_tcs2eZdZdZ  ddfd d ZdddZZS)VariationalNetworkModelac The end-to-end variational network (or simply e2e-VarNet) based on Sriram et. al., "End-to-end variational networks for accelerated MRI reconstruction". It comprises several cascades each consisting of refinement and data consistency steps. The network takes in the under-sampled kspace and estimates the ground-truth reconstruction. Modified and adopted from: https://github.com/facebookresearch/fastMRI Args: coil_sensitivity_model: A convolutional model for learning coil sensitivity maps. An example is :py:class:`monai.apps.reconstruction.networks.nets.coil_sensitivity_model.CoilSensitivityModel`. refinement_model: A convolutional network used in the refinement step of e2e-VarNet. An example is :py:class:`monai.apps.reconstruction.networks.nets.complex_unet.ComplexUnet`. num_cascades: Number of cascades. Each cascade is a :py:class:`monai.apps.reconstruction.networks.blocks.varnetblock.VarNetBlock` which consists of refinement and data consistency steps. spatial_dims: number of spatial dimensions. coil_sensitivity_model nn.Modulerefinement_model num_cascadesint spatial_dimscs8t||_tfddt|D|_||_dS)Ncsg|] }ttqS)rcopydeepcopy).0ir rp/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/apps/reconstruction/networks/nets/varnet.py 6sz4VariationalNetworkModel.__init__..)super__init__r nn ModuleListrangecascadesr)selfr r rr __class__rrr-s  z VariationalNetworkModel.__init__ masked_kspacermaskreturncCsL|||}|}|jD] }|||||}q ttt||jddd}|S)ad Args: masked_kspace: The under-sampled kspace. It's a 2D kspace (B,C,H,W,2) with the last dimension being 2 (for real/imaginary parts) and C denoting the coil dimension. 3D data will have the shape (B,C,H,W,D,2). mask: The under-sampling mask with shape (1,1,1,W,1) for 2D data or (1,1,1,1,D,1) for 3D data. Returns: The reconstructed image which is the root sum of squares (rss) of the absolute value of the inverse fourier of the predicted kspace (note that rss combines coil images into one image). )r) spatial_dim)r clonerrrrr)rr"r#Zsensitivity_mapsZ kspace_predcascadeZ output_imagerrrforward9s  zVariationalNetworkModel.forward)r r )r r r r rrrr)r"rr#rr$r)__name__ __module__ __qualname____doc__rr) __classcell__rrr rrs  r) __future__rrtorch.nnrtorchrZ'monai.apps.reconstruction.complex_utilsrZ#monai.apps.reconstruction.mri_utilsrZ5monai.apps.reconstruction.networks.blocks.varnetblockr!monai.networks.blocks.fft_utils_trModulerrrrrs