| import os |
| import sys |
| from tqdm import tqdm |
| import shutil |
| import argparse |
| import logging |
| import numpy as np |
| from skimage import io |
| from scipy.ndimage import zoom |
|
|
| import torch |
| import torch.nn as nn |
| from torchvision import transforms |
| from torch.utils.data import DataLoader |
| from networks.compare_models import build_model_from_name |
| from dataloaders.BRATS_dataloader_new import Hybrid as MyDataset |
| from dataloaders.BRATS_dataloader_new import ToTensor |
| from networks.mynet import TwoBranch |
| from utils import bright, trunc |
| from skimage.metrics import mean_squared_error, peak_signal_noise_ratio, structural_similarity |
|
|
|
|
| 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]') |
|
|
| |
| from option import args |
| test_data_path = args.root_path |
| snapshot_path = "model/" + args.exp + "/" |
|
|
| os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu |
|
|
|
|
| def normalize_output(out_img): |
| out_img = (out_img - out_img.min())/(out_img.max() - out_img.min() + 1e-8) |
| return out_img |
|
|
|
|
| if __name__ == "__main__": |
| |
| 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)) |
|
|
| network = TwoBranch(args).cuda() |
| device = torch.device('cuda') |
| network.to(device) |
|
|
| if len(args.gpu.split(',')) > 1: |
| network = nn.DataParallel(network) |
|
|
| 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 = [], [], [] |
|
|
| 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 |
| |
|
|
| 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_in.png', bright(t2_in_img,0,0.8)) |
| 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)) |
|
|
|
|
| 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) |
| t2_MSE_all.append(MSE) |
| t2_PSNR_all.append(PSNR) |
| t2_SSIM_all.append(SSIM) |
| print("[t2 MRI] MSE:", MSE, "PSNR:", PSNR, "SSIM:", SSIM) |
|
|
| |
| |
|
|
| 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()) |
|
|
|
|
