| """
|
| SAR-DDPM Inference on real SAR images.
|
| """
|
|
|
| import argparse
|
| import torch
|
| import os
|
| import cv2
|
| import numpy as np
|
|
|
| import torch.nn.functional as F
|
|
|
| from guided_diffusion import dist_util, logger
|
| from guided_diffusion.image_datasets import load_data
|
| from guided_diffusion.resample import create_named_schedule_sampler
|
| from guided_diffusion.script_util import (
|
| sr_model_and_diffusion_defaults,
|
| sr_create_model_and_diffusion,
|
| args_to_dict,
|
| add_dict_to_argparser,
|
| )
|
| from guided_diffusion.train_util import TrainLoop
|
| from torch.utils.data import DataLoader
|
| from torch.optim import AdamW
|
|
|
| from valdata import ValData, ValDataNew, ValDataNewReal
|
| from skimage.metrics import peak_signal_noise_ratio as psnr
|
| from skimage.metrics import structural_similarity as ssim
|
|
|
|
|
|
|
| val_dir = 'path_to_validation_data/'
|
| base_path = 'path_to_save_results/'
|
| resume_checkpoint_clean = './weights/sar_ddpm.pt'
|
|
|
|
|
|
|
|
|
| def main():
|
| args = create_argparser().parse_args()
|
|
|
| print(args)
|
|
|
|
|
| model_clean, diffusion = sr_create_model_and_diffusion(
|
| **args_to_dict(args, sr_model_and_diffusion_defaults().keys())
|
| )
|
|
|
|
|
| print(torch.device('cuda'))
|
|
|
| schedule_sampler = create_named_schedule_sampler(args.schedule_sampler, diffusion)
|
|
|
|
|
| val_data = DataLoader(ValDataNewReal(dataset_path=val_dir), batch_size=1, shuffle=False, num_workers=1)
|
|
|
| device0 = torch.device("cuda:0")
|
|
|
| model_clean.load_state_dict(torch.load(resume_checkpoint_clean, map_location="cuda:0"))
|
|
|
|
|
| model_clean.to(device0)
|
|
|
|
|
|
|
|
|
| params = list(model_clean.parameters())
|
|
|
| print('model clean device:')
|
| print(next(model_clean.parameters()).device)
|
|
|
|
|
|
|
| with torch.no_grad():
|
| number = 0
|
|
|
|
|
| for batch_id1, data_var in enumerate(val_data):
|
| number = number+1
|
| clean_batch, model_kwargs1 = data_var
|
|
|
| single_img = model_kwargs1['SR'].to(dist_util.dev())
|
|
|
| count = 0
|
| [t1,t2,max_r,max_c] = single_img.size()
|
|
|
| N =9
|
|
|
| val_inputv = single_img.clone()
|
|
|
| for row in range(0,max_r,100):
|
| for col in range(0,max_c,100):
|
|
|
|
|
| val_inputv[:,:,:row,:col] = single_img[:,:,max_r-row:,max_c-col:]
|
| val_inputv[:,:,row:,col:] = single_img[:,:,:max_r-row,:max_c-col]
|
| val_inputv[:,:,row:,:col] = single_img[:,:,:max_r-row,max_c-col:]
|
| val_inputv[:,:,:row,col:] = single_img[:,:,max_r-row:,:max_c-col]
|
|
|
| model_kwargs = {}
|
| for k, v in model_kwargs1.items():
|
| if('Index' in k):
|
| img_name=v
|
| elif('SR' in k):
|
| model_kwargs[k] = val_inputv.to(dist_util.dev())
|
| else:
|
| model_kwargs[k]= v.to(dist_util.dev())
|
|
|
|
|
|
|
| sample = diffusion.p_sample_loop(
|
| model_clean,
|
| (clean_batch.shape[0], 3, 256,256),
|
| clip_denoised=True,
|
| model_kwargs=model_kwargs,
|
| )
|
|
|
|
|
|
|
| if count==0:
|
| sample_new = (1.0/N)*sample
|
| else :
|
| sample_new[:,:,max_r-row:,max_c-col:] = sample_new[:,:,max_r-row:,max_c-col:] + (1.0/N)*sample[:,:,:row,:col]
|
| sample_new[:,:,:max_r-row,:max_c-col] = sample_new[:,:,:max_r-row,:max_c-col] + (1.0/N)*sample[:,:,row:,col:]
|
| sample_new[:,:,:max_r-row,max_c-col:] = sample_new[:,:,:max_r-row,max_c-col:] + (1.0/N)*sample[:,:,row:,:col]
|
| sample_new[:,:,max_r-row:,:max_c-col] = sample_new[:,:,max_r-row:,:max_c-col] + (1.0/N)*sample[:,:,:row,col:]
|
|
|
| count += 1
|
|
|
| sample_new = ((sample_new + 1) * 127.5)
|
| sample_new = sample_new.clamp(0, 255).to(torch.uint8)
|
| sample_new = sample_new.permute(0, 2, 3, 1)
|
| sample_new = sample_new.contiguous().cpu().numpy()
|
| sample_new = sample_new[0][:,:,::-1]
|
|
|
| sample_new = cv2.cvtColor(sample_new, cv2.COLOR_BGR2GRAY)
|
| print(img_name[0])
|
| cv2.imwrite(base_path+'pred_'+img_name[0],sample_new)
|
|
|
|
|
|
|
|
|
|
|
|
|
| def create_argparser():
|
| defaults = dict(
|
| data_dir= val_dir,
|
| schedule_sampler="uniform",
|
| lr=1e-4,
|
| weight_decay=0.0,
|
| lr_anneal_steps=0,
|
| batch_size=2,
|
| microbatch=1,
|
| ema_rate="0.9999",
|
| log_interval=100,
|
| save_interval=200,
|
| use_fp16=False,
|
| fp16_scale_growth=1e-3,
|
| )
|
| defaults.update(sr_model_and_diffusion_defaults())
|
| parser = argparse.ArgumentParser()
|
| add_dict_to_argparser(parser, defaults)
|
| return parser
|
|
|
| if __name__ == "__main__":
|
| main()
|
|
|
