MRI_recon/code/Frequency-Diffusion/FSMNet/dataloaders/m4raw_std_dataloader.py

from __future__ import print_function, division
from typing import Dict, NamedTuple, Optional, Sequence, Tuple, Union
import sys
sys.path.append('.')
from glob import glob
import os
os.environ['OPENBLAS_NUM_THREADS'] = '1'
import numpy as np
import torch
from torch.utils.data import Dataset

import h5py
from matplotlib import pyplot as plt
from dataloaders.m4_utils import ifft2c, fft2c, complex_abs
from dataloaders.kspace_subsample import create_mask_for_mask_type

from .albu_transform import get_albu_transforms

import argparse, time
from torch.utils.data import DataLoader
from skimage.metrics import mean_squared_error, peak_signal_noise_ratio, structural_similarity

def normalize(data, mean, stddev, eps=0.0):
    """
    Normalize the given tensor.

    Applies the formula (data - mean) / (stddev + eps).

    Args:
        data (torch.Tensor): Input data to be normalized.
        mean (float): Mean value.
        stddev (float): Standard deviation.
        eps (float, default=0.0): Added to stddev to prevent dividing by zero.

    Returns:
        torch.Tensor: Normalized tensor
    """
    return (data - mean) / (stddev + eps)

def normalize_instance_dim(data, eps=0.0):
    """
    Normalize the given tensor  with instance norm/

    Applies the formula (data - mean) / (stddev + eps), where mean and stddev
    are computed from the data itself.

    Args:
        data (torch.Tensor): Input data to be normalized
        eps (float): Added to stddev to prevent dividing by zero

    Returns:
        torch.Tensor: Normalized tensor
    """
    mean = data.mean(dim=(1, 2, 3), keepdim=True)   # B, C, H, W
    std = data.std(dim=(1, 2, 3), keepdim=True)

    return normalize(data, mean, std, eps), mean, std


def normalize_instance(data, eps=0.0):
    """
    Normalize the given tensor  with instance norm/

    Applies the formula (data - mean) / (stddev + eps), where mean and stddev
    are computed from the data itself.

    Args:
        data (torch.Tensor): Input data to be normalized
        eps (float): Added to stddev to prevent dividing by zero

    Returns:
        torch.Tensor: Normalized tensor
    """
    mean = data.mean()
    std = data.std()

    return normalize(data, mean, std, eps), mean, std

def undersample_mri(kspace, _MRIDOWN):                                                                                        
    # print("kspace shape:", kspace.shape) ## [18, 4, 256, 256, 2]
    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
    elif _MRIDOWN == "12X":
        mask_type_str, center_fraction, MRIDOWN = "equispaced", 0.03, 12

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

    shape = [256, 256, 1]
    mask = ff(shape, seed=1337) ## [1, 256, 1]

    mask = mask[:, :, 0] # [1, 256]

    masked_kspace = kspace * mask[None, None, :, :, None]

    return masked_kspace, mask.unsqueeze(-1)


def apply_mask(data, mask_func, seed=None, padding=None):
    """
    Subsample given k-space by multiplying with a mask.

    Args:
        data (torch.Tensor): The input k-space data. This should have at least 3 dimensions, where
            dimensions -3 and -2 are the spatial dimensions, and the final dimension has size
            2 (for complex values).
        mask_func (callable): A function that takes a shape (tuple of ints) and a random
            number seed and returns a mask.
        seed (int or 1-d array_like, optional): Seed for the random number generator.

    Returns:
        (tuple): tuple containing:
            masked data (torch.Tensor): Subsampled k-space data
            mask (torch.Tensor): The generated mask
    """
    shape = np.array(data.shape)
    shape[:-3] = 1
    mask = mask_func(shape, seed)
    if padding is not None:
        mask[:, :, : padding[0]] = 0
        mask[:, :, padding[1] :] = 0  # padding value inclusive on right of zeros

    masked_data = data * mask + 0.0  # the + 0.0 removes the sign of the zeros

    return masked_data, mask

def complex_center_crop(data, shape):
    """
    Apply a center crop to the input image or batch of complex images.

    Args:
        data (torch.Tensor): The complex input tensor to be center cropped. It
            should have at least 3 dimensions and the cropping is applied along
            dimensions -3 and -2 and the last dimensions should have a size of
            2.
        shape (int): The output shape. The shape should be smaller than
            the corresponding dimensions of data.

    Returns:
        torch.Tensor: The center cropped image
    """
    assert 0 < shape[0] <= data.shape[-3]
    assert 0 < shape[1] <= data.shape[-2]

    w_from = (data.shape[-3] - shape[0]) // 2   #80
    h_from = (data.shape[-2] - shape[1]) // 2   #80
    w_to = w_from + shape[0]  #240
    h_to = h_from + shape[1]  #240

    return data[..., w_from:w_to, h_from:h_to, :]



def to_tensor(data):
    """
    Convert numpy array to PyTorch tensor.
    
    For complex arrays, the real and imaginary parts are stacked along the last
    dimension.

    Args:
        data (np.array): Input numpy array.

    Returns:
        torch.Tensor: PyTorch version of data.
    """
    if np.iscomplexobj(data):
        data = np.stack((data.real, data.imag), axis=-1)

    return torch.from_numpy(data)

def rss(data, dim=0):
    """
    Compute the Root Sum of Squares (RSS).

    RSS is computed assuming that dim is the coil dimension.

    Args:
        data (torch.Tensor): The input tensor
        dim (int): The dimensions along which to apply the RSS transform

    Returns:
        torch.Tensor: The RSS value.
    """
    return torch.sqrt((data ** 2).sum(dim))

def read_h5(file_name, _MRIDOWN, use_kspace):
    crop_size=[240,240]

    hf = h5py.File(file_name)
    volume_kspace = hf['kspace'][()]
    
    slice_kspace = volume_kspace
    slice_kspace = to_tensor(slice_kspace)
    import imageio as io

    target = ifft2c(slice_kspace)
    target = complex_center_crop(target, crop_size)
    target = complex_abs(target)
    target = rss(target, dim=1)

    if not use_kspace:
        # print("use_kspace is False")

        # masked_kspace, mask = apply_mask(slice_kspace, mask_func, seed=123456)
        masked_kspace, mask = undersample_mri(slice_kspace, _MRIDOWN)
        lq_image = ifft2c(masked_kspace)
        lq_image = complex_center_crop(lq_image, crop_size)
        lq_image = complex_abs(lq_image)
        lq_image = rss(lq_image, dim=1)

    else:
        lq_image = target

    lq_image_list=[]
    mean_list=[]
    std_list=[]
    for i in range(lq_image.shape[0]):
        image, mean, std = normalize_instance(lq_image[i], eps=1e-11)
        image = image.clamp(-6, 6)
        lq_image_list.append(image)
        mean_list.append(mean)
        std_list.append(std)



    target_list=[]

    for i in range(lq_image.shape[0]):
        target_slice = normalize(target[i], mean_list[i], std_list[i], eps=1e-11)
        target_slice = target_slice.clamp(-6, 6)
        target_list.append(target_slice)
    
    return torch.stack(lq_image_list), torch.stack(target_list), torch.stack(mean_list), torch.stack(std_list)





class M4Raw_TrainSet(Dataset):
    def __init__(self, args, use_kspace=False, DEBUG=False):
        mask_func = create_mask_for_mask_type(
            args.MASKTYPE, args.CENTER_FRACTIONS, args.ACCELERATIONS,
        )
        start_time = time.time()

        self.albu_transforms = get_albu_transforms("train", (240, 240))

        self._MRIDOWN = args.MRIDOWN
        self.input_normalize = args.input_normalize
        input_list1 = sorted(glob(os.path.join(args.root_path, 'multicoil_train', '*_T102.h5')))
        input_list2 = [path.replace('_T102.h5','_T101.h5') for path in input_list1]
        input_list3 = [path.replace('_T102.h5','_T103.h5') for path in input_list1]
        if DEBUG:
            input_list1 = input_list1[:2]
            input_list2 = input_list2[:2]
            input_list3 = input_list3[:2]

        T1_input_list = [input_list1, input_list2, input_list3]

        input_list1 = sorted(glob(os.path.join(args.root_path, 'multicoil_train', '*_T202.h5')))
        input_list2 = [path.replace('_T202.h5','_T201.h5') for path in input_list1]
        input_list3 = [path.replace('_T202.h5','_T203.h5') for path in input_list1]
        if DEBUG:
            input_list1 = input_list1[:2]
            input_list2 = input_list2[:2]
            input_list3 = input_list3[:2]


        T2_input_list = [input_list1, input_list2, input_list3]
        
        
        self.T1_input_list = T1_input_list
        self.T2_input_list = T2_input_list
        self.T1_images = np.zeros([len(input_list1),len(T1_input_list), 18, 240, 240])
        self.T2_images = np.zeros([len(input_list2),len(T2_input_list), 18, 240, 240])
        self.T2_masked_images = np.zeros([len(input_list2),len(T2_input_list), 18, 240, 240])
        self.T2_mean = np.zeros([len(input_list2),len(T2_input_list), 18])
        self.T2_std = np.zeros([len(input_list2),len(T2_input_list), 18])

        print('TrainSet loading...')
        for i in range(len(self.T1_input_list)):
            for j, path in enumerate(T1_input_list[i]):
                _, self.T1_images[j][i], _, _ = read_h5(path, self._MRIDOWN, use_kspace=use_kspace)

        self.T1_labels = np.mean(self.T1_images, axis=1)

        for i in range(len(self.T2_input_list)):
            for j, path in enumerate(T2_input_list[i]):
                self.T2_masked_images[j][i], self.T2_images[j][i], self.T2_mean[j][i], self.T2_std[j][i]  = read_h5(path, self._MRIDOWN, use_kspace=use_kspace)
                # lq_image_list, target_list
                 

        # self.T2_labels = np.mean(self.T2_images, axis=1)  # TODO
        self.T2_labels = np.mean(self.T2_images, axis=1)
        self.T2_images = self.T2_masked_images

        print(f'Finish loading with time = {time.time() - start_time}s')

        # print("T1 image original shape:", self.T1_images.shape)  # T1 image original shape: (128, 3, 18, 256, 256)
        # print("T2 image original shape:", self.T2_images.shape)

        N, _, S, H, W = self.T1_images.shape
        self.fname_slices = []

        for i in range(N):
            for j in range(S):
                self.fname_slices.append((i, j))

        
        self.T1_images = self.T1_images.transpose(0,2,1,3,4).reshape(-1,len(T1_input_list),240,240)
        self.T2_images = self.T2_images.transpose(0,2,1,3,4).reshape(-1,len(T2_input_list),240,240)
        self.T1_labels = self.T1_labels.reshape(-1,1,240,240)
        self.T2_labels = self.T2_labels.reshape(-1,1,240,240)
        self.T2_mean = self.T2_mean.reshape(-1,3)
        self.T2_std = self.T2_std.reshape(-1,3)

        print("Train data shape:", self.T1_images.shape)
        # breakpoint()


    def __len__(self):
        return len(self.T1_images)

    def __getitem__(self, idx):
        T1_images = self.T1_images[idx]
        T2_images = self.T2_images[idx]
        T1_labels = self.T1_labels[idx]
        T2_labels = self.T2_labels[idx]

        fname = self.fname_slices[idx][0]
        slice = self.fname_slices[idx][1]


        choices = np.random.choice([i for i in range(len(self.T1_input_list))],1) ## 每次都是从三个repetition中选择一个作为input. 
        # choices = np.random.choice([0],1) ## 用第一个repetition作为输入图像进行测试
        T1_images = T1_images[choices]
        T2_images = T2_images[choices]
        t2_mean = self.T2_mean[idx][choices]
        t2_std = self.T2_std[idx][choices]

        # (1, 240, 240)
        sample = self.albu_transforms(image=T1_images[0], image2=T2_images[0],
                                      image3=T1_labels[0], image4=T2_labels[0])

        # breakpoint()
        t1_in = np.expand_dims(sample['image'], 0)
        t2_in = np.expand_dims(sample['image2'], 0)
        t1 = np.expand_dims(sample['image3'], 0)
        t2 = np.expand_dims(sample['image4'], 0)

        # sample_stats = {"t2_mean": t2_mean, "t2_std": t2_std}

        # print("t1_in shape:", t1_in.shape, "t1 shape:", t1.shape, "t2_in shape:", t2_in.shape, "t2 shape:", t2.shape)

        # breakpoint()
        sample = {
                  'fname': fname,
                  'slice': slice,

                  't1_in': t1_in.astype(np.float32),
                  't1': t1.astype(np.float32),
                   "t2_mean": t2_mean, "t2_std": t2_std,

                  't2_in': t2_in.astype(np.float32),
                  't2': t2.astype(np.float32)}

        return sample #, sample_stats



class M4Raw_TestSet(Dataset):
    def __init__(self, args, use_kspace=False, DEBUG=False):
        mask_func = create_mask_for_mask_type(
            args.MASKTYPE, args.CENTER_FRACTIONS, args.ACCELERATIONS,
        )
        self.use_kspace = use_kspace
        self._MRIDOWN = args.MRIDOWN

        self.input_normalize = args.input_normalize
        input_list1 = sorted(glob(os.path.join(args.root_path, 'multicoil_val' + '*_T102.h5')))
        input_list2 = [path.replace('_T102.h5','_T101.h5') for path in input_list1]
        input_list3 = [path.replace('_T102.h5','_T103.h5') for path in input_list1]
        if DEBUG:
            input_list1 = input_list1[:2]
            input_list2 = input_list2[:2]
            input_list3 = input_list3[:2]


        T1_input_list = [input_list1, input_list2, input_list3]

        input_list1 = sorted(glob(os.path.join(args.root_path, 'multicoil_val', '*_T202.h5')))
        input_list2 = [path.replace('_T202.h5','_T201.h5') for path in input_list1]
        input_list3 = [path.replace('_T202.h5','_T203.h5') for path in input_list1]
        if DEBUG:
            input_list1 = input_list1[:2]
            input_list2 = input_list2[:2]
            input_list3 = input_list3[:2]


        T2_input_list = [input_list1,input_list2,input_list3]
        
        self.T1_input_list = T1_input_list
        self.T2_input_list = T2_input_list
        self.T1_images = np.zeros([len(input_list1),len(T1_input_list), 18, 240, 240])
        self.T2_images = np.zeros([len(input_list2),len(T2_input_list), 18, 240, 240])
        self.T2_masked_images = np.zeros([len(input_list2),len(T2_input_list), 18, 240, 240])
        self.T2_mean = np.zeros([len(input_list2),len(T2_input_list), 18])
        self.T2_std = np.zeros([len(input_list2),len(T2_input_list), 18])
        print('TestSet loading...')

        for i in range(len(self.T1_input_list)):
            for j, path in enumerate(T1_input_list[i]):
                _, self.T1_images[j][i], _, _ = read_h5(path,  self._MRIDOWN,  use_kspace=use_kspace)

        self.T1_labels = np.mean(self.T1_images, axis=1)

        for i in range(len(self.T2_input_list)):
            for j, path in enumerate(T2_input_list[i]):
                self.T2_masked_images[j][i], self.T2_images[j][i], self.T2_mean[j][i], self.T2_std[j][i] = read_h5(path,
                                                                                                                   self._MRIDOWN,
                                                                                                                   use_kspace=use_kspace)
                # lq_image_list, target_list

        # self.T2_labels = np.mean(self.T2_images, axis=1)  # TODO
        self.T2_labels = np.mean(self.T2_images, axis=1)
        self.T2_images = self.T2_masked_images

        print('Finish loading')
        N, _, S, H, W = self.T1_images.shape
        self.fname_slices = []

        for i in range(N):
            for j in range(S):
                self.fname_slices.append((i, j))


        self.T1_images = self.T1_images.transpose(0,2,1,3,4).reshape(-1,len(T1_input_list),240,240)
        self.T2_images = self.T2_images.transpose(0,2,1,3,4).reshape(-1,len(T2_input_list),240,240)
        self.T1_labels = self.T1_labels.reshape(-1,1,240,240)
        self.T2_labels = self.T2_labels.reshape(-1,1,240,240)
        self.T2_mean = self.T2_mean.reshape(-1,3)
        self.T2_std = self.T2_std.reshape(-1,3)
        print("Train data shape:", self.T1_images.shape)


    def __len__(self):
        return len(self.T1_images)

    def __getitem__(self, idx):

        T1_images = self.T1_images[idx]
        T2_images = self.T2_images[idx]
        T1_labels = self.T1_labels[idx]
        T2_labels = self.T2_labels[idx]


        fname = self.fname_slices[idx][0]
        slice = self.fname_slices[idx][1]



        # choices = np.random.choice([i for i in range(len(self.T1_input_list))],1) ## 每次都是从三个repetition中选择一个作为input. 
        choices = np.random.choice([0],1) ## 用第一个repetition作为输入图像进行测试


        T1_images = T1_images[choices]
        T2_images = T2_images[choices]

        t2_mean = self.T2_mean[idx][choices]
        t2_std = self.T2_std[idx][choices]


        # breakpoint()
        # import imageio as io
        # io.imsave('T1_images.png', (T1_images[0]*255).astype(np.uint8))
        # io.imsave('T1_labels.png', (T1_labels[0]*255).astype(np.uint8))
        # io.imsave('T2_images.png', (T2_images[0]*255).astype(np.uint8))
        # io.imsave('T2_labels.png', (T2_labels[0]*255).astype(np.uint8))
        # breakpoint()

        t1_in = T1_images
        t1 = T1_labels
        t2_in = T2_images
        t2 = T2_labels

        # sample_stats = {"t2_mean": t2_mean, "t2_std": t2_std}

        # print("Test t1_in shape:", t1_in.shape, "t1 shape:", t1.shape, "t2_in shape:", t2_in.shape, "t2 shape:", t2.shape)

        # breakpoint()
        sample = {
                  'fname': fname,
                  'slice': slice,

                  't1_in': t1_in.astype(np.float32),
                  't1': t1.astype(np.float32),
                   "t2_mean": t2_mean, "t2_std": t2_std,

                  't2_in': t2_in.astype(np.float32),
                  't2': t2.astype(np.float32)}

        return sample #, sample_stats



def compute_metrics(image, labels):
    MSE = mean_squared_error(labels, image)/np.var(labels)
    PSNR = peak_signal_noise_ratio(labels, image)
    SSIM = structural_similarity(labels, image)

    # print("metrics:", MSE, PSNR, SSIM)

    return MSE, PSNR, SSIM

def complex_abs_eval(data):
    return (data[0:1, :, :] ** 2 + data[1:2, :, :] ** 2).sqrt()


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--root_path', type=str, default='/data/qic99/MRI_recon/')
    parser.add_argument('--MRIDOWN', type=str, default='4X', help='MRI down-sampling rate')
    parser.add_argument('--kspace_refine', type=str, default='False', help='whether use the image reconstructed from kspace network.')
    args = parser.parse_args()

    db_test = M4Raw_TestSet(args)
    testloader = DataLoader(db_test, batch_size=4, shuffle=False, num_workers=4, pin_memory=True)

    t1_MSE_all, t1_PSNR_all, t1_SSIM_all = [], [], []
    t2_MSE_all, t2_PSNR_all, t2_SSIM_all = [], [], []

    save_dir = "./visualize_images/"

    for i_batch, sampled_batch in enumerate(testloader):
        # t1_in, t2_in = sampled_batch['t1_in'].cuda(), sampled_batch['t2_in'].cuda()
        # t1, t2 = sampled_batch['t1_labels'].cuda(), sampled_batch['t2_labels'].cuda()

        t1_in, t2_in = sampled_batch['ref_image_sub'].cuda(), sampled_batch['tag_image_sub'].cuda()
        t1, t2 = sampled_batch['ref_image_full'].cuda(), sampled_batch['tag_image_full'].cuda()

        # breakpoint()
        for j in range(t1_in.shape[0]):
            # t1_in_img = (np.clip(complex_abs_eval(t1_in[j])[0].cpu().numpy(), 0, 1) * 255).astype(np.uint8)
            # t1_img = (np.clip(complex_abs_eval(t1[j])[0].cpu().numpy(), 0, 1) * 255).astype(np.uint8)
            # t2_in_img = (np.clip(complex_abs_eval(t2_in[j])[0].cpu().numpy(), 0, 1) * 255).astype(np.uint8)
            # t2_img = (np.clip(complex_abs_eval(t2[j])[0].cpu().numpy(), 0, 1) * 255).astype(np.uint8)
            # breakpoint()
            t1_in_img = (np.clip(t1_in[j][0].cpu().numpy(), 0, 1) * 255).astype(np.uint8)
            t1_img = (np.clip(t1[j][0].cpu().numpy(), 0, 1) * 255).astype(np.uint8)
            t2_in_img = (np.clip(t2_in[j][0].cpu().numpy(), 0, 1) * 255).astype(np.uint8)
            t2_img = (np.clip(t2[j][0].cpu().numpy(), 0, 1) * 255).astype(np.uint8)

            # t1_in_img = (np.clip(t1_in[j][0].cpu().numpy(), 0, 1) * 255).astype(np.uint8)
            # t1_img = (np.clip(t1[j][0].cpu().numpy(), 0, 1) * 255).astype(np.uint8)
            # t2_in_img = (np.clip(t2_in[j][0].cpu().numpy(), 0, 1) * 255).astype(np.uint8)
            # t2_img = (np.clip(t2[j][0].cpu().numpy(), 0, 1) * 255).astype(np.uint8)


            # print(t1_in_img.shape, t1_img.shape)
            t1_MSE, t1_PSNR, t1_SSIM = compute_metrics(t1_in_img, t1_img)
            t2_MSE, t2_PSNR, t2_SSIM = compute_metrics(t2_in_img, t2_img)


            t1_MSE_all.append(t1_MSE)
            t1_PSNR_all.append(t1_PSNR)
            t1_SSIM_all.append(t1_SSIM)

            t2_MSE_all.append(t2_MSE)
            t2_PSNR_all.append(t2_PSNR)
            t2_SSIM_all.append(t2_SSIM)


    # print("t1_PSNR:", t1_PSNR_all)
    print("t1_PSNR:", round(np.array(t1_PSNR_all).mean(), 4))
    print("t1_NMSE:", round(np.array(t1_MSE_all).mean(), 4))
    print("t1_SSIM:", round(np.array(t1_SSIM_all).mean(), 4))

    print("t2_PSNR:", round(np.array(t2_PSNR_all).mean(), 4))
    print("t2_NMSE:", round(np.array(t2_MSE_all).mean(), 4))
    print("t2_SSIM:", round(np.array(t2_SSIM_all).mean(), 4))