MRI_recon/code/Frequency-Diffusion/dataset/kspace_subsample.py

"""
2023/10/16,
preprocess kspace data with the undersampling mask in the fastMRI project.
"""

import contextlib

import numpy as np
import torch


@contextlib.contextmanager
def temp_seed(rng, seed):
    state = rng.get_state()
    rng.seed(seed)
    try:
        yield
    finally:
        rng.set_state(state)


def create_mask_for_mask_type(mask_type_str, center_fractions, accelerations):
    if mask_type_str == "random":
        return RandomMaskFunc(center_fractions, accelerations)
    elif mask_type_str == "equispaced":
        return EquispacedMaskFunc(center_fractions, accelerations)
    else:
        raise Exception(f"{mask_type_str} not supported")


## mri related
def mri_fourier_transform_2d(image, mask):
    '''
  image: input tensor [B, H, W, C]
  mask: mask tensor [H, W]
  '''
    spectrum = torch.fft.fftn(image, dim=(1, 2), norm='ortho')
    # K-space spectrum has been shifted to shift the zero-frequency component to the center of the spectrum
    spectrum = torch.fft.fftshift(spectrum, dim=(1, 2))
    # Downsample k-space
    masked_spectrum = spectrum * mask[None, :, :, None]
    return spectrum, masked_spectrum


## mri related
def mri_inver_fourier_transform_2d(spectrum):
    '''
  image: input tensor [B, H, W, C]
  '''
    spectrum = torch.fft.ifftshift(spectrum, dim=(1, 2))
    image = torch.fft.ifftn(spectrum, dim=(1, 2), norm='ortho')
    # image = torch.fft.fftshift(image, dim=[1, 2])

    return image


def add_gaussian_noise(kspace, snr):
    ### 根据SNR确定noise的放大比例
    num_pixels = kspace.shape[0]*kspace.shape[1]*kspace.shape[2]*kspace.shape[3]
    psr = torch.sum(torch.abs(kspace.real)**2)/num_pixels
    pnr = psr/(np.power(10, snr/10))
    noise_r = torch.randn_like(kspace.real)*np.sqrt(pnr)

    psim = torch.sum(torch.abs(kspace.imag)**2)/num_pixels
    pnim = psim/(np.power(10, snr/10))
    noise_im = torch.randn_like(kspace.imag)*np.sqrt(pnim)

    noise = noise_r + 1j*noise_im
    noisy_kspace = kspace + noise

    return noisy_kspace

 
# def mri_fft(raw_mri, _SNR):
#     mri = torch.tensor(raw_mri)[None, :, :, None].to(torch.float32)
#     spectrum = torch.fft.fftn(mri, dim=(1, 2), norm='ortho')
#     # K-space spectrum has been shifted to shift the zero-frequency component to the center of the spectrum
#     kspace = torch.fft.fftshift(spectrum, dim=(1, 2))

#     if _SNR > 0:
#         noisy_kspace = add_gaussian_noise(kspace, _SNR)
#     else:
#         noisy_kspace = kspace

#     noisy_mri = mri_inver_fourier_transform_2d(noisy_kspace)
#     noisy_mri = torch.sqrt(torch.real(noisy_mri)**2 + torch.imag(noisy_mri)**2)

#     return noisy_kspace[0].permute(2, 0, 1), noisy_mri[0].permute(2, 0, 1), \
#         kspace[0].permute(2, 0, 1), mri[0].permute(2, 0, 1)

def mri_fft_m4raw(lq_mri, hq_mri):
    # breakpoint()
    lq_mri = torch.tensor(lq_mri[0])[None, :, :, None].to(torch.float32)
    lq_mri_spectrum = torch.fft.fftn(lq_mri, dim=(1, 2), norm='ortho')
    lq_mri_spectrum = torch.fft.fftshift(lq_mri_spectrum, dim=(1, 2))

    lq_mri = mri_inver_fourier_transform_2d(lq_mri_spectrum)
    lq_mri = torch.cat([torch.real(lq_mri), torch.imag(lq_mri)], dim=-1)
    lq_kspace = torch.cat([torch.real(lq_mri_spectrum), torch.imag(lq_mri_spectrum)], dim=-1)


    hq_mri = torch.tensor(hq_mri[0])[None, :, :, None].to(torch.float32)
    hq_mri_spectrum = torch.fft.fftn(hq_mri, dim=(1, 2), norm='ortho')
    hq_mri_spectrum = torch.fft.fftshift(hq_mri_spectrum, dim=(1, 2))

    hq_mri = mri_inver_fourier_transform_2d(hq_mri_spectrum)
    hq_mri = torch.cat([torch.real(hq_mri), torch.imag(hq_mri)], dim=-1)
    hq_kspace = torch.cat([torch.real(hq_mri_spectrum), torch.imag(hq_mri_spectrum)], dim=-1)

    # breakpoint()
    return lq_kspace[0], lq_mri[0].permute(2, 0, 1), \
           hq_kspace[0], hq_mri[0].permute(2, 0, 1)

def mri_fft(raw_mri, _SNR):
    mri = torch.tensor(raw_mri)[None, :, :, None].to(torch.float32)
    spectrum = torch.fft.fftn(mri, dim=(1, 2), norm='ortho')
    # K-space spectrum has been shifted to shift the zero-frequency component to the center of the spectrum
    kspace = torch.fft.fftshift(spectrum, dim=(1, 2))

    if _SNR > 0:
        noisy_kspace = add_gaussian_noise(kspace, _SNR)
    else:
        noisy_kspace = kspace

    noisy_mri = mri_inver_fourier_transform_2d(noisy_kspace)
    noisy_mri = torch.cat([torch.real(noisy_mri), torch.imag(noisy_mri)], dim=-1)
    noisy_kspace = torch.cat([torch.real(noisy_kspace), torch.imag(noisy_kspace)], dim=-1)

    raw_ksapce = torch.cat([torch.real(kspace), torch.imag(kspace)], dim=-1)
    raw_mri = mri_inver_fourier_transform_2d(kspace)
    raw_mri = torch.cat([torch.real(raw_mri), torch.imag(raw_mri)], dim=-1)

    # breakpoint()
    return noisy_kspace[0], noisy_mri[0].permute(2, 0, 1), \
        raw_ksapce[0], raw_mri[0].permute(2, 0, 1)

def undersample_mri(raw_mri, _MRIDOWN, _SNR):
    mri = torch.tensor(raw_mri)[None, :, :, None].to(torch.float32)
    raw_spectrum = torch.fft.fftn(mri, dim=(1, 2), norm='ortho')
    # K-space spectrum has been shifted to shift the zero-frequency component to the center of the spectrum
    raw_kspace = torch.fft.fftshift(raw_spectrum, dim=(1, 2))

    if not _MRIDOWN == "0X":
        if _MRIDOWN == "4X":
            mask_type_str, center_fraction, MRIDOWN = "random", 0.1, 4
        elif _MRIDOWN == "8X":
            mask_type_str, center_fraction, MRIDOWN = "equispaced", 0.04, 8

        ff = create_mask_for_mask_type(mask_type_str, [center_fraction], [MRIDOWN]) ## 0.2 for MRIDOWN=2, 0.1 for MRIDOWN=4, 0.04 for MRIDOWN=8
        shape = [240, 240, 1]
        mask = ff(shape, seed=1337)
        mask = mask[:, :, 0] # [1, 240]
        
        ### under-sample the kspace data.
        kspace, masked_kspace = mri_fourier_transform_2d(mri, mask)

        ### add low-field noise to the kspace data.
        if _SNR > 0:
            noisy_kspace = add_gaussian_noise(masked_kspace, _SNR)
        else:
            noisy_kspace = masked_kspace

    else:
        if _SNR > 0:
            noisy_kspace = add_gaussian_noise(raw_kspace, _SNR)
        else:
            noisy_kspace = raw_kspace
    
        mask = torch.ones([1,240])
        # breakpoint()
    ### conver the corrupted kspace data back to noisy MRI image.
    noisy_mri = mri_inver_fourier_transform_2d(noisy_kspace)
    noisy_mri = torch.cat([torch.real(noisy_mri), torch.imag(noisy_mri)], dim=-1)

    noisy_kspace_ = torch.cat([torch.real(noisy_kspace), torch.imag(noisy_kspace)], dim=-1)

    raw_mri = mri_inver_fourier_transform_2d(raw_kspace)
    raw_mri = torch.cat([torch.real(raw_mri), torch.imag(raw_mri)], dim=-1)
    raw_kspace = torch.cat([torch.real(raw_kspace), torch.imag(raw_kspace)], dim=-1)

    return noisy_kspace_[0], noisy_mri[0].permute(2, 0, 1), \
        raw_kspace[0], raw_mri[0].permute(2, 0, 1), mask.unsqueeze(-1)

# def undersample_mri(raw_mri, _MRIDOWN, _SNR):
#     mri = torch.tensor(raw_mri)[None, :, :, None].to(torch.float32)
#     if _MRIDOWN == "4X":
#         mask_type_str, center_fraction, MRIDOWN = "random", 0.1, 4
#     elif _MRIDOWN == "8X":
#         mask_type_str, center_fraction, MRIDOWN = "equispaced", 0.04, 8

#     ff = create_mask_for_mask_type(mask_type_str, [center_fraction], [MRIDOWN]) ## 0.2 for MRIDOWN=2, 0.1 for MRIDOWN=4, 0.04 for MRIDOWN=8

#     shape = [240, 240, 1]
#     mask = ff(shape, seed=1337)
#     mask = mask[:, :, 0] # [1, 240]
#     # print("mask:", mask.shape)
#     # print("original MRI:", mri)

#     # print("original MRI:", mri.shape)
#     ### under-sample the kspace data.
#     kspace, masked_kspace = mri_fourier_transform_2d(mri, mask)
#     ### add low-field noise to the kspace data.
#     if _SNR > 0:
#         noisy_kspace = add_gaussian_noise(masked_kspace, _SNR)
#     else:
#         noisy_kspace = masked_kspace

#     ### conver the corrupted kspace data back to noisy MRI image.
#     noisy_mri = mri_inver_fourier_transform_2d(noisy_kspace)
#     noisy_mri = torch.sqrt(torch.real(noisy_mri)**2 + torch.imag(noisy_mri)**2)

#     return noisy_kspace[0].permute(2, 0, 1), noisy_mri[0].permute(2, 0, 1), \
#         kspace[0].permute(2, 0, 1), mri[0].permute(2, 0, 1), mask.unsqueeze(-1)



class MaskFunc(object):
    """
    An object for GRAPPA-style sampling masks.

    This crates a sampling mask that densely samples the center while
    subsampling outer k-space regions based on the undersampling factor.
    """

    def __init__(self, center_fractions, accelerations):
        """
        Args:
            center_fractions (List[float]): Fraction of low-frequency columns to be
                retained. If multiple values are provided, then one of these
                numbers is chosen uniformly each time. 
            accelerations (List[int]): 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.
        """
        if len(center_fractions) != len(accelerations):
            raise ValueError(
                "Number of center fractions should match number of accelerations"
            )

        self.center_fractions = center_fractions
        self.accelerations = accelerations
        self.rng = np.random

    def choose_acceleration(self):
        """Choose acceleration based on class parameters."""
        choice = self.rng.randint(0, len(self.accelerations))
        center_fraction = self.center_fractions[choice]
        acceleration = self.accelerations[choice]

        return center_fraction, acceleration


class RandomMaskFunc(MaskFunc):
    """
    RandomMaskFunc creates a sub-sampling mask of a given shape.

    The mask 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 RandomMaskFunc object is called.

    For 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.
    """

    def __call__(self, shape, seed=None):
        """
        Create the mask.

        Args:
            shape (iterable[int]): The shape of the mask to be created. The
                shape should have at least 3 dimensions. Samples are drawn
                along the second last dimension.
            seed (int, optional): Seed for the random number generator. Setting
                the seed ensures the same mask is generated each time for the
                same shape. The random state is reset afterwards.
                
        Returns:
            torch.Tensor: A mask of the specified shape.
        """
        if len(shape) < 3:
            raise ValueError("Shape should have 3 or more dimensions")

        with temp_seed(self.rng, seed):
            num_cols = shape[-2]
            center_fraction, acceleration = self.choose_acceleration()

            # create the mask
            num_low_freqs = int(round(num_cols * center_fraction))
            prob = (num_cols / acceleration - num_low_freqs) / (
                num_cols - num_low_freqs
            )
            mask = self.rng.uniform(size=num_cols) < prob
            pad = (num_cols - num_low_freqs + 1) // 2
            mask[pad : pad + num_low_freqs] = True

            # reshape the mask
            mask_shape = [1 for _ in shape]
            mask_shape[-2] = num_cols
            mask = torch.from_numpy(mask.reshape(*mask_shape).astype(np.float32))

        return mask


class EquispacedMaskFunc(MaskFunc):
    """
    EquispacedMaskFunc creates a sub-sampling mask of a given shape.

    The mask 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 tovlow-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.

    Note that this function may not give equispaced samples (documented in
    https://github.com/facebookresearch/fastMRI/issues/54), which will require
    modifications to standard GRAPPA approaches. Nonetheless, this aspect of
    the function has been preserved to match the public multicoil data. 
    """

    def __call__(self, shape, seed):
        """
        Args:
            shape (iterable[int]): The shape of the mask to be created. The
                shape should have at least 3 dimensions. Samples are drawn
                along the second last dimension.
            seed (int, optional): Seed for the random number generator. Setting
                the seed ensures the same mask is generated each time for the
                same shape. The random state is reset afterwards.

        Returns:
            torch.Tensor: A mask of the specified shape.
        """
        if len(shape) < 3:
            raise ValueError("Shape should have 3 or more dimensions")

        with temp_seed(self.rng, seed):
            center_fraction, acceleration = self.choose_acceleration()
            num_cols = shape[-2]
            num_low_freqs = int(round(num_cols * center_fraction))

            # create the mask
            mask = np.zeros(num_cols, dtype=np.float32)
            pad = (num_cols - num_low_freqs + 1) // 2
            mask[pad : pad + num_low_freqs] = True

            # determine acceleration rate by adjusting for the number of low frequencies
            adjusted_accel = (acceleration * (num_low_freqs - num_cols)) / (
                num_low_freqs * acceleration - num_cols
            )
            offset = self.rng.randint(0, round(adjusted_accel))

            accel_samples = np.arange(offset, num_cols - 1, adjusted_accel)
            accel_samples = np.around(accel_samples).astype(np.uint)
            mask[accel_samples] = True

            # reshape the mask
            mask_shape = [1 for _ in shape]
            mask_shape[-2] = num_cols
            mask = torch.from_numpy(mask.reshape(*mask_shape).astype(np.float32))

        return mask