U PhG@sddlmZddlmZddlZddlmZddlmZddlm Z ddl m Z ddl m Z mZddlmZGd d d ejZdS) ) annotations)SequenceN)Tensor)root_sum_of_squares_t) ComplexUnet)$reshape_batch_channel_to_channel_dimreshape_channel_to_batch_dim)ifftn_centered_tc s~eZdZdZddddddfdd difdd d d d f dddddddddd fdd ZdddddZddddddZZS) CoilSensitivityModela| This class uses a convolutional model to learn coil sensitivity maps for multi-coil MRI reconstruction. The convolutional model is :py:class:`monai.apps.reconstruction.networks.nets.complex_unet` by default but can be specified by the user as well. Learning is done on the center of the under-sampled kspace (that region is fully sampled). The data being a (complex) 2-channel tensor is a requirement for using this model. Modified and adopted from: https://github.com/facebookresearch/fastMRI Args: spatial_dims: number of spatial dimensions. features: six integers as numbers of features. denotes number of channels in each layer. act: activation type and arguments. Defaults to LeakyReLU. norm: feature normalization type and arguments. Defaults to instance norm. bias: whether to have a bias term in convolution blocks. Defaults to True. dropout: dropout ratio. Defaults to 0.0. upsample: upsampling mode, available options are ``"deconv"``, ``"pixelshuffle"``, ``"nontrainable"``. coil_dim: coil dimension in the data conv_net: the learning model used to estimate the coil sensitivity maps. default is :py:class:`monai.apps.reconstruction.networks.nets.complex_unet`. The only requirement on the model is to have 2 as input and output number of channels. ) r @r LeakyReLUg?T)negative_slopeinplaceinstanceaffinegdeconvNintz Sequence[int]z str | tupleboolz float | tuplestrznn.Module | None) spatial_dimsfeaturesactnormbiasdropoutupsamplecoil_dimconv_netc s|t| dkr,t|||||||d|_n@dd| D} | dddkrftd| ddd| |_||_||_dS) N)rrrrrrr cSsg|] }|jqSshape).0pr#r#s/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/apps/reconstruction/networks/nets/coil_sensitivity_model.py Qsz1CoilSensitivityModel.__init__..rrr z!in_channels should be 2 but it's .)super__init__rr" parameters ValueErrorrr!) selfrrrrrrr r!r"params __class__r#r(r,7s"  zCoilSensitivityModel.__init__rztuple[int, int])maskreturncCsV|jdd}}|d|ddfr.|d7}q|d|ddfrJ|d8}q.|d|fS)a Extracts the size of the fully-sampled part of the kspace. Note that when a kspace is under-sampled, a part of its center is fully sampled. This part is called the Auto Calibration Region (ACR). ACR is used for sensitivity map computation. Args: mask: the under-sampling mask of shape (..., S, 1) where S denotes the sampling dimension Returns: A tuple containing (1) left index of the region (2) right index of the region Note: Suppose the mask is of shape (1,1,20,1). If this function returns 8,12 as left and right indices, then it means that the fully-sampled center region has size 4 starting from 8 to 12. r .Nrr$)r/r3leftrightr#r#r(get_fully_sampled_regionXs   z-CoilSensitivityModel.get_fully_sampled_region) masked_kspacer3r4c Cs||\}}||}t|}|jd|dd}|d|||ddf|d|||ddf<t||jdd}t|\}}||}t||}|t ||j d |j }|S) a Args: masked_kspace: the under-sampled kspace (which is the input measurement). Its shape is (B,C,H,W,2) for 2D data or (B,C,H,W,D,2) for 3D data. 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: predicted coil sensitivity maps with shape (B,C,H,W,2) for 2D data or (B,C,H,W,D,2) for 3D data. r5rr .NT)r is_complex) spatial_dim) r8torch zeros_liker%r rrr"rrr! unsqueeze) r/r9r3r6r7Z num_low_freqsxstartbr#r#r(forwardss  0   zCoilSensitivityModel.forward)__name__ __module__ __qualname____doc__r,r8rB __classcell__r#r#r1r(r s  $!r ) __future__rcollections.abcrr<torch.nnnnrZ#monai.apps.reconstruction.mri_utilsrZ4monai.apps.reconstruction.networks.nets.complex_unetr-monai.apps.reconstruction.networks.nets.utilsrr!monai.networks.blocks.fft_utils_tr Moduler r#r#r#r( s