import os import sys from tqdm import tqdm from tensorboardX import SummaryWriter import logging import time import torch.optim as optim from torch.utils.data import DataLoader from networks.mynet import TwoBranch from networks_time.mynet import DiffTwoBranch from utils.option import args from dataloaders.fastmri import build_dataset from frequency_diffusion.degradation.k_degradation import get_ksu_kernel, apply_tofre, apply_to_spatial from utils.lpips import LPIPS from utils.metric import nmse, psnr, ssim, AverageMeter from collections import defaultdict 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 from utils.utils import * # --num_timesteps 30 --image_size 320 --use_kspace True --use_time_model True --ACCELERATIONS 4 --gpu 0 --phase train DEBUG = False use_kspace = args.use_kspace frequency_distortion = False use_time_model = args.use_time_model num_timesteps = args.num_timesteps image_size = 320 distortion_sigma = 10 / 255 if args.phase == 'test': kspace_masks = np.load(f"./dataloaders/example_mask/kspace_{args.ACCELERATIONS[0]}_mask.npy") 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], ) # args.ACCELERATIONS = [4] or [8] np.save(f"./dataloaders/example_mask/kspace_{args.ACCELERATIONS[0]}_mask.npy", kspace_masks) kspace_masks = torch.from_numpy(np.asarray(kspace_masks)).cuda() print("kspace kernels shape:", kspace_masks.shape) # (1, 1, 320, 320) @torch.no_grad() def evaluate(model, data_loader, device): model.eval() nmse_meter, psnr_meter, ssim_meter = AverageMeter(), AverageMeter(), AverageMeter() direct_nmse, direct_psnr, direct_ssim = AverageMeter(), AverageMeter(), AverageMeter() output_dic = defaultdict(dict) target_dic = defaultdict(dict) input_dic = defaultdict(dict) direct_dic = defaultdict(dict) for id, data in enumerate(data_loader): pd, pdfs, _ = data name = os.path.basename(pdfs[4][0]).split('.')[0] target = pdfs[1].to(device) mean, std = pdfs[2], pdfs[3] fname = pdfs[4] slice_num = pdfs[5] mean = mean.unsqueeze(1).unsqueeze(2).to(device) std = std.unsqueeze(1).unsqueeze(2).to(device) pd_img = pd[1].unsqueeze(1).to(device) pdfs_img = pdfs[0].unsqueeze(1).to(device) pdfs_img_origin = pdfs_img.clone() # Degradation if use_kspace: b = pd_img.size(0) t = torch.randint(num_timesteps - 1, num_timesteps, (b,), device=device).long() # t-1 mask = kspace_masks[t] fft, mask = apply_tofre(target.clone(), mask) fft = fft * mask + 0.0 pdfs_img = apply_to_spatial(fft) # print("pdfs_img shape:", pdfs_img.shape, pdfs_img.min(), pdfs_img.max()) while t >= 0: if use_time_model: outputs = model(pdfs_img, pd_img, t)['img_out'] else: outputs = model(pdfs_img, pd_img)['img_out'] if t == num_timesteps - 1: direct_recon = outputs if t == 0: mask = kspace_masks[0] # last one pdfs_img = outputs else: k_full = kspace_masks[-1] faded_recon_sample_fre, k_full = apply_tofre(pdfs_img, 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 faded_recon_sample_fre = faded_recon_sample_fre + fre_amend # * (1-k_residual) # faded_recon_sample_fre = faded_recon_sample_fre * kt_sub_1 outputs = apply_to_spatial(faded_recon_sample_fre) pdfs_img = outputs t = t - 1 else: outputs = model(pdfs_img, pd_img)['img_out'] target = target * std + mean inputs = pdfs_img_origin.squeeze(1) * std + mean outputs = outputs.squeeze(1) * std + mean direct_recon = direct_recon.squeeze(1) * std + mean # print("target/outputs shape: ", target.shape, outputs.shape) # our_nmse = nmse(target[0].cpu().numpy(), outputs[0].cpu().numpy()) # our_psnr = psnr(target[0].cpu().numpy(), outputs[0].cpu().numpy()) # our_ssim = ssim(target[0].cpu().numpy(), outputs[0].cpu().numpy()) # print('name:{}, slice:{}, nmse:{}, psnr:{}, ssim:{}'.format(name, slice_num[0], our_nmse, our_psnr, our_ssim)) for i, f in enumerate(fname): output_dic[f][slice_num[i]] = outputs[i] target_dic[f][slice_num[i]] = target[i] input_dic[f][slice_num[i]] = inputs[i] direct_dic[f][slice_num[i]] = direct_recon[i] if id > 100: break for name in output_dic.keys(): f_output = torch.stack([v for _, v in output_dic[name].items()]) f_target = torch.stack([v for _, v in target_dic[name].items()]) our_nmse = nmse(f_target.cpu().numpy(), f_output.cpu().numpy()) our_psnr = psnr(f_target.cpu().numpy(), f_output.cpu().numpy()) our_ssim = ssim(f_target.cpu().numpy(), f_output.cpu().numpy()) nmse_meter.update(our_nmse, 1) psnr_meter.update(our_psnr, 1) ssim_meter.update(our_ssim, 1) direct_nmse.update( nmse(f_target.cpu().numpy(), torch.stack([v for _, v in direct_dic[name].items()]).cpu().numpy()), 1) direct_psnr.update( psnr(f_target.cpu().numpy(), torch.stack([v for _, v in direct_dic[name].items()]).cpu().numpy()), 1) direct_ssim.update( ssim(f_target.cpu().numpy(), torch.stack([v for _, v in direct_dic[name].items()]).cpu().numpy()), 1) print("==> Evaluate Metric") print("Direct Results ----------") print("NMSE: {:.4}".format(direct_nmse.avg)) print("PSNR: {:.4}".format(direct_psnr.avg)) print("SSIM: {:.4}".format(direct_ssim.avg)) print("------------------") print("==> Evaluate Metric") print("Results ----------") print("NMSE: {:.4}".format(nmse_meter.avg)) print("PSNR: {:.4}".format(psnr_meter.avg)) print("SSIM: {:.4}".format(ssim_meter.avg)) print("------------------") model.train() return {'NMSE': nmse_meter.avg, 'PSNR': psnr_meter.avg, 'SSIM': ssim_meter.avg} 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 frequency_distortion: snapshot_path = snapshot_path.rstrip("/") + '_no_distortion/' 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() # network = build_model_from_name(args).cuda() device = torch.device('cuda') network.to(device) lpips_loss = LPIPS().eval().to(device) if len(args.gpu.split(',')) > 1: network = nn.DataParallel(network) # network = nn.SyncBatchNorm.convert_sync_batchnorm(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 = build_dataset(args, mode='train', use_kspace=use_kspace) db_test = build_dataset(args, mode='val', use_kspace=use_kspace) trainloader = DataLoader(db_train, batch_size=batch_size, shuffle=True, num_workers=4, pin_memory=True) testloader = DataLoader(db_test, batch_size=1, shuffle=False, num_workers=4, pin_memory=True) mask = None 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) scheduler1 = optim.lr_scheduler.StepLR(optimizer1, step_size=20000, gamma=0.5) writer = SummaryWriter(snapshot_path + '/log') iter_num = 0 max_epoch = max_iterations // len(trainloader) + 1 best_status = {'NMSE': 10000000, 'PSNR': 0, 'SSIM': 0} fft_weight = 0.01 criterion = nn.L1Loss().to(device, non_blocking=True) freloss = Frequency_Loss().to(device, non_blocking=True) for epoch_num in tqdm(range(max_epoch), ncols=70): time1 = time.time() start_time = time.time() for i_batch, sampled_batch in enumerate(trainloader): time2 = time.time() pd, pdfs, _ = sampled_batch target = pdfs[1] mean, std = pdfs[2], pdfs[3] pd_img = pd[1].unsqueeze(1) pdfs_img = pdfs[0].unsqueeze(1) target = target.unsqueeze(1) b = pd_img.size(0) pd_img = pd_img.to(device) # [4, 1, 320, 320] pdfs_img = pdfs_img.to(device) # [4, 1, 320, 320] target = target.to(device) # [4, 1, 320, 320] time3 = time.time() # Degradation if use_kspace: t = torch.randint(0, num_timesteps, (b,), device=device).long() mask = kspace_masks[t] fft, mask = apply_tofre(target.clone(), mask) fft = fft * mask # 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) pdfs_img = apply_to_spatial(fft) # breakpoint() if use_time_model: outputs = network(pdfs_img, pd_img, t) else: outputs = network(pdfs_img, pd_img) loss = criterion(outputs['img_out'], target) + \ fft_weight * freloss(outputs['img_fre'], target, mask) + \ criterion(outputs['img_fre'], target) + \ 0.01 * lpips_loss(outputs['img_out'], target).mean() 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() time5 = time.time() # summary iter_num = iter_num + 1 print_iter = 100 # if not DEBUG else 5 if iter_num % print_iter == 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() eval_result = evaluate(network, testloader, device) if eval_result['PSNR'] > best_status['PSNR']: best_status = {'NMSE': eval_result['NMSE'], 'PSNR': eval_result['PSNR'], 'SSIM': eval_result['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 saved:', best_checkpoint_path) logging.info( f"average MSE: {eval_result['NMSE']} average PSNR: {eval_result['PSNR']} average SSIM: {eval_result['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()