from tqdm import tqdm from tensorboardX import SummaryWriter import logging, time, os, sys import torch.optim as optim 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 RandomPadCrop, ToTensor from networks.mynet import TwoBranch from utils.option import args from skimage.metrics import mean_squared_error, peak_signal_noise_ratio, structural_similarity 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 train_data_path = args.root_path test_data_path = args.root_path snapshot_path = "model/" + args.exp + "/" os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu batch_size = args.batch_size * len(args.gpu.split(',')) max_iterations = args.max_iterations base_lr = args.base_lr # --use_time_model True --use_kspace True --ACCELERATIONS 4 --MRIDOWN 4X --low_field_SNR 20 --input_normalize mean_std DEBUG = args.DEBUG use_time_model = args.use_time_model use_kspace = args.use_kspace # PSNR: 30.138548551934974 average SSIM: 0.770964106980312 # PSNR: 31.325274490046855 average SSIM: 0.8589609042898623 4X if not # PSNR: 29.846184815515585 average SSIM: 0.8797758188214125 -> 31.18, 0.77 # PSNR: 28.494279128317515 average SSIM: 0.8179950512965841 8X if not # kspace_refine = True # Albu with 41.33, w/ 42.00 # mask_vacant = False frequency_distortion = True num_timesteps = args.num_timesteps #30 image_size = args.image_size #240 distortion_sigma = 10/255 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() else: accelerate_mask = None # 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 ) np.save(f"./dataloaders/example_mask/brats_{args.ACCELERATIONS[0]}_kspace_mask.npy", kspace_masks) kspace_masks = torch.from_numpy(np.asarray(kspace_masks)).cuda() 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 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_train = MyDataset(split='train', MRIDOWN=args.MRIDOWN, SNR=args.low_field_SNR, transform=transforms.Compose([RandomPadCrop(), ToTensor()]), base_dir=train_data_path, input_normalize = args.input_normalize, use_kspace=use_kspace) 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) trainloader = DataLoader(db_train, batch_size=batch_size, shuffle=True, num_workers=4, pin_memory=True) fixtrainloader = DataLoader(db_train, batch_size=1, shuffle=False, num_workers=4, pin_memory=True) testloader = DataLoader(db_test, batch_size=1, shuffle=False, num_workers=4, pin_memory=True) if args.phase == 'train': network.train() params = list(network.parameters()) optimizer1 = optim.AdamW(params, lr=base_lr, betas=(0.9, 0.999), weight_decay=1e-4) if not use_kspace: scheduler1 = optim.lr_scheduler.StepLR(optimizer1, step_size=20000, gamma=0.5) else: scheduler1 = optim.lr_scheduler.StepLR(optimizer1, step_size=40000, gamma=0.5) writer = SummaryWriter(snapshot_path + '/log') iter_num = 0 max_epoch = max_iterations // len(trainloader) + 1 if use_kspace: max_epoch = max_epoch * num_timesteps best_status = {'T1_NMSE': 10000000, 'T1_PSNR': 0, 'T1_SSIM': 0, 'T2_NMSE': 10000000, 'T2_PSNR': 0, 'T2_SSIM': 0} fft_weight = 0.01 criterion = nn.L1Loss().to(device, non_blocking=True) freloss = Frequency_Loss().to(device, non_blocking=True) start_time = time.time() mask = None for epoch_num in tqdm(range(max_epoch), ncols=70): time1 = time.time() debug_time = False for i_batch, (sampled_batch, sample_stats) in enumerate(trainloader): time2 = time.time() 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() # Degradation if use_kspace: b = t1_in.shape[0] t = torch.randint(0, num_timesteps, (b,), device=device).long() mask = kspace_masks[t] target_fft, _ = apply_tofre(t2.clone(), mask) fft, mask = apply_tofre(t2_in.clone(), mask) # if np.random.rand() > (1 / (1 + num_timesteps)): fft = target_fft * mask + fft * (1 - mask) # Seems too easy # Frequency Noise if frequency_distortion: fft_magnitude = torch.abs(fft) # 幅度 fft_phase = torch.angle(fft) # 相位 # Add noise to unmasked frequencies to maintain the stochasticity that diffusion models typically rely on. sigma = distortion_sigma * torch.abs(torch.randn(1)).item() noise = torch.randn_like(fft_magnitude) * sigma noise_magnitude = noise * fft_magnitude * mask # + noise * (1 - mask) fft_magnitude += noise_magnitude sigma = distortion_sigma / 2 * torch.abs(torch.randn(1)).item() noise = torch.randn_like(fft_phase) * sigma noise_pha = noise * fft_phase * mask # + noise * (1 - mask) fft_phase += noise_pha fft = fft_magnitude * torch.exp(1j * fft_phase) t2_in = apply_to_spatial(fft) time3 = time.time() if use_time_model and use_kspace: outputs = network(t2_in, t1, t) else: outputs = network(t2_in, t1) loss = criterion(outputs['img_out'], t2) + criterion(outputs['img_fre'], t2) + \ fft_weight * freloss(outputs['img_fre'], t2, mask) time4 = time.time() optimizer1.zero_grad() loss.backward() if args.clip_grad == "True": ### clip the gradients to a small range. torch.nn.utils.clip_grad_norm_(network.parameters(), 0.01) optimizer1.step() scheduler1.step() if debug_time: print("Optimizer Step Time: ", time.time() - time2) time5 = time.time() # summary iter_num = iter_num + 1 if iter_num % 100 == 0: logging.info('iteration %d [%.2f sec]: learning rate : %f loss : %f ' % (iter_num, time.time()-start_time, scheduler1.get_lr()[0], loss.item())) if DEBUG: break if iter_num % 20000 == 0: save_mode_path = os.path.join(snapshot_path, 'iter_' + str(iter_num) + '.pth') torch.save({'network': network.state_dict()}, save_mode_path) logging.info("save model to {}".format(save_mode_path)) if iter_num > max_iterations: break time1 = time.time() ## ================ Evaluate ================ logging.info(f'Epoch {epoch_num} Evaluation:') # print() t1_MSE_all, t1_PSNR_all, t1_SSIM_all = [], [], [] t2_MSE_all, t2_PSNR_all, t2_SSIM_all = [], [], [] t2_MSE_first_step, t2_PSNR_first_step, t2_SSIM_first_step = [], [], [] t1_MSE_krecon, t1_PSNR_krecon, t1_SSIM_krecon = [], [], [] t2_MSE_krecon, t2_PSNR_krecon, t2_SSIM_krecon = [], [], [] ids = 0 for (sampled_batch, sample_stats) in testloader: 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() t_merge = torch.cat([t1_in, t2_in], dim=1) if use_kspace: t = num_timesteps - 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, t)['img_out'] else: outputs = network(t2_in, t1)['img_out'] if t == num_timesteps - 1: first_step_recon = outputs if t == 0: mask = kspace_masks[0] # last one t2_in = outputs else: k_full = kspace_masks[-1] # True 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] # current one k_residual = kt_sub_1 - kt recon_sample_fre, k_residual = apply_tofre(outputs, k_residual) # fft = target_fft * mask + fft * (1 - mask) t2_in_fre = t2_in_fre * (1 - k_residual) + recon_sample_fre * k_residual # substitute outputs = apply_to_spatial(t2_in_fre) t2_in = outputs t = t - 1 t2_out = t2_in else: t2_out = network(t2_in, t1)['img_out'] t1_out = None if args.input_normalize == "mean_std": 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) 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_first_step_recon_img = (np.clip(first_step_recon.data.cpu().numpy()[0, 0] * t2_std + t2_mean, 0, 1) * 255).astype(np.uint8) else: if t1_out is not None: t1_img = (np.clip(t1.data.cpu().numpy()[0, 0], 0, 1) * 255).astype(np.uint8) t1_out_img = (np.clip(t1_out.data.cpu().numpy()[0, 0], 0, 1) * 255).astype(np.uint8) t1_krecon_img = (np.clip(t1_krecon.data.cpu().numpy()[0, 0], 0, 1) * 255).astype(np.uint8) t2_img = (np.clip(t2.data.cpu().numpy()[0, 0], 0, 1) * 255).astype(np.uint8) t2_out_img = (np.clip(t2_out.data.cpu().numpy()[0, 0], 0, 1) * 255).astype(np.uint8) t2_krecon_img = (np.clip(t2_krecon.data.cpu().numpy()[0, 0], 0, 1) * 255).astype(np.uint8) t2_first_step_recon_img = (np.clip(first_step_recon.data.cpu().numpy()[0, 0], 0, 1) * 255).astype(np.uint8) if t1_out is not None: MSE = mean_squared_error(t1_img, t1_out_img) PSNR = peak_signal_noise_ratio(t1_img, t1_out_img) SSIM = structural_similarity(t1_img, t1_out_img) t1_MSE_all.append(MSE) t1_PSNR_all.append(PSNR) t1_SSIM_all.append(SSIM) MSE = mean_squared_error(t1_img, t1_krecon_img) PSNR = peak_signal_noise_ratio(t1_img, t1_krecon_img) SSIM = structural_similarity(t1_img, t1_krecon_img) t1_MSE_krecon.append(MSE) t1_PSNR_krecon.append(PSNR) t1_SSIM_krecon.append(SSIM) if t2_out is not None: 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) MSE = mean_squared_error(t2_img, t2_first_step_recon_img) PSNR = peak_signal_noise_ratio(t2_img, t2_first_step_recon_img) SSIM = structural_similarity(t2_img, t2_first_step_recon_img) t2_MSE_first_step.append(MSE) t2_PSNR_first_step.append(PSNR) t2_SSIM_first_step.append(SSIM) MSE = mean_squared_error(t2_img, t2_krecon_img) PSNR = peak_signal_noise_ratio(t2_img, t2_krecon_img) SSIM = structural_similarity(t2_img, t2_krecon_img) t2_MSE_krecon.append(MSE) t2_PSNR_krecon.append(PSNR) t2_SSIM_krecon.append(SSIM) ids += 1 if ids > 100: break if t1_out is not None: t1_mse = np.array(t1_MSE_all).mean() t1_psnr = np.array(t1_PSNR_all).mean() t1_ssim = np.array(t1_SSIM_all).mean() t1_krecon_mse = np.array(t1_MSE_krecon).mean() t1_krecon_psnr = np.array(t1_PSNR_krecon).mean() t1_krecon_ssim = np.array(t1_SSIM_krecon).mean() t2_mse = np.array(t2_MSE_all).mean() t2_psnr = np.array(t2_PSNR_all).mean() t2_ssim = np.array(t2_SSIM_all).mean() t2_first_step_mse = np.array(t2_MSE_first_step).mean() t2_first_step_psnr = np.array(t2_PSNR_first_step).mean() t2_first_step_ssim = np.array(t2_SSIM_first_step).mean() t2_krecon_mse = np.array(t2_MSE_krecon).mean() t2_krecon_psnr = np.array(t2_PSNR_krecon).mean() t2_krecon_ssim = np.array(t2_SSIM_krecon).mean() if t2_psnr > best_status['T2_PSNR']: best_status = {'T2_NMSE': t2_mse, 'T2_PSNR': t2_psnr, 'T2_SSIM': t2_ssim} best_checkpoint_path = os.path.join(snapshot_path, 'best_checkpoint.pth') torch.save({'network': network.state_dict()}, best_checkpoint_path) print('New Best Network:') logging.info(f"[T2 First MRI:] average MSE: {t2_first_step_mse} average PSNR: {t2_first_step_psnr} average SSIM: {t2_first_step_ssim}") logging.info(f"[T2 MRI:] average MSE: {t2_mse} average PSNR: {t2_psnr} average SSIM: {t2_ssim}") print("Snapshot_path = ", snapshot_path) if iter_num > max_iterations: break print(best_status) save_mode_path = os.path.join(snapshot_path, 'iter_' + str(max_iterations) + '.pth') torch.save({'network': network.state_dict()}, save_mode_path) logging.info("save model to {}".format(save_mode_path)) writer.close()