MRI_recon/code/Frequency-Diffusion/FSMNet/dataloaders/m4raw_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, time
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.math 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
from dataloaders.kspace_subsample import undersample_mri, mri_fft, mri_fft_m4raw
from tqdm import tqdm

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 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='None', use_kspace=False):
    hf = h5py.File(file_name)
    volume_kspace = hf['kspace'][()]
    slice_kspace = volume_kspace
    slice_kspace2 = to_tensor(slice_kspace)

    slice_image = ifft2c(slice_kspace2)
    slice_image_abs = complex_abs(slice_image)
    slice_image_rss = rss(slice_image_abs, dim=1)
    slice_image_rss = np.abs(slice_image_rss.numpy())
    slice_image_rss = normal(slice_image_rss)

    if _MRIDOWN == 'None' or use_kspace:
        masked_image_rss = slice_image_rss

    else:
        # print("Undersample MRI")
        # Undersample MRI
        masked_kspace, mask = undersample_mri(slice_kspace2, _MRIDOWN)  # Masked

        masked_image     = ifft2c(masked_kspace)
        masked_image_abs = complex_abs(masked_image)
        masked_image_rss = rss(masked_image_abs, dim=1)
        masked_image_rss = np.abs(masked_image_rss.numpy())
        masked_image_rss = normal(masked_image_rss)
        
    return slice_image_rss, masked_image_rss


DEBUG = True

class M4Raw_TrainSet(Dataset):
    def __init__(self, root_path, MRIDOWN, kspace_refine='False', use_kspace=False):

        self.use_kspace = use_kspace
        self.kspace_refine = kspace_refine
        start_time = time.time()

        input_list1 = sorted(glob(os.path.join(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(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, 256, 256])
        self.T2_images = np.zeros([len(input_list2),len(T2_input_list), 18, 256, 256])
        self.T2_masked_images = np.zeros([len(input_list2),len(T2_input_list), 18, 256, 256])

        """
        读取kspace network重建的图像
        """
        if kspace_refine == 'True':
            krecon_list1 = sorted(glob(os.path.join(root_path + 'multicoil_train' + '/*_T102_recon_kspace_round2_images.npy')))
            krecon_list2 = [path.replace('_T102','_T101') for path in krecon_list1]
            krecon_list3 = [path.replace('_T102','_T103') for path in krecon_list1]
            T1_krecon_list = [krecon_list1, krecon_list2, krecon_list3]

            krecon_list1 = sorted(glob(os.path.join(root_path + 'multicoil_train' + '/*_T202_recon_kspace_round2_images.npy')))
            krecon_list2 = [path.replace('_T202','_T201') for path in krecon_list1]
            krecon_list3 = [path.replace('_T202','_T203') for path in krecon_list1]
            T2_krecon_list = [krecon_list1, krecon_list2, krecon_list3]

            self.T1_krecon_list = T1_krecon_list
            self.T2_krecon_list = T2_krecon_list
            
            self.T1_krecon = np.zeros([len(input_list1), len(T1_krecon_list), 18, 240, 240]).astype(np.float32)
            self.T2_krecon = np.zeros([len(input_list2), len(T2_krecon_list), 18, 240, 240]).astype(np.float32)



        print('TrainSet loading...')
        for i in tqdm(range(len(self.T1_input_list))):
            for j, path in enumerate(T1_input_list[i]):
                self.T1_images[j][i], _ = read_h5(path, use_kspace=use_kspace)
                # self.fname_slices[i].append(path)       # each coil

            if kspace_refine == 'True':
                for k, path in enumerate(T1_krecon_list[i]):
                    self.T1_krecon[k][i] = np.load(path).astype(np.float32)/255.0
        self.T1_labels = np.mean(self.T1_images, axis=1) # multi-coil mean

        for i in tqdm(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] = read_h5(path, _MRIDOWN=MRIDOWN, use_kspace=use_kspace)
            if kspace_refine == 'True':
                for k, path in enumerate(T2_krecon_list[i]):
                    self.T2_krecon[k][i] = np.load(path).astype(np.float32)/255.0

        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))
                # print(f'nan value at {i}, {j}, {k}, {l}')

        self.T1_images = self.T1_images.transpose(0,2,1,3,4).reshape(-1,len(T1_input_list),256,256)[:, :, 8:248, 8:248]
        self.T2_images = self.T2_images.transpose(0,2,1,3,4).reshape(-1,len(T2_input_list),256,256)[:, :, 8:248, 8:248]
        self.T1_labels = self.T1_labels.reshape(-1,1,256,256)[:, :, 8:248, 8:248]
        self.T2_labels = self.T2_labels.reshape(-1,1,256,256)[:, :, 8:248, 8:248]

        # Train data shape: (2304, 3, 240, 240)

        # T1 N, 3, 240, 240

        if kspace_refine == 'True':
            self.T1_krecon = self.T1_krecon.transpose(0,2,1,3,4).reshape(-1,len(T1_krecon_list),240,240)
            self.T2_krecon = self.T2_krecon.transpose(0,2,1,3,4).reshape(-1,len(T2_krecon_list),240,240)


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

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

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

        ## 每次都是从三个repetition中选择一个作为input.
        choices = np.random.choice([i for i in range(len(self.T1_input_list))],1)
        T1_images = T1_images[choices]
        T2_images = T2_images[choices]

        t1_kspace_in, t1_in, t1_kspace, t1_img = mri_fft_m4raw(T1_images, T1_labels)
        t2_kspace_in, t2_in, t2_kspace, t2_img = mri_fft_m4raw(T2_images, T2_labels)

        
        # normalize
        t1_img, t1_mean, t1_std = normalize_instance(t1_img)
        t1_in = normalize(t1_in, t1_mean, t1_std)
        # t1_mean = 0
        # t1_std = 1

        t2_img, t2_mean, t2_std = normalize_instance(t2_img)
        t2_in = normalize(t2_in, t2_mean, t2_std)



        # filter value that greater or less than 6
        t1_img = torch.clamp(t1_img, -6, 6)
        t2_img = torch.clamp(t2_img, -6, 6)
        t1_in = torch.clamp(t1_in, -6, 6)
        t2_in = torch.clamp(t2_in, -6, 6)


        # t2_mean = 0
        # t2_std = 1

        # t1_img: torch.Size([2, 240, 240]) torch.float32 tensor(0.9775) tensor(-9.0143e-08)
        # t2_img: torch.Size([1, 240, 240]) torch.float32 tensor(22.1929) tensor(-0.3244)

        # t1_img: torch.Size([1, 240, 240]) torch.float32 tensor(5.1340) tensor(1.7756e-06)
        # t2_img: torch.Size([1, 240, 240]) torch.float32 tensor(4.4957) tensor(2.8719e-05)
        # t1_in:  torch.Size([1, 240, 240]) torch.float32 tensor(5.2390) tensor(0.0003)
        # t2_in:  torch.Size([1, 240, 240]) torch.float32 tensor(4.7321) tensor(4.5622e-05)

        # print("t1_img:", t1_img.shape, t1_img.dtype, t1_img.max(), t1_img.min())
        # print("t2_img:", t2_img.shape, t2_img.dtype, t2_img.max(), t2_img.min())
        # print("t1_in:", t1_in.shape, t1_in.dtype, t1_in.max(), t1_in.min())
        # print("t2_in:", t2_in.shape, t2_in.dtype, t2_in.max(), t2_in.min())  # t1_img: torch.Size([1, 240, 240]) torch.float32 tensor(20.5561) tensor(-0.2671)
        # print()

        # How to get mean and std of the training data?
        # fname, slice
        sample =  {
                'fname': fname,
                'slice': slice,

                'ref_kspace_full': t1_kspace,
                'ref_kspace_sub': t1_kspace_in,
                'ref_image_full': t1_img,
                'ref_image_sub': t1_in,
                't1_mean': t1_mean,
                't1_std': t1_std,

                'tag_kspace_full': t2_kspace,
                'tag_kspace_sub': t2_kspace_in,
                'tag_image_full': t2_img,
                'tag_image_sub': t2_in,
                't2_mean': t2_mean,
                't2_std': t2_std,

                }

        return sample



class M4Raw_TestSet(Dataset):
    def __init__(self, root_path, MRIDOWN, kspace_refine='False', use_kspace=False):

        self.kspace_refine = kspace_refine

        input_list1 = sorted(glob(os.path.join(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(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, 256, 256])
        self.T2_images = np.zeros([len(input_list2),len(T2_input_list), 18, 256, 256])
        self.T2_masked_images = np.zeros([len(input_list2),len(T2_input_list), 18, 256, 256])

        """
        读取kspace network重建的图像
        """
        if kspace_refine == 'True':
            krecon_list1 = sorted(glob(os.path.join(root_path + 'multicoil_val' + '/*_T102_recon_kspace_round2_images.npy')))
            krecon_list2 = [path.replace('_T102','_T101') for path in krecon_list1]
            krecon_list3 = [path.replace('_T102','_T103') for path in krecon_list1]
            T1_krecon_list = [krecon_list1, krecon_list2, krecon_list3]

            krecon_list1 = sorted(glob(os.path.join(root_path + 'multicoil_val' + '/*_T202_recon_kspace_round2_images.npy')))
            krecon_list2 = [path.replace('_T202','_T201') for path in krecon_list1]
            krecon_list3 = [path.replace('_T202','_T203') for path in krecon_list1]
            T2_krecon_list = [krecon_list1, krecon_list2, krecon_list3]

            self.T1_krecon_list = T1_krecon_list
            self.T2_krecon_list = T2_krecon_list
            
            self.T1_krecon = np.zeros([len(input_list1), len(T1_krecon_list), 18, 240, 240]).astype(np.float32)
            self.T2_krecon = np.zeros([len(input_list2), len(T2_krecon_list), 18, 240, 240]).astype(np.float32)


        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, use_kspace=use_kspace)

            if kspace_refine == 'True':
                for k, path in enumerate(T1_krecon_list[i]):
                    self.T1_krecon[k][i] = np.load(path).astype(np.float32)/255.0
        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] = read_h5(path, _MRIDOWN = MRIDOWN, use_kspace=use_kspace)

            if kspace_refine == 'True':
                for k, path in enumerate(T2_krecon_list[i]):
                    self.T2_krecon[k][i] = np.load(path).astype(np.float32)/255.0
        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),256,256)[:, :, 8:248, 8:248]
        self.T2_images = self.T2_images.transpose(0,2,1,3,4).reshape(-1,len(T2_input_list),256,256)[:, :, 8:248, 8:248]
        self.T1_labels = self.T1_labels.reshape(-1,1,256,256)[:, :, 8:248, 8:248]
        self.T2_labels = self.T2_labels.reshape(-1,1,256,256)[:, :, 8:248, 8:248]
        print("Test data shape:", self.T1_images.shape)

        if kspace_refine == 'True':
            self.T1_krecon = self.T1_krecon.transpose(0,2,1,3,4).reshape(-1,len(T1_krecon_list),240,240)
            self.T2_krecon = self.T2_krecon.transpose(0,2,1,3,4).reshape(-1,len(T2_krecon_list),240,240)


    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]

        # print("T1_labels:", T1_labels.shape, T1_labels.dtype, T1_labels.max(), T1_labels.min())
        # print("T2_labels:", T2_labels.shape, T2_labels.dtype, T2_labels.max(), T2_labels.min())

        choices = np.random.choice([0],1) ## 用第一个repetition作为输入图像进行测试
        T1_images = T1_images[choices]
        T2_images = T2_images[choices]

        t1_kspace_in, t1_in, t1_kspace, t1_img = mri_fft_m4raw(T1_images, T1_labels)
        t2_kspace_in, t2_in, t2_kspace, t2_img = mri_fft_m4raw(T2_images, T2_labels)

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

        # normalize
        t1_img, t1_mean, t1_std = normalize_instance(t1_img)
        t1_in = normalize(t1_in, t1_mean, t1_std)
        # t1_mean = 0
        # t1_std = 1

        t2_img, t2_mean, t2_std = normalize_instance(t2_img)
        t2_in = normalize(t2_in, t2_mean, t2_std)

        # filter value that greater or less than 6
        t1_img = torch.clamp(t1_img, -6, 6)
        t2_img = torch.clamp(t2_img, -6, 6)
        t1_in = torch.clamp(t1_in, -6, 6)
        t2_in = torch.clamp(t2_in, -6, 6)


        # print("t1_img:", t1_img.shape, t1_img.dtype, t1_img.max(), t1_img.min())
        # print("in dataset t2_img:", t2_img.shape, t2_img.dtype, t2_img.max(), t2_img.min())
        # print("t1_in:", t1_in.shape, t1_in.dtype, t1_in.max(), t1_in.min())
        # print("t2_in:", t2_in.shape, t2_in.dtype, t2_in.max(), t2_in.min())  # t1_img: torch.Size([1, 240, 240]) torch.float32 tensor(20.5561) tensor(-0.2671)
        # print()

        # fname, slice
        sample = {
            'fname': fname,
            'slice': slice,

            'ref_kspace_full': t1_kspace,
            'ref_kspace_sub': t1_kspace_in,
            'ref_image_full': t1_img,
            'ref_image_sub': t1_in,
            't1_mean': t1_mean,
            't1_std': t1_std,

            'tag_kspace_full': t2_kspace,
            'tag_kspace_sub': t2_kspace_in,
            'tag_image_full': t2_img,
            'tag_image_sub': t2_in,
            't2_mean': t2_mean,
            't2_std': t2_std,

        }

        return sample


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