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