MRI_recon/code/Frequency-Diffusion/dataset/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

import argparse
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(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 normal(x):
#     y = np.zeros_like(x)
#     for i in range(y.shape[0]):
#         x_min = x[i].min()
#         x_max = x[i].max()
#         y[i] = (x[i] - x_min)/(x_max-x_min)
#     return y



# 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

#     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 = [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, mask_func):
    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



    masked_kspace, mask = apply_mask(slice_kspace, mask_func, seed=123456)
    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)
    # breakpoint()
    # io.imsave('lq_image.png', lq_image[0].numpy().astype(np.uint8))
    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 = ifft2c(slice_kspace)
    target = complex_center_crop(target, crop_size)
    target = complex_abs(target)
    target = rss(target, dim=1)
    # io.imsave('target1.png', target[10].numpy().astype(np.uint8))
    # breakpoint()
    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):
        mask_func = create_mask_for_mask_type(
            args.MASKTYPE, args.CENTER_FRACTIONS, args.ACCELERATIONS,
        )

        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]
        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]
        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, mask_func)

        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_images[j][i], self.T2_masked_images[j][i], self.T2_mean[j][i], self.T2_std[j][i]  = read_h5(path, mask_func)
                 
                 

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

        print('Finish loading')
        
        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]
        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}


        # breakpoint()
        sample = {'image_in': t1_in, 
                  'image': t1, 

                  'target_in': t2_in, 
                  'target': t2}

        return sample, sample_stats



class M4Raw_TestSet(Dataset):
    def __init__(self, args):
        mask_func = create_mask_for_mask_type(
            args.MASKTYPE, args.CENTER_FRACTIONS, args.ACCELERATIONS,
        )

        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]
        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]
        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, mask_func)

        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_images[j][i], self.T2_masked_images[j][i], self.T2_mean[j][i], self.T2_std[j][i]  = read_h5(path, mask_func)
                 
                 

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

        print('Finish loading')
        
        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]
        # 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}


        # breakpoint()
        sample = {'image_in': t1_in, 
                  'image': t1, 

                  'target_in': t2_in, 
                  'target': t2}

        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))