o  iG@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 d)fdd Zd*d$d%Zd+d'd(ZZS),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_slopeinplaceinstanceaffinegdeconvN spatial_dimsintfeatures Sequence[int]act str | tuplenormbiasbooldropout float | tupleupsamplestrcoil_dimconv_netnn.Module | Nonec s|t| durt|||||||d|_n dd| D} | dddkr3td| ddd| |_||_||_dS) N)rrrrrr r"cSsg|]}|jqSshape).0pr'r'/home/dell461/cl/sdc2/last_ska_mid/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$) selfrrrrrr r"r$r%params __class__r'r,r07s"    zCoilSensitivityModel.__init__maskrreturntuple[int, int]cCsv|jdd}}|d|ddfr|d7}|d|ddfs|d|ddfr5|d8}|d|ddfs(|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()r3r7leftrightr'r'r,get_fully_sampled_regionXs z-CoilSensitivityModel.get_fully_sampled_region masked_kspacec 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. r:rr .NT)r is_complex) spatial_dim) r=torch zeros_liker)r rrr%rrr$ unsqueeze) r3r>r7r;r<Z num_low_freqsxstartbr'r'r,forwardss  0   zCoilSensitivityModel.forward)rrrrrrrrrrr r!r"r#r$rr%r&)r7rr8r9)r>rr7rr8r)__name__ __module__ __qualname____doc__r0r=rG __classcell__r'r'r5r,r s   !r ) __future__rcollections.abcrrAtorch.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