def test(args):
model = MIRNet_DGF()
# summary(model,[[3,128,128],[0]])
# exit()
if args.data_type == 'rgb':
load_data = load_data_split
elif args.data_type == 'filter':
load_data = load_data_filter
checkpoint_dir = args.checkpoint
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# try:
checkpoint = load_checkpoint(checkpoint_dir, device == 'cuda', 'latest')
start_epoch = checkpoint['epoch']
global_step = checkpoint['global_iter']
state_dict = checkpoint['state_dict']
model.load_state_dict(state_dict)
print('=> loaded checkpoint (epoch {}, global_step {})'.format(start_epoch, global_step))
# except:
# print('=> no checkpoint file to be loaded.') # model.load_state_dict(state_dict)
# exit(1)
model.eval()
model = model.to(device)
trans = transforms.ToPILImage()
torch.manual_seed(0)
all_noisy_imgs = scipy.io.loadmat(args.noise_dir)['siddplus_valid_noisy_srgb']
all_clean_imgs = scipy.io.loadmat(args.gt)['siddplus_valid_gt_srgb']
# noisy_path = sorted(glob.glob(args.noise_dir+ "/*.png"))
# clean_path = [ i.replace("noisy","clean") for i in noisy_path]
i_imgs, _,_,_ = all_noisy_imgs.shape
psnrs = []
ssims = []
# print(noisy_path)
for i_img in range(i_imgs):
noise = transforms.ToTensor()(Image.fromarray(all_noisy_imgs[i_img]))
image_noise, image_noise_hr = load_data(noise, args.burst_length)
image_noise_hr = image_noise_hr.to(device)
burst_noise = image_noise.to(device)
begin = time.time()
_, pred = model(burst_noise,image_noise_hr)
pred = pred.detach().cpu()
gt = transforms.ToTensor()((Image.fromarray(all_clean_imgs[i_img])))
gt = gt.unsqueeze(0)
psnr_t = calculate_psnr(pred, gt)
ssim_t = calculate_ssim(pred, gt)
psnrs.append(psnr_t)
ssims.append(ssim_t)
print(i_img, " UP : PSNR : ", str(psnr_t), " : SSIM : ", str(ssim_t))
if args.save_img != '':
if not os.path.exists(args.save_img):
os.makedirs(args.save_img)
plt.figure(figsize=(15, 15))
plt.imshow(np.array(trans(pred[0])))
plt.title("denoise KPN DGF " + args.model_type, fontsize=25)
image_name = str(i_img) + "_"
plt.axis("off")
plt.suptitle(image_name + " UP : PSNR : " + str(psnr_t) + " : SSIM : " + str(ssim_t), fontsize=25)
plt.savefig(os.path.join(args.save_img, image_name + "_" + args.checkpoint + '.png'), pad_inches=0)
print(" AVG : PSNR : "+ str(np.mean(psnrs))+" : SSIM : "+ str(np.mean(ssims)))
def test(args):
model = MIRNet()
save_img = args.save_img
checkpoint_dir = "checkpoints/mir"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# try:
checkpoint = load_checkpoint(checkpoint_dir, device == 'cuda', 'latest')
start_epoch = checkpoint['epoch']
global_step = checkpoint['global_iter']
state_dict = checkpoint['state_dict']
model.load_state_dict(state_dict)
print('=> loaded checkpoint (epoch {}, global_step {})'.format(
start_epoch, global_step))
# except:
# print('=> no checkpoint file to be loaded.') # model.load_state_dict(state_dict)
# exit(1)
model.eval()
model = model.to(device)
trans = transforms.ToPILImage()
torch.manual_seed(0)
noisy_path = sorted(glob.glob(args.noise_dir + "/2_*.png"))
clean_path = [i.replace("noisy", "clean") for i in noisy_path]
print(noisy_path)
for i in range(len(noisy_path)):
noise = transforms.ToTensor()(Image.open(
noisy_path[i]).convert('RGB'))[:, 0:args.image_size,
0:args.image_size].unsqueeze(0)
noise = noise.to(device)
begin = time.time()
print(noise.size())
pred = model(noise)
pred = pred.detach().cpu()
gt = transforms.ToTensor()(Image.open(
clean_path[i]).convert('RGB'))[:, 0:args.image_size,
0:args.image_size]
gt = gt.unsqueeze(0)
psnr_t = calculate_psnr(pred, gt)
ssim_t = calculate_ssim(pred, gt)
print(i, " UP : PSNR : ", str(psnr_t), " : SSIM : ", str(ssim_t))
if save_img != '':
if not os.path.exists(args.save_img):
os.makedirs(args.save_img)
plt.figure(figsize=(15, 15))
plt.imshow(np.array(trans(pred[0])))
plt.title("denoise KPN DGF " + args.model_type, fontsize=25)
image_name = noisy_path[i].split("/")[-1].split(".")[0]
plt.axis("off")
plt.suptitle(image_name + " UP : PSNR : " + str(psnr_t) +
" : SSIM : " + str(ssim_t),
fontsize=25)
plt.savefig(os.path.join(
args.save_img, image_name + "_" + args.checkpoint + '.png'),
pad_inches=0)
def test(args):
if args.model_type == "DGF":
model = MIRNet_DGF(n_colors=args.n_colors,
out_channels=args.out_channels)
elif args.model_type == "noise":
model = MIRNet_noise(n_colors=args.n_colors,
out_channels=args.out_channels)
else:
print(" Model type not valid")
return
checkpoint_dir = args.checkpoint
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# try:
checkpoint = load_checkpoint(checkpoint_dir, device == 'cuda', 'latest')
start_epoch = checkpoint['epoch']
global_step = checkpoint['global_iter']
state_dict = checkpoint['state_dict']
# new_state_dict = OrderedDict()
# for k, v in state_dict.items():
# name = "model." + k # remove `module.`
# new_state_dict[name] = v
model.load_state_dict(state_dict)
print('=> loaded checkpoint (epoch {}, global_step {})'.format(
start_epoch, global_step))
# except:
# print('=> no checkpoint file to be loaded.') # model.load_state_dict(state_dict)
# exit(1)
model.eval()
model = model.to(device)
trans = transforms.ToPILImage()
torch.manual_seed(0)
all_noisy_imgs = scipy.io.loadmat(
args.noise_dir)['ValidationNoisyBlocksRaw']
all_clean_imgs = scipy.io.loadmat(args.gt_dir)['ValidationGtBlocksRaw']
# noisy_path = sorted(glob.glob(args.noise_dir+ "/*.png"))
# clean_path = [ i.replace("noisy","clean") for i in noisy_path]
i_imgs, i_blocks, _, _ = all_noisy_imgs.shape
psnrs = []
ssims = []
# print(noisy_path)
for i_img in range(i_imgs):
for i_block in range(i_blocks):
noise = transforms.ToTensor()(pack_raw(
all_noisy_imgs[i_img][i_block]))
image_noise, image_noise_hr = load_data(noise, args.burst_length)
image_noise_hr = image_noise_hr.to(device)
burst_noise = image_noise.to(device)
begin = time.time()
_, pred = model(burst_noise, image_noise_hr)
pred = pred.detach().cpu()
gt = transforms.ToTensor()(
(pack_raw(all_clean_imgs[i_img][i_block])))
gt = gt.unsqueeze(0)
psnr_t = calculate_psnr(pred, gt)
ssim_t = calculate_ssim(pred, gt)
psnrs.append(psnr_t)
ssims.append(ssim_t)
print(i_img, " UP : PSNR : ", str(psnr_t), " : SSIM : ",
str(ssim_t))
if args.save_img != '':
if not os.path.exists(args.save_img):
os.makedirs(args.save_img)
plt.figure(figsize=(15, 15))
plt.imshow(np.array(trans(pred[0])))
plt.title("denoise KPN DGF " + args.model_type, fontsize=25)
image_name = str(i_img)
plt.axis("off")
plt.suptitle(image_name + " UP : PSNR : " + str(psnr_t) +
" : SSIM : " + str(ssim_t),
fontsize=25)
plt.savefig(os.path.join(
args.save_img,
image_name + "_" + args.checkpoint + '.png'),
pad_inches=0)
print(" AVG : PSNR : " + str(np.mean(psnrs)) + " : SSIM : " +
str(np.mean(ssims)))