MRI_recon/code/Frequency-Diffusion/FSMNet/test_brats.py

import os
import sys
from tqdm import tqdm
import argparse
import logging
from skimage import io

from torchvision import transforms
from torch.utils.data import DataLoader
from dataloaders.BRATS_dataloader_new import Hybrid as MyDataset
from dataloaders.BRATS_dataloader_new import ToTensor
from networks.mynet import TwoBranch
from skimage.metrics import mean_squared_error, peak_signal_noise_ratio, structural_similarity
from utils.option import args


def normalise_mse(gt, pred):
    """Compute Normalized Mean Squared Error (NMSE)"""
    return np.linalg.norm(gt - pred) ** 2 / np.linalg.norm(gt) ** 2



parser = argparse.ArgumentParser()
parser.add_argument('--root_path', type=str, default='/home/xiaohan/datasets/BRATS_dataset/BRATS_2020_images/selected_images/')
parser.add_argument('--MRIDOWN', type=str, default='4X', help='MRI down-sampling rate')
parser.add_argument('--low_field_SNR', type=int, default=15, help='SNR of the simulated low-field image')
parser.add_argument('--phase', type=str, default='test', help='Name of phase')
parser.add_argument('--gpu', type=str, default='0', help='GPU to use')
parser.add_argument('--exp', type=str, default='msl_model', help='model_name')
parser.add_argument('--seed', type=int, default=1337, help='random seed')
parser.add_argument('--base_lr', type=float, default=0.0002, help='maximum epoch numaber to train')

parser.add_argument('--model_name', type=str, default='unet_single', help='model_name')
parser.add_argument('--relation_consistency', type=str, default='False', help='regularize the consistency of feature relation')
parser.add_argument('--norm', type=str, default='False', help='Norm Layer between UNet and Transformer')
parser.add_argument('--input_normalize', type=str, default='mean_std', help='choose from [min_max, mean_std, divide]')
parser.add_argument('--test_sample', default="Ksample", help="Ksample | ColdDiffusion | DDPM")

# args = parser.parse_args()

test_data_path = args.root_path
snapshot_path = "model/" + args.exp + "/"

os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu


from utils.utils import *
from frequency_diffusion.degradation.k_degradation import get_ksu_kernel, apply_tofre, apply_to_spatial
from networks_time.mynet import DiffTwoBranch

DEBUG = args.DEBUG
use_time_model = args.use_time_model
use_kspace = args.use_kspace
use_t2_in = True

num_timesteps = 5
image_size    = 240

if args.MRIDOWN == "4X":
    accelerate_mask = np.load("./dataloaders/example_mask/brats_4X_mask.npy")
    accelerate_mask = torch.from_numpy(accelerate_mask).unsqueeze(0).clone().float()
    print("accelerate_mask shape =", accelerate_mask.shape)
else:
    accelerate_mask = None

k_file = f"./dataloaders/example_mask/brats_{args.ACCELERATIONS[0]}_kspace_mask.npy"
if os.path.exists(k_file):
    kspace_masks = np.load(k_file)
    kspace_masks = torch.from_numpy(np.asarray(kspace_masks)).cuda()

else:
    # Output a list of k-space kernels
    kspace_masks = get_ksu_kernel(num_timesteps, image_size,
                                  ksu_routine="LogSamplingRate",
                                  accelerated_factor=args.ACCELERATIONS[0],
                                  accelerate_mask=accelerate_mask
                                  )
    kspace_masks = torch.from_numpy(np.asarray(kspace_masks)).cuda()


def normalize_output(out_img):
    out_img = (out_img - out_img.min())/(out_img.max() - out_img.min() + 1e-8)
    return out_img


test_sample = args.test_sample  # Ksample | ColdDiffusion | DDPM

if __name__ == "__main__":
    ## make logger file
    if use_kspace:
        snapshot_path = snapshot_path.rstrip("/") + f'_t{num_timesteps}_kspace/'

    if use_time_model:
        snapshot_path = snapshot_path.rstrip("/") + '_time/'

    if not isinstance(args.test_tag, type(None)):
        snapshot_path = snapshot_path.rstrip("/") + f'_{args.test_tag}/'

    if not os.path.exists(snapshot_path):
        os.makedirs(snapshot_path)

    logging.basicConfig(filename=snapshot_path + "/log.txt", level=logging.INFO,
                        format='[%(asctime)s.%(msecs)03d] %(message)s', datefmt='%H:%M:%S')
    logging.getLogger().addHandler(logging.StreamHandler(sys.stdout))
    logging.info(str(args))

    if use_time_model:
        network = DiffTwoBranch(args).cuda()
    else:
        network = TwoBranch(args).cuda()


    device = torch.device('cuda')
    network.to(device)

    if len(args.gpu.split(',')) > 1:
        network = nn.DataParallel(network)

    n_parameters = sum(p.numel() for p in network.parameters() if p.requires_grad)
    print('number of params: %.2f M' % (n_parameters / 1024 / 1024))

    db_test = MyDataset(split='test', MRIDOWN=args.MRIDOWN, SNR=args.low_field_SNR, 
                        transform=transforms.Compose([ToTensor()]),
                        base_dir=test_data_path, input_normalize = args.input_normalize)
    testloader = DataLoader(db_test, batch_size=1, shuffle=False, num_workers=2, pin_memory=True)

    if args.phase == 'test':

        save_mode_path = os.path.join(snapshot_path, 'best_checkpoint.pth')
        print('load weights from ' + save_mode_path)
        checkpoint = torch.load(save_mode_path)
        network.load_state_dict(checkpoint['network'])
        network.eval()
        cnt = 0
        save_path = snapshot_path + '/result_case/'
        feature_save_path = snapshot_path + '/feature_visualization/'
        if not os.path.exists(save_path):
            os.makedirs(save_path)
        if not os.path.exists(feature_save_path):
            os.makedirs(feature_save_path)


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

        for (sampled_batch, sample_stats) in tqdm(testloader, ncols=70):
            cnt += 1

            print('processing ' + str(cnt) + ' image')
            t1_in, t1, t2_in, t2 = sampled_batch['image_in'].cuda(), sampled_batch['image'].cuda(), \
                                   sampled_batch['target_in'].cuda(), sampled_batch['target'].cuda()
            t1_krecon, t2_krecon = sampled_batch['image_krecon'].cuda(), sampled_batch['target_krecon'].cuda()

            t2_mean = sample_stats['t2_mean'].data.cpu().numpy()[0]
            t2_std = sample_stats['t2_std'].data.cpu().numpy()[0]


            t1_out, t2_out = None, None

            if use_kspace:
                b = t2.shape[0]
                t = torch.randint(num_timesteps - 1, num_timesteps, (b,), device=device).long()  # t-1

                mask = kspace_masks[t]
                target_fft, _ = apply_tofre(t2.clone(), mask)
                fft, mask = apply_tofre(t2_in.clone(), mask)

                fft = target_fft * mask + fft * (1 - mask)  # Seems too easy
                t2_in = apply_to_spatial(fft)


                while t >= 0:
                    if use_time_model:
                        outputs = network(t2_in, t1_in, t)['img_out']
                    else:
                        outputs = network(t2_in, t1_in)['img_out']

                    if t == 0:
                        mask = kspace_masks[0]  # last one
                        t2_in = outputs
                        if use_time_model:
                            t2_out_2 = network(t2_in, t1_in, t)['img_out']
                        else:
                            t2_out_2 = network(t2_in, t1_in)['img_out']
                    else:

                        if test_sample == "Ksample":   # Ksample | ColdDiffusion | DDPM

                            k_full = kspace_masks[-1]
                            t2_in_fre, k_full = apply_tofre(t2_in, k_full)

                            with torch.no_grad():

                                kt_sub_1 = kspace_masks[t - 1]  # get_kspace_kernels(t - 2).cuda()
                                kt = kspace_masks[t]  # self.get_kspace_kernels(t - 1).cuda()  # last one
                                k_residual = kt_sub_1 - kt

                                recon_sample_fre, k_residual = apply_tofre(outputs, k_residual)
                                # fre_amend = recon_sample_fre * k_residual

                                t2_in_fre = t2_in_fre * (1 - k_residual) + recon_sample_fre * k_residual

                                outputs = apply_to_spatial(t2_in_fre)
                                t2_in = outputs

                        elif test_sample == "ColdDiffusion":
                            k_full = kspace_masks[-1]
                            # t2_in_fre, k_full = apply_tofre(t2_in, k_full)

                            with torch.no_grad():

                                kt_sub_1 = kspace_masks[t - 1]  # get_kspace_kernels(t - 2).cuda()
                                kt = kspace_masks[t]  # self.get_kspace_kernels(t - 1).cuda()  # last one

                                k_residual = kt_sub_1 - kt

                                recon_sample_fre, k_residual = apply_tofre(outputs, k_residual)

                                x_t_hat_fre = recon_sample_fre * kt
                                x_t_sub_1_hat_fre = recon_sample_fre * kt_sub_1

                                x_t_hat = apply_to_spatial(x_t_hat_fre)
                                x_t_sub_1_hat = apply_to_spatial(x_t_sub_1_hat_fre)

                                outputs = t2_in - x_t_hat + x_t_sub_1_hat

                                t2_in = outputs

                        elif test_sample == "DDPM":

                            with torch.no_grad():

                                kt_sub_1 = kspace_masks[t - 1]  # get_kspace_kernels(t - 2).cuda()

                                recon_sample_fre, kt_sub_1 = apply_tofre(outputs, kt_sub_1)
                                fre_new = recon_sample_fre * kt_sub_1

                                outputs = apply_to_spatial(fre_new)
                                t2_in = outputs

                t = t - 1
                t2_out = outputs

            else:
                t2_out = network(t2_in, t1_in)['img_out']
                t2_out_2 = network(t2_in, t1_in)['img_out']

            t1_mean = sample_stats['t1_mean'].data.cpu().numpy()[0]
            t1_std = sample_stats['t1_std'].data.cpu().numpy()[0]
            t2_mean = sample_stats['t2_mean'].data.cpu().numpy()[0]
            t2_std = sample_stats['t2_std'].data.cpu().numpy()[0]

            if t1_out is not None:
                t1_img = (np.clip(t1.data.cpu().numpy()[0, 0] * t1_std + t1_mean, 0, 1) * 255).astype(np.uint8)
                t1_out_img = (np.clip(t1_out.data.cpu().numpy()[0, 0] * t1_std + t1_mean, 0, 1) * 255).astype(np.uint8)
                t1_krecon_img = (np.clip(t1_krecon.data.cpu().numpy()[0, 0] * t1_std + t1_mean, 0, 1) * 255).astype(np.uint8)

            t1_img = (np.clip(t1.data.cpu().numpy()[0, 0] * t1_std + t1_mean, 0, 1) * 255).astype(np.uint8)
            t2_in_img = (np.clip(t2_in.data.cpu().numpy()[0, 0] * t2_std + t2_mean, 0, 1) * 255).astype(np.uint8)
            t2_img = (np.clip(t2.data.cpu().numpy()[0, 0] * t2_std + t2_mean, 0, 1) * 255).astype(np.uint8)
            t2_out_img = (np.clip(t2_out.data.cpu().numpy()[0, 0] * t2_std + t2_mean, 0, 1) * 255).astype(np.uint8)
            t2_krecon_img = (np.clip(t2_krecon.data.cpu().numpy()[0, 0] * t2_std + t2_mean, 0, 1) * 255).astype(np.uint8)
            t2_out_2_img = (np.clip(t2_out_2.data.cpu().numpy()[0, 0] * t2_std + t2_mean, 0, 1) * 255).astype(np.uint8)


            io.imsave(save_path + str(cnt) + '_t1.png', bright(t1_img,0,0.8))
            io.imsave(save_path + str(cnt) + '_t2.png', bright(t2_img,0,0.8))
            io.imsave(save_path + str(cnt) + '_t2_original.png', t2_img)
            io.imsave(save_path + str(cnt) + '_t2_in.png', bright(t2_in_img,0,0.8))
            io.imsave(save_path + str(cnt) + '_t2_out_original.png', t2_out_img)
            io.imsave(save_path + str(cnt) + '_t2_out.png', bright(t2_out_img,0,0.8))
            io.imsave(save_path + str(cnt) + '_t2_out2.png', bright(t2_out_2_img,0,0.8))

            # ------------------------------------
            # NMSE: 5.4534  ± 1.5515
            # PSNR: 39.2132 ± 1.6888
            # SSIM: 0.9792  ± 0.0054
            # ------------------------------------
            # Save Path: model/FSMNet_BraTS_8x_kspace//result_case/

            if t2_out is not None:
                t2_out_img[t2_out_img < 0.0] = 0.0
                t2_img[t2_img < 0.0] = 0.0
                MSE  = mean_squared_error(t2_img, t2_out_img)
                PSNR = peak_signal_noise_ratio(t2_img, t2_out_img)
                SSIM = structural_similarity(t2_img, t2_out_img)
                nmse = normalise_mse(t2_img/255, t2_out_img/255)

                t2_MSE_all.append(MSE)
                t2_PSNR_all.append(PSNR)
                t2_SSIM_all.append(SSIM)
                t2_NMSE_all.append(nmse)

                print("[t2 MRI] MSE:", MSE, "PSNR:", PSNR, "SSIM:", SSIM, "NMSE:", nmse)


        print("===> Evaluate Metric <===")
        print("Results")
        print("-" * 36)
        print(f"{test_sample} NMSE: {np.array(t2_NMSE_all).mean() * 100:.4f}  ± {np.array(t2_NMSE_all).std() * 100 :.4f}")
        # print(f"MSE: {np.array(t2_MSE_all).mean():.4f} ± {np.array(t2_MSE_all).std():.4f}")
        print(f"{test_sample} PSNR: {np.array(t2_PSNR_all).mean():.4f} ± {np.array(t2_PSNR_all).std():.4f}")
        print(f"{test_sample} SSIM: {np.array(t2_SSIM_all).mean():.4f}  ± {np.array(t2_SSIM_all).std():.4f}")
        print("-" * 36)
        print(f"Save Path: {save_path}")



        # print("[T2 MRI:] average MSE:", np.array(t2_MSE_all).mean(), "average PSNR:", np.array(t2_PSNR_all).mean(), "average SSIM:", np.array(t2_SSIM_all).mean())
        # print("[T2 MRI:] average MSE:", np.array(t2_MSE_all).std(), "average PSNR:", np.array(t2_PSNR_all).std(), "average SSIM:", np.array(t2_SSIM_all).std())