U PhG@sRddlmZddlZddlmZddlmZddlmZmZGdddej Z dS)) annotationsN)Tensor)sensitivity_map_expandsensitivity_map_reducecsTeZdZdZddddfdd Zddddd d d Zdddddd d dZZS) VarNetBlockaQ A variational block based on Sriram et. al., "End-to-end variational networks for accelerated MRI reconstruction". It applies data consistency and refinement to the intermediate kspace and combines those results. Modified and adopted from: https://github.com/facebookresearch/fastMRI Args: refinement_model: the model used for refinement (typically a U-Net but can be any deep learning model that performs well when the input and output are in image domain (e.g., a convolutional network). spatial_dims: is 2 for 2D data and is 3 for 3D data z nn.Moduleint)refinement_model spatial_dimscsTt||_||_ttd|_ddt |dD}| dt |dS)NcSsg|]}dqS)r ).0_r r j/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/apps/reconstruction/networks/blocks/varnetblock.py (sz(VarNetBlock.__init__..zeros) super__init__modelr nn Parametertorchones dc_weightrangeregister_bufferr)selfr r Z buffer_shape __class__r rr"s  zVarNetBlock.__init__r)x ref_kspacemaskreturncCst||||j|jS)a! Applies data consistency to input x. Suppose x is an intermediate estimate of the kspace and ref_kspace is the reference under-sampled measurement. This function returns mask * (x - ref_kspace). View this as the residual between the original under-sampled kspace and the estimate given by the network. Args: x: 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). ref_kspace: original under-sampled kspace with the same shape as x. 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: Output of DC block with the same shape as x )rwhererr)rr r!r"r r rsoft_dc+szVarNetBlock.soft_dc)current_kspacer!r" sens_mapsr#cCs@||||}t|t|||jd||jd}|||}|S)a Args: current_kspace: Predicted kspace from the previous block. 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). ref_kspace: reference kspace for applying data consistency (is the under-sampled kspace in MRI reconstruction). Its shape is the same as current_kspace. 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. sens_maps: coil sensitivity maps with the same shape as current_kspace Returns: Output of VarNetBlock with the same shape as current_kspace )r )r%rrrr )rr&r!r"r'Zdc_outZrefinement_outoutputr r rforward<s zVarNetBlock.forward)r)__name__ __module__ __qualname____doc__rr%r) __classcell__r r rrrs  r) __future__rrtorch.nnrrZ-monai.apps.reconstruction.networks.nets.utilsrrModulerr r r r s