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