o  iW@sRddlmZddlZddlmZddlmZddlmZmZGdddej Z dS)) annotationsN)Tensor)sensitivity_map_expandsensitivity_map_reducecs8eZdZdZddfdd ZdddZdddZZS) 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 refinement_model nn.Module spatial_dimsintcsTt||_||_ttd|_ddt |dD}| dt |dS)NcSsg|]}dqS)r ).0_r r w/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/apps/reconstruction/networks/blocks/varnetblock.py (sz(VarNetBlock.__init__..zeros) super__init__modelr nn Parametertorchones dc_weightrangeregister_bufferr)selfrr buffer_shape __class__r rr"s zVarNetBlock.__init__xr 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_dccurrent_kspace sens_mapscCs@||||}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)rr r r )r"rr#rr$rr%r) r(rr#rr$rr)rr%r)__name__ __module__ __qualname____doc__rr'r+ __classcell__r r r rrs  r) __future__rrtorch.nnrrZ-monai.apps.reconstruction.networks.nets.utilsrrModulerr r r rs