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 import matplotlib.pyplot as plt use_new_dataloader = True if use_new_dataloader: from dataloaders.m4raw_std_dataloader import M4Raw_TestSet, M4Raw_TrainSet, normalize, normalize_instance_dim else: from dataloaders.m4raw_dataloader import M4Raw_TestSet, M4Raw_TrainSet, normalize, normalize_instance_dim 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 # import imsave from skimage.io import imsave 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 * frequency_distortion = True num_timesteps = args.num_timesteps image_size = args.image_size distortion_sigma = 10/255 DEBUG = args.DEBUG use_kspace = args.use_kspace use_time_model = args.use_time_model # Baseline ------------------------------------ # NMSE: 3.3329 ± 0.4915 # PSNR: 34.0418 ± 0.6790 # SSIM: 0.8699 ± 0.0143 # ------------------------------------ # Save Path: model/FSMNet_m4raw_4x//result_case/ # ------------------------------------ # NMSE: 3.1966 ± 0.3939 # PSNR: 34.2115 ± 0.5978 # SSIM: 0.8910 ± 0.0118 # ------------------------------------ # Save Path: model/FSMNet_m4raw_4x_kspace//result_case/ # num_timesteps = 5 image_size = 240 # if args.phase == 'test': # kspace_masks = np.load(f"./dataloaders/example_mask/m4raw_{args.ACCELERATIONS[0]}_mask.npy") # kspace_masks = torch.from_numpy(np.asarray(kspace_masks)).cuda() # # else: use_in_mean_std = False # 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/m4raw_{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() print_i = 1 nmse_meter = AverageMeter() psnr_meter = AverageMeter() ssim_meter = AverageMeter() output_dic = defaultdict(dict) target_dic = defaultdict(dict) input_dic = defaultdict(dict) for id, sampled_batch in enumerate(data_loader): if use_new_dataloader: t1_img, t1_in = sampled_batch['t1'], sampled_batch['t1_in'] t2_img, t2_in = sampled_batch['t2'], sampled_batch['t2_in'] else: t1_img, t1_in = sampled_batch['ref_image_full'], sampled_batch['ref_image_sub'] t2_img, t2_in = sampled_batch['tag_image_full'], sampled_batch['tag_image_sub'] t1_img = t1_img.to(device) # t1_in = t1_in.to(device) t2_img = t2_img.to(device) t2_in = t2_in.to(device) mean, std = sampled_batch['t2_mean'], sampled_batch['t2_std'] fname = sampled_batch['fname'] slice_num = sampled_batch['slice'] mean = mean.unsqueeze(1).to(device) std = std.unsqueeze(1).to(device) t2_in_origin = t2_in.clone() # Degradation if use_kspace: b = 1 t = torch.randint(num_timesteps - 1, num_timesteps, (b,), device=device).long() # t-1 mask = kspace_masks[t] fft, mask = apply_tofre(t2_in.clone(), mask) # t2_img fft = fft * mask + 0.0 t2_in = apply_to_spatial(fft) # Save for debug: # put t2_in and t2_in_oringin side by side t2_in_origin = t2_in.clone() # print("test fft/mask = ", fft.shape, fft.shape) if use_in_mean_std: t2_in = t2_in * std + mean t2_img = t2_img * std + mean # print("Test After restore:", t2_in.shape, t2_img.shape) t2_in, mean, std = normalize_instance_dim(t2_in, eps=1e-11) t2_img = normalize(t2_img, mean=mean, stddev=std, eps=1e-11) t2_in = t2_in.float() t2_img = t2_img.float() mean = mean[0] std = std[0] # print("Test Re-normalize:", t2_in.shape, t2_img.shape) # print("in put t2_in shape:", t2_in.shape, t2_in_origin.shape) while t >= 0: if use_time_model: outputs = model(t2_in, t1_img, t)['img_out'] else: outputs = model(t2_in, t1_img)['img_out'] if t == 0: mask = kspace_masks[0] # last one t2_in = outputs else: k_full = kspace_masks[-1] faded_recon_sample_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 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) t2_in = outputs t = t-1 else: outputs = model(t2_in, t1_img)['img_out'] if print_i: t2_in_save = torch.cat([t2_in, t2_in_origin, t2_img], dim=3).cpu().numpy()[0, 0] t2_in_save = (t2_in_save - t2_in_save.min()) / (t2_in_save.max() - t2_in_save.min()) # t2_in_save = np.stack([t2_in_save, t2_in_save, t2_in_save], axis=2) # save to file os.makedirs("./debug", exist_ok=True) save_path = f"./debug/{use_kspace}_{fname[0]}_{slice_num[0]}.png" plt.imsave(save_path, t2_in_save, cmap='gray') print_i = 0 print("print_i") t2_img = t2_img.squeeze(1) * std + mean inputs = t2_in_origin.squeeze(1) * std + mean outputs = outputs.squeeze(1) * std + mean # print("output:", t2_img.shape, outputs.shape, inputs.shape) for i, f in enumerate(fname): output_dic[f][slice_num[i]] = outputs[i] target_dic[f][slice_num[i]] = t2_img[i] input_dic[f][slice_num[i]] = inputs[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) 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: if use_new_dataloader: snapshot_path = snapshot_path.rstrip("/") + f'_t{num_timesteps}_new_kspace/' else: snapshot_path = snapshot_path.rstrip("/") + f'_t{num_timesteps}/' if not isinstance(args.test_tag, type(None)): snapshot_path = snapshot_path.rstrip("/") + f'_{args.test_tag}/' 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: model = DiffTwoBranch(args).cuda() else: model = TwoBranch(args).cuda() # model = build_model_from_name(args).cuda() device = torch.device('cuda') model.to(device) lpips_loss = LPIPS().eval().to(device) if len(args.gpu.split(',')) > 1: model = nn.DataParallel(model) # model = nn.SyncBatchNorm.convert_sync_batchnorm(model) n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad) print('number of params: %.2f M' % (n_parameters / 1024 / 1024)) # db_train = M4Raw_TrainSet(args.root_path, args.MRIDOWN, use_kspace=use_kspace) # build_dataset(args, mode='train') # db_test = M4Raw_TestSet(args.root_path, args.MRIDOWN, use_kspace=use_kspace) # build_dataset(args, mode='val') db_train = M4Raw_TrainSet(args, use_kspace=use_kspace, DEBUG=DEBUG) # build_dataset(args, mode='train') db_test = M4Raw_TestSet(args, use_kspace=use_kspace, DEBUG=DEBUG) # 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) if args.phase == 'train': model.train() params = list(model.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) t = 0 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() # T1 is the reference image, T2 is the target image t1_img, t1_in = sampled_batch['t1'], sampled_batch['t1_in'] t2_img, t2_in = sampled_batch['t2'], sampled_batch['t2_in'] t1_img = t1_img.to(device) t1_in = t1_in.to(device) t2_img = t2_img.to(device) t2_in = t2_in.to(device) time3 = time.time() # Degradation if use_kspace: t2_origin = t2_in.clone() b = t1_in.size(0) t = torch.randint(0, num_timesteps, (b,), device=device).long() mask = kspace_masks[t] fft, mask = apply_tofre(t2_in.clone(), mask) # TODO t2_img 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) t2_in = apply_to_spatial(fft) if use_in_mean_std: mean, std = sampled_batch['t2_mean'], sampled_batch['t2_std'] mean = mean.unsqueeze(1).unsqueeze(1).to(device) std = std.unsqueeze(1).unsqueeze(1).to(device) t2_in = t2_in * std + mean t2_img = t2_img * std + mean # print("Train After restore:", t2_in.shape, t2_img.shape, mean.shape, std.shape) t2_in, mean, std = normalize_instance_dim(t2_in, eps=1e-11) t2_img = normalize(t2_img, mean=mean, stddev=std, eps=1e-11).detach() t2_in = t2_in.float() t2_img = t2_img.float() # print("Train After restore:", t2_in.shape, t2_img.shape) # breakpoint() if use_time_model: outputs = model(t2_in, t1_img, t) # ['img_out'] else: outputs = model(t2_in, t1_img) # ['img_out'] loss = criterion(outputs['img_out'], t2_img) + \ fft_weight * freloss(outputs['img_fre'], t2_img) + \ criterion(outputs['img_fre'], t2_img) # 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_(model.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': model.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(model, 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': model.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': model.state_dict()}, save_mode_path) logging.info("save model to {}".format(save_mode_path)) writer.close()