o  i/@sddlmZddlmZddlmZddlZddlm Z ddl m Z m Z ddl mZddlmZdd lmZdd lmZdd lmZdd lmZGd ddeZGdddeZGdddeZdS)) annotations)abstractmethod)SequenceN)Tensor) complex_absconvert_to_tensor_complex)root_sum_of_squares)NdarrayOrTensor)ifftn_centered)RandomizableTransform)TransformBackends)convert_to_tensorc@s8eZdZdZ  ddd d ZedddZdddZdS) KspaceMaskaK A basic class for under-sampling mask setup. It provides common features for under-sampling mask generators. For example, RandomMaskFunc and EquispacedMaskFunc (two mask transform objects defined right after this module) both inherit MaskFunc to properly setup properties like the acceleration factor. Tcenter_fractionsSequence[float] accelerations spatial_dimsint is_complexboolcCs4t|t|kr td||_||_||_||_dS)a Args: center_fractions: Fraction of low-frequency columns to be retained. If multiple values are provided, then one of these numbers is chosen uniformly each time. accelerations: Amount of under-sampling. This should have the same length as center_fractions. If multiple values are provided, then one of these is chosen uniformly each time. spatial_dims: Number of spatial dims (e.g., it's 2 for a 2D data; it's also 2 for pseudo-3D datasets like the fastMRI dataset). The last spatial dim is selected for sampling. For the fastMRI dataset, k-space has the form (...,num_slices,num_coils,H,W) and sampling is done along W. For a general 3D data with the shape (...,num_coils,H,W,D), sampling is done along D. is_complex: if True, then the last dimension will be reserved for real/imaginary parts. zONumber of center fractions should match number of accelerationsN)len ValueErrorrrrr)selfrrrrrl/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/apps/reconstruction/transforms/array.py__init__'s zKspaceMask.__init__kspacer returnSequence[Tensor]cCst)a This is an extra instance to allow for defining new mask generators. For creating other mask transforms, define a new class and simply override __call__. See an example of this in :py:class:`monai.apps.reconstruction.transforms.array.RandomKspacemask`. Args: kspace: The input k-space data. The shape is (...,num_coils,H,W,2) for complex 2D inputs and (...,num_coils,H,W,D) for real 3D data. )NotImplementedError)rrrrr__call__Js zKspaceMask.__call__cCs0|jdt|j}|j|}|j|}||fS)a If multiple values are provided for center_fractions and accelerations, this function selects one value uniformly for each training/test sample. Returns: A tuple containing (1) center_fraction: chosen fraction of center kspace lines to exclude from under-sampling (2) acceleration: chosen acceleration factor r)Rrandintrrr)rchoicecenter_fraction accelerationrrrrandomize_choose_accelerationYs  z(KspaceMask.randomize_choose_accelerationN)rT)rrrrrrrrrr rr)rr)__name__ __module__ __qualname____doc__rrr!r'rrrrrs  # rc@"eZdZdZejgZd ddZdS) RandomKspaceMaska This k-space mask transform under-samples the k-space according to a random sampling pattern. Precisely, it uniformly selects a subset of columns from the input k-space data. If the k-space data has N columns, the mask picks out: 1. N_low_freqs = (N * center_fraction) columns in the center corresponding to low-frequencies 2. The other columns are selected uniformly at random with a probability equal to: prob = (N / acceleration - N_low_freqs) / (N - N_low_freqs). This ensures that the expected number of columns selected is equal to (N / acceleration) It is possible to use multiple center_fractions and accelerations, in which case one possible (center_fraction, acceleration) is chosen uniformly at random each time the transform is called. Example: If accelerations = [4, 8] and center_fractions = [0.08, 0.04], then there is a 50% probability that 4-fold acceleration with 8% center fraction is selected and a 50% probability that 8-fold acceleration with 4% center fraction is selected. Modified and adopted from: https://github.com/facebookresearch/fastMRI/tree/master/fastmri rr rrcCst|}|j}|d}|jr|d}|\}}tt||}|||||}|jj|d|k} ||dd} d| | | |<dd|D} |jrR|| d<n|| d<t| j | t j } | |} t| } | |_ ttt| |j|jd }tt||j dd }| |fS) a Args: kspace: The input k-space data. The shape is (...,num_coils,H,W,2) for complex 2D inputs and (...,num_coils,H,W,D) for real 3D data. The last spatial dim is selected for sampling. For the fastMRI dataset, k-space has the form (...,num_slices,num_coils,H,W) and sampling is done along W. For a general 3D data with the shape (...,num_coils,H,W,D), sampling is done along D. Returns: A tuple containing (1) the under-sampled kspace (2) absolute value of the inverse fourier of the under-sampled kspace )sizerTcSg|]}dqSr2r.0_rrr z-RandomKspaceMask.__call__..rr spatial_dim)rshaperr'rroundr"uniformr reshapeastypenpfloat32maskrr rr)rrkspace_t spatial_sizenum_colsr%r& num_low_freqsprobrDpad mask_shapemasked masked_kspacemasked_kspace_ifftmasked_kspace_ifft_rssrrrr!s4  zRandomKspaceMask.__call__Nr(r)r*r+r,r TORCHbackendr!rrrrr.ksr.c@r-) EquispacedKspaceMaska This k-space mask transform under-samples the k-space according to an equi-distant sampling pattern. Precisely, it selects an equi-distant subset of columns from the input k-space data. If the k-space data has N columns, the mask picks out: 1. N_low_freqs = (N * center_fraction) columns in the center corresponding to low-frequencies 2. The other columns are selected with equal spacing at a proportion that reaches the desired acceleration rate taking into consideration the number of low frequencies. This ensures that the expected number of columns selected is equal to (N / acceleration) It is possible to use multiple center_fractions and accelerations, in which case one possible (center_fraction, acceleration) is chosen uniformly at random each time the EquispacedMaskFunc object is called. Example: If accelerations = [4, 8] and center_fractions = [0.08, 0.04], then there is a 50% probability that 4-fold acceleration with 8% center fraction is selected and a 50% probability that 8-fold acceleration with 4% center fraction is selected. Modified and adopted from: https://github.com/facebookresearch/fastMRI/tree/master/fastmri rr rrcCsRt|}|j}|d}|jr|d}|\}}tt||}tj|tjd}||dd} d|| | |<||||||} |j dt| } t | |d| } t |  tj} d|| <dd |D} |jrr|| d<n|| d<t|j| tj}||}t|}||_ttt||j|jd }tt||j dd }||fS) a Args: kspace: The input k-space data. The shape is (...,num_coils,H,W,2) for complex 2D inputs and (...,num_coils,H,W,D) for real 3D data. The last spatial dim is selected for sampling. For the fastMRI multi-coil dataset, k-space has the form (...,num_slices,num_coils,H,W) and sampling is done along W. For a general 3D data with the shape (...,num_coils,H,W,D), sampling is done along D. Returns: A tuple containing (1) the under-sampled kspace (2) absolute value of the inverse fourier of the under-sampled kspace r/r0)dtyper2rTrcSr3r4rr5rrrr8 r9z1EquispacedKspaceMask.__call__..r:r;)rr=rr'rr>rBzerosrCr"r#arangearoundrAuintr r@rDrr rr)rrrErFrGr%r&rHrDrJZadjusted_acceloffsetZ accel_samplesrKrLrMrNrOrrrr!s<  zEquispacedKspaceMask.__call__Nr(rPrrrrrSsrS) __future__rabcrcollections.abcrnumpyrBtorchrZ'monai.apps.reconstruction.complex_utilsrrZ#monai.apps.reconstruction.mri_utilsrmonai.config.type_definitionsr Zmonai.data.fft_utilsr monai.transforms.transformr monai.utils.enumsr monai.utils.type_conversionr rr.rSrrrrs         NY