def load_image(infile, n_channels):
img_name, ext = os.path.splitext(os.path.basename(infile))
print("Input File: %s" % (img_name + ext))
img_L = util.imread_uint(infile, n_channels=n_channels)
img_L = util.uint2single(img_L)
img_L = util.single2tensor4(img_L)
print('Brisque Score of input image : %f' % (brisq.get_score(infile)))
return img_name, ext, img_L
def add_JPEG_noise(img):
quality_factor = random.randint(30, 95)
img = cv2.cvtColor(util.single2uint(img), cv2.COLOR_RGB2BGR)
result, encimg = cv2.imencode(
'.jpg', img, [int(cv2.IMWRITE_JPEG_QUALITY), quality_factor])
img = cv2.imdecode(encimg, 1)
img = cv2.cvtColor(util.uint2single(img), cv2.COLOR_BGR2RGB)
return img
def __getitem__(self, index: int) -> Dict[str, Union[str, torch.Tensor]]:
# get H image
img_path = self.img_paths[index]
img_H = util.imread_uint(img_path, self.n_channels)
H, W = img_H.shape[:2]
if self.opt['phase'] == 'train':
self.count += 1
# crop
rnd_h = random.randint(0, max(0, H - self.patch_size))
rnd_w = random.randint(0, max(0, W - self.patch_size))
patch_H = img_H[rnd_h:rnd_h + self.patch_size,
rnd_w:rnd_w + self.patch_size, :]
# augmentation
patch_H = util.augment_img(patch_H, mode=np.random.randint(0, 8))
# HWC to CHW, numpy(uint) to tensor
img_H = util.uint2tensor3(patch_H)
img_L: torch.Tensor = img_H.clone()
# get noise level
noise_level: torch.FloatTensor = torch.FloatTensor(
[np.random.uniform(self.sigma[0], self.sigma[1])]) / 255.0
# add noise
noise = torch.randn(img_L.size()).mul_(noise_level).float()
img_L.add_(noise)
else:
img_H = util.uint2single(img_H)
img_L = np.copy(img_H)
# add noise
np.random.seed(seed=0)
img_L += np.random.normal(0, self.sigma / 255.0, img_L.shape)
noise_level = torch.FloatTensor([self.sigma / 255.0])
img_H, img_L = util.single2tensor3(img_H), util.single2tensor3(
img_L)
return {
'y': img_L,
'y_gt': img_H,
'sigma': noise_level.unsqueeze(1).unsqueeze(1),
'path': img_path
}
def __getitem__(self, index):
# ------------------------------------
# get H image
# ------------------------------------
H_path = self.paths_H[index]
H_resolution = [int(s) for s in self.sizes_H[index].split('_')]
img_H = _read_img_lmdb(self.h_env, H_path, H_resolution)
img_H = util.uint2single(img_H)
img_H = util.modcrop(img_H, self.sf)
# ----------------------------------------------
# get down4 sharp image(img_h0) to be target in train for deblur
# ----------------------------------------------
img_h0 = util.imresize_np(img_H, 1 / 4, True)
# ------------------------------------
# get L image
# ------------------------------------
L_path = self.paths_L[index]
L_resolution = [int(s) for s in self.sizes_L[index].split('_')]
img_L = _read_img_lmdb(self.l_env, L_path, L_resolution)
img_L = util.uint2single(img_L)
# ------------------------------------
# L/H pairs, HWC to CHW, numpy to tensor
# ------------------------------------
img_deblur = [img_L,img_h0]
img_deblur = util.generate_pyramid(*img_deblur,n_scales=3)
img_deblur =util.np2tensor(*img_deblur)
img_L = img_deblur[0]
img_H = util.single2tensor3(img_H)
img_L0 = img_L[0]
return {'L0':img_L0, 'L': img_L, 'H': img_H, 'L_path': L_path, 'H_path': H_path, 'ls': img_L, 'hs': img_deblur[1]}
def __getitem__(self, index):
# ------------------------------------
# get L image
# ------------------------------------
L_path = self.paths_L[index]
L_resolution = [int(s) for s in self.sizes_L[index].split('_')]
img_L = _read_img_lmdb(self.l_env, L_path, L_resolution)
img_L = util.uint2single(img_L)
# ------------------------------------
# L/H pairs, HWC to CHW, numpy to tensor
# ------------------------------------
img_deblur = [img_L] #,img_h0]
img_deblur = util.generate_pyramid(*img_deblur, n_scales=3)
img_deblur = util.np2tensor(*img_deblur)
img_L = img_deblur[0]
return {'L0': img_L[0], 'ls': img_L, 'L': img_L, 'L_path': L_path}
def __getitem__(self, index):
H_path = 'toy.png'
if self.opt['phase'] == 'train':
patch_H = self.H_data[index]
# --------------------------------
# augmentation - flip and/or rotate
# --------------------------------
mode = np.random.randint(0, 8)
patch_H = util.augment_img(patch_H, mode=mode)
patch_H = util.uint2tensor3(patch_H)
patch_L = patch_H.clone()
# ------------------------------------
# add noise
# ------------------------------------
noise = torch.randn(patch_L.size()).mul_(self.sigma / 255.0)
patch_L.add_(noise)
else:
H_path = self.paths_H[index]
img_H = util.imread_uint(H_path, self.n_channels)
img_H = util.uint2single(img_H)
img_L = np.copy(img_H)
# ------------------------------------
# add noise
# ------------------------------------
np.random.seed(seed=0)
img_L += np.random.normal(0, self.sigma_test / 255.0, img_L.shape)
patch_L, patch_H = util.single2tensor3(img_L), util.single2tensor3(
img_H)
L_path = H_path
return {'L': patch_L, 'H': patch_H, 'L_path': L_path, 'H_path': H_path}
def main():
#0. global config
#scale factor
sf = 4
stage = 8
patch_size = [32, 32]
patch_num = [3, 3]
#1. local PSF
#shape: gx,gy,kw,kw,3
all_PSFs = load_kernels('./data')
#2. local model
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = net(n_iter=8,
h_nc=64,
in_nc=4,
out_nc=3,
nc=[64, 128, 256, 512],
nb=2,
sf=sf,
act_mode="R",
downsample_mode='strideconv',
upsample_mode="convtranspose")
#model.proj.load_state_dict(torch.load('./data/usrnet_pretrain.pth'),strict=True)
model.train()
for _, v in model.named_parameters():
v.requires_grad = True
model = model.to(device)
#positional lambda, mu for HQS, set as free trainable parameters here.
ab_buffer = np.ones(
(patch_num[0], patch_num[1], 2 * stage, 3), dtype=np.float32) * 0.1
ab = torch.tensor(ab_buffer, device=device, requires_grad=True)
params = []
params += [{"params": [ab], "lr": 0.0005}]
for key, value in model.named_parameters():
params += [{"params": [value], "lr": 0.0001}]
optimizer = torch.optim.Adam(params, lr=0.0001, betas=(0.9, 0.999))
scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
step_size=1000,
gamma=0.9)
#3.load training data
imgs_H = glob.glob('/home/xiu/databag/deblur/images/DIV2K_train/*.png',
recursive=True)
imgs_L = glob.glob('/home/xiu/databag/deblur/images/DIV2K_lr/*.png',
recursive=True)
imgs_H.sort()
imgs_L.sort()
global_iter = 0
N_maxiter = 200000
#def get_train_pairs()
for i in range(N_maxiter):
t0 = time.time()
#draw random image.
img_idx = np.random.randint(len(imgs_H))
img_H = cv2.imread(imgs_H[img_idx])
#img2 = imgs_L[img_idx]
#img_L = cv2.imread(img2)
#draw random patch from image
#a. without img_L
#draw random kernel
PSF_grid = draw_random_kernel(all_PSFs, patch_num)
patch_L, patch_H, patch_psf = draw_training_pair(
img_H, PSF_grid, sf, patch_num, patch_size)
#b. with img_L
#patch_L, patch_H, patch_psf,px_start, py_start,block_expand = draw_training_pair(img_H, PSF_grid, sf, patch_num, patch_size, img_L)
t_data = time.time() - t0
x = util.uint2single(patch_L)
x = util.single2tensor4(x)
x_gt = util.uint2single(patch_H)
x_gt = util.single2tensor4(x_gt)
k_local = []
for h_ in range(patch_num[1]):
for w_ in range(patch_num[0]):
k_local.append(util.single2tensor4(patch_psf[w_, h_]))
k = torch.cat(k_local, dim=0)
[x, x_gt, k] = [el.to(device) for el in [x, x_gt, k]]
ab_patch = F.softplus(ab)
ab_patch_v = []
for h_ in range(patch_num[1]):
for w_ in range(patch_num[0]):
ab_patch_v.append(ab_patch[w_:w_ + 1, h_])
ab_patch_v = torch.cat(ab_patch_v, dim=0)
x_E = model.forward_patchwise_SR(x, k, ab_patch_v, patch_num,
[patch_size[0], patch_size[1]], sf)
loss = F.l1_loss(x_E, x_gt)
optimizer.zero_grad()
loss.backward()
optimizer.step()
scheduler.step()
t_iter = time.time() - t0 - t_data
print('[iter:{}] loss:{:.4f}, data_time:{:.2f}s, net_time:{:.2f}s'.
format(global_iter + 1, loss.item(), t_data, t_iter))
patch_L = cv2.resize(patch_L,
dsize=None,
fx=sf,
fy=sf,
interpolation=cv2.INTER_NEAREST)
#patch_L = patch_L[block_expand*sf:-block_expand*sf,block_expand*sf:-block_expand*sf]
patch_E = util.tensor2uint((x_E))
show = np.hstack((patch_H, patch_L, patch_E))
cv2.imshow('H,L,E', show)
key = cv2.waitKey(1)
global_iter += 1
# for logging model weight.
# if global_iter % 100 ==0:
# torch.save(model.state_dict(),'./logs/uabcnet_{}.pth'.format(global_iter))
if key == ord('q'):
break
if key == ord('s'):
torch.save(model.state_dict(), './logs/uabcnet.pth')
torch.save(model.state_dict(), './logs/uabcnet.pth')
def main():
# ----------------------------------------
# Preparation
# ----------------------------------------
model_name = 'usrnet' # 'usrgan' | 'usrnet' | 'usrgan_tiny' | 'usrnet_tiny'
testset_name = 'set5' # test set, 'set5' | 'srbsd68'
test_sf = [4] if 'gan' in model_name else [
2, 3, 4
] # scale factor, from {1,2,3,4}
show_img = False # default: False
save_L = True # save LR image
save_E = True # save estimated image
save_LEH = False # save zoomed LR, E and H images
# ----------------------------------------
# load testing kernels
# ----------------------------------------
# kernels = hdf5storage.loadmat(os.path.join('kernels', 'kernels.mat'))['kernels']
kernels = loadmat(os.path.join('kernels', 'kernels_12.mat'))['kernels']
n_channels = 1 if 'gray' in model_name else 3 # 3 for color image, 1 for grayscale image
model_pool = 'model_zoo' # fixed
testsets = 'testsets' # fixed
results = 'results' # fixed
noise_level_img = 0 # fixed: 0, noise level for LR image
noise_level_model = noise_level_img # fixed, noise level of model, default 0
result_name = testset_name + '_' + model_name
model_path = os.path.join(model_pool, model_name + '.pth')
# ----------------------------------------
# L_path = H_path, E_path, logger
# ----------------------------------------
L_path = os.path.join(
testsets,
testset_name) # L_path and H_path, fixed, for Low-quality images
E_path = os.path.join(results,
result_name) # E_path, fixed, for Estimated images
util.mkdir(E_path)
logger_name = result_name
utils_logger.logger_info(logger_name,
log_path=os.path.join(E_path,
logger_name + '.log'))
logger = logging.getLogger(logger_name)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# ----------------------------------------
# load model
# ----------------------------------------
if 'tiny' in model_name:
model = net(n_iter=6,
h_nc=32,
in_nc=4,
out_nc=3,
nc=[16, 32, 64, 64],
nb=2,
act_mode="R",
downsample_mode='strideconv',
upsample_mode="convtranspose")
else:
model = net(n_iter=8,
h_nc=64,
in_nc=4,
out_nc=3,
nc=[64, 128, 256, 512],
nb=2,
act_mode="R",
downsample_mode='strideconv',
upsample_mode="convtranspose")
model.load_state_dict(torch.load(model_path), strict=True)
model.eval()
for key, v in model.named_parameters():
v.requires_grad = False
number_parameters = sum(map(lambda x: x.numel(), model.parameters()))
model = model.to(device)
logger.info('Model path: {:s}'.format(model_path))
logger.info('Params number: {}'.format(number_parameters))
logger.info('Model_name:{}, image sigma:{}'.format(model_name,
noise_level_img))
logger.info(L_path)
L_paths = util.get_image_paths(L_path)
# --------------------------------
# read images
# --------------------------------
test_results_ave = OrderedDict()
test_results_ave['psnr_sf_k'] = []
for sf in test_sf:
for k_index in range(kernels.shape[1]):
test_results = OrderedDict()
test_results['psnr'] = []
kernel = kernels[0, k_index].astype(np.float64)
## other kernels
# kernel = utils_deblur.blurkernel_synthesis(h=25) # motion kernel
# kernel = utils_deblur.fspecial('gaussian', 25, 1.6) # Gaussian kernel
# kernel = sr.shift_pixel(kernel, sf) # pixel shift; optional
# kernel /= np.sum(kernel)
util.surf(kernel) if show_img else None
idx = 0
for img in L_paths:
# --------------------------------
# (1) classical degradation, img_L
# --------------------------------
idx += 1
img_name, ext = os.path.splitext(os.path.basename(img))
img_H = util.imread_uint(
img, n_channels=n_channels) # HR image, int8
img_H = util.modcrop(img_H, np.lcm(sf, 8)) # modcrop
# generate degraded LR image
img_L = ndimage.filters.convolve(img_H,
kernel[..., np.newaxis],
mode='wrap') # blur
img_L = sr.downsample_np(
img_L, sf,
center=False) # downsample, standard s-fold downsampler
img_L = util.uint2single(img_L) # uint2single
np.random.seed(seed=0) # for reproducibility
img_L += np.random.normal(0, noise_level_img,
img_L.shape) # add AWGN
util.imshow(util.single2uint(img_L)) if show_img else None
x = util.single2tensor4(img_L)
k = util.single2tensor4(kernel[..., np.newaxis])
sigma = torch.tensor(noise_level_model).float().view(
[1, 1, 1, 1])
[x, k, sigma] = [el.to(device) for el in [x, k, sigma]]
# --------------------------------
# (2) inference
# --------------------------------
x = model(x, k, sf, sigma)
# --------------------------------
# (3) img_E
# --------------------------------
img_E = util.tensor2uint(x)
if save_E:
util.imsave(
img_E,
os.path.join(
E_path, img_name + '_x' + str(sf) + '_k' +
str(k_index + 1) + '_' + model_name + '.png'))
# --------------------------------
# (4) img_LEH
# --------------------------------
img_L = util.single2uint(img_L)
if save_LEH:
k_v = kernel / np.max(kernel) * 1.2
k_v = util.single2uint(
np.tile(k_v[..., np.newaxis], [1, 1, 3]))
k_v = cv2.resize(k_v, (3 * k_v.shape[1], 3 * k_v.shape[0]),
interpolation=cv2.INTER_NEAREST)
img_I = cv2.resize(
img_L, (sf * img_L.shape[1], sf * img_L.shape[0]),
interpolation=cv2.INTER_NEAREST)
img_I[:k_v.shape[0], -k_v.shape[1]:, :] = k_v
img_I[:img_L.shape[0], :img_L.shape[1], :] = img_L
util.imshow(np.concatenate([img_I, img_E, img_H], axis=1),
title='LR / Recovered / Ground-truth'
) if show_img else None
util.imsave(
np.concatenate([img_I, img_E, img_H], axis=1),
os.path.join(
E_path, img_name + '_x' + str(sf) + '_k' +
str(k_index + 1) + '_LEH.png'))
if save_L:
util.imsave(
img_L,
os.path.join(
E_path, img_name + '_x' + str(sf) + '_k' +
str(k_index + 1) + '_LR.png'))
psnr = util.calculate_psnr(
img_E, img_H, border=sf**2) # change with your own border
test_results['psnr'].append(psnr)
logger.info(
'{:->4d}--> {:>10s} -- x{:>2d} --k{:>2d} PSNR: {:.2f}dB'.
format(idx, img_name + ext, sf, k_index, psnr))
ave_psnr_k = sum(test_results['psnr']) / len(test_results['psnr'])
logger.info(
'------> Average PSNR(RGB) of ({}) scale factor: ({}), kernel: ({}) sigma: ({}): {:.2f} dB'
.format(testset_name, sf, k_index + 1, noise_level_model,
ave_psnr_k))
test_results_ave['psnr_sf_k'].append(ave_psnr_k)
logger.info(test_results_ave['psnr_sf_k'])
def main():
#0. global config
#scale factor
sf = 4
stage = 5
patch_size = [32, 32]
patch_num = [2, 2]
#1. local PSF
#shape: gx,gy,kw,kw,3
all_PSFs = load_kernels('./data')
#2. local model
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = net(n_iter=5,
h_nc=64,
in_nc=4,
out_nc=3,
nc=[64, 128, 256, 512],
nb=2,
sf=sf,
act_mode="R",
downsample_mode='strideconv',
upsample_mode="convtranspose")
#loaded_state_dict= torch.load('./data/uabcnet_final.pth')
loaded_state_dict = torch.load('./data/uabcnet_finetune.pth')
model.load_state_dict(loaded_state_dict, strict=True)
model.eval()
for _, v in model.named_parameters():
v.requires_grad = False
model = model.to(device)
#positional lambda, mu for HQS, set as free trainable parameters here.
ab_buffer = np.loadtxt('./data/ab_finetune.txt').reshape(
(patch_num[0], patch_num[1], 2 * stage, 3)).astype(np.float32)
#ab[2x2,2*stage,3]
#ab_buffer = np.ones((patch_num[0],patch_num[1],2*stage,3),dtype=np.float32)*0.1
ab = torch.tensor(ab_buffer, device=device, requires_grad=False)
ab = F.softplus(ab)
#3.load training data
imgs_H = glob.glob('./DIV2K_train/*.png', recursive=True)
imgs_H.sort()
global_iter = 0
N_maxiter = 1000
PSF_grid = using_AC254_lens(all_PSFs, patch_num)
all_PSNR = []
out_folder = 'finetune'
for i in range(N_maxiter):
#draw random image.
img_idx = np.random.randint(len(imgs_H))
img_H = cv2.imread(imgs_H[img_idx])
img_H = np.pad(img_H, [(12, 12), (12, 12), (0, 0)])
croppatch = imgpatch(img_H, 280, 280, 24)
patches = croppatch.crop(img_H, 1)
patch_E_list = []
patch_L_list = []
# patch_L,patch_H,patch_psf = draw_training_pair(img_H,PSF_grid,sf,patch_num,patch_size)
for piece in range(len(patches)):
patch_L, patch_H, patch_psf = draw_testing_pair(
patches[piece], PSF_grid, sf, patch_num, patch_size)
x = util.uint2single(patch_L)
x = util.single2tensor4(x)
x_gt = util.uint2single(patch_H)
x_gt = util.single2tensor4(x_gt)
k_local = []
for h_ in range(patch_num[1]):
for w_ in range(patch_num[0]):
k_local.append(util.single2tensor4(patch_psf[w_, h_]))
k = torch.cat(k_local, dim=0)
[x, x_gt, k] = [el.to(device) for el in [x, x_gt, k]]
ab_patch = F.softplus(ab)
ab_patch_v = []
for h_ in range(patch_num[1]):
for w_ in range(patch_num[0]):
ab_patch_v.append(ab_patch[w_:w_ + 1, h_])
ab_patch_v = torch.cat(ab_patch_v, dim=0)
x_E = model.forward_patchwise_SR(x, k, ab_patch_v, patch_num,
[patch_size[0], patch_size[1]],
sf)
patch_L = cv2.resize(patch_L,
dsize=None,
fx=sf,
fy=sf,
interpolation=cv2.INTER_NEAREST)
patch_E = util.tensor2uint((x_E))
patch_E_list.append(patch_E[np.newaxis, :])
patch_L_list.append((patch_L[None, ...]))
print(piece)
img_E = croppatch.merge(patch_E_list)
# croppatch_E = imgpatch(np.zeros_like(img_H), 256, 256, 0)
# croppatch_E.crop(np.zeros_like(img_H),1)
# img_E = croppatch_E.merge(patch_E_list)
psnr, ssim = cal_psnrssim(img_E, img_H, 255)
print(psnr)
print(ssim)
cv2.imwrite(
os.path.join('./result', out_folder,
'resultE-{:04d}.png'.format(i + 1)), img_E)
cv2.imwrite(
os.path.join('./result', out_folder,
'resultH-{:04d}.png'.format(i + 1)), img_H)
all_PSNR.append(psnr)
#show = np.hstack((patch_H,patch_L,patch_E))
np.savetxt(os.path.join('./result', out_folder, 'psnr.txt'), all_PSNR)
def __getitem__(self, index):
# -------------------------------------
# get H image
# -------------------------------------
H_path = self.paths_H[index]
img_H = util.imread_uint(H_path, self.n_channels)
L_path = H_path
if self.opt['phase'] == 'train':
"""
# --------------------------------
# get L/H/M patch pairs
# --------------------------------
"""
H, W = img_H.shape[:2]
# ---------------------------------
# randomly crop the patch
# ---------------------------------
rnd_h = random.randint(0, max(0, H - self.patch_size))
rnd_w = random.randint(0, max(0, W - self.patch_size))
patch_H = img_H[rnd_h:rnd_h + self.patch_size,
rnd_w:rnd_w + self.patch_size, :]
# ---------------------------------
# augmentation - flip, rotate
# ---------------------------------
mode = np.random.randint(0, 8)
patch_H = util.augment_img(patch_H, mode=mode)
# ---------------------------------
# HWC to CHW, numpy(uint) to tensor
# ---------------------------------
img_H = util.uint2tensor3(patch_H)
img_L = img_H.clone()
# ---------------------------------
# get noise level
# ---------------------------------
# noise_level = torch.FloatTensor([np.random.randint(self.sigma_min, self.sigma_max)])/255.0
noise_level = torch.FloatTensor(
[np.random.uniform(self.sigma_min, self.sigma_max)]) / 255.0
# ---------------------------------
# add noise
# ---------------------------------
noise = torch.randn(img_L.size()).mul_(noise_level).float()
img_L.add_(noise)
else:
"""
# --------------------------------
# get L/H/sigma image pairs
# --------------------------------
"""
img_H = util.uint2single(img_H)
img_L = np.copy(img_H)
np.random.seed(seed=0)
img_L += np.random.normal(0, self.sigma_test / 255.0, img_L.shape)
noise_level = torch.FloatTensor([self.sigma_test / 255.0])
# ---------------------------------
# L/H image pairs
# ---------------------------------
img_H, img_L = util.single2tensor3(img_H), util.single2tensor3(
img_L)
noise_level = noise_level.unsqueeze(1).unsqueeze(1)
return {
'L': img_L,
'H': img_H,
'C': noise_level,
'L_path': L_path,
'H_path': H_path
}
def main():
# ----------------------------------------
# Preparation
# ----------------------------------------
noise_level_img = 30 # noise level for noisy image
noise_level_model = noise_level_img # noise level for model
model_name = 'ffdnet_color' # 'ffdnet_gray' | 'ffdnet_color' | 'ffdnet_color_clip' | 'ffdnet_gray_clip'
testset_name = 'CBSD68' # test set, 'bsd68' | 'cbsd68' | 'set12'
need_degradation = True # default: True
show_img = False # default: False
task_current = 'dn' # 'dn' for denoising | 'sr' for super-resolution
sf = 1 # unused for denoising
if 'color' in model_name:
n_channels = 3 # setting for color image
nc = 96 # setting for color image
nb = 12 # setting for color image
else:
n_channels = 1 # setting for grayscale image
nc = 64 # setting for grayscale image
nb = 15 # setting for grayscale image
if 'clip' in model_name:
use_clip = True # clip the intensities into range of [0, 1]
else:
use_clip = False
model_pool = 'model_zoo' # fixed
testsets = 'testsets' # fixed
results = 'results' # fixed
result_name = testset_name + '_' + model_name
border = sf if task_current == 'sr' else 0 # shave boader to calculate PSNR and SSIM
model_path = os.path.join(model_pool, model_name + '.pth')
# ----------------------------------------
# L_path, E_path, H_path
# ----------------------------------------
L_path = os.path.join(testsets,
testset_name) # L_path, for Low-quality images
H_path = L_path # H_path, for High-quality images
E_path = os.path.join(results, result_name) # E_path, for Estimated images
util.mkdir(E_path)
if H_path == L_path:
need_degradation = True
logger_name = result_name
utils_logger.logger_info(logger_name,
log_path=os.path.join(E_path,
logger_name + '.log'))
logger = logging.getLogger(logger_name)
need_H = True if H_path is not None else False
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# ----------------------------------------
# load model
# ----------------------------------------
from models.network_ffdnet import FFDNet as net
model = net(in_nc=n_channels,
out_nc=n_channels,
nc=nc,
nb=nb,
act_mode='R')
model.load_state_dict(torch.load(model_path), strict=True)
model.eval()
for k, v in model.named_parameters():
v.requires_grad = False
model = model.to(device)
logger.info('Model path: {:s}'.format(model_path))
test_results = OrderedDict()
test_results['psnr'] = []
test_results['ssim'] = []
logger.info('model_name:{}, model sigma:{}, image sigma:{}'.format(
model_name, noise_level_img, noise_level_model))
logger.info(L_path)
L_paths = util.get_image_paths(L_path)
H_paths = util.get_image_paths(H_path) if need_H else None
for idx, img in enumerate(L_paths):
# ------------------------------------
# (1) img_L
# ------------------------------------
img_name, ext = os.path.splitext(os.path.basename(img))
# logger.info('{:->4d}--> {:>10s}'.format(idx+1, img_name+ext))
img_L = util.imread_uint(img, n_channels=n_channels)
img_L = util.uint2single(img_L)
if need_degradation: # degradation process
np.random.seed(seed=0) # for reproducibility
img_L += np.random.normal(0, noise_level_img / 255., img_L.shape)
if use_clip:
img_L = util.uint2single(util.single2uint(img_L))
util.imshow(util.single2uint(img_L),
title='Noisy image with noise level {}'.format(
noise_level_img)) if show_img else None
img_L = util.single2tensor4(img_L)
img_L = img_L.to(device)
sigma = torch.full((1, 1, 1, 1),
noise_level_model / 255.).type_as(img_L)
# ------------------------------------
# (2) img_E
# ------------------------------------
img_E = model(img_L, sigma)
img_E = util.tensor2uint(img_E)
if need_H:
# --------------------------------
# (3) img_H
# --------------------------------
img_H = util.imread_uint(H_paths[idx], n_channels=n_channels)
img_H = img_H.squeeze()
# --------------------------------
# PSNR and SSIM
# --------------------------------
psnr = util.calculate_psnr(img_E, img_H, border=border)
ssim = util.calculate_ssim(img_E, img_H, border=border)
test_results['psnr'].append(psnr)
test_results['ssim'].append(ssim)
logger.info('{:s} - PSNR: {:.2f} dB; SSIM: {:.4f}.'.format(
img_name + ext, psnr, ssim))
util.imshow(np.concatenate([img_E, img_H], axis=1),
title='Recovered / Ground-truth') if show_img else None
# ------------------------------------
# save results
# ------------------------------------
util.imsave(img_E, os.path.join(E_path, img_name + ext))
if need_H:
ave_psnr = sum(test_results['psnr']) / len(test_results['psnr'])
ave_ssim = sum(test_results['ssim']) / len(test_results['ssim'])
logger.info(
'Average PSNR/SSIM(RGB) - {} - PSNR: {:.2f} dB; SSIM: {:.4f}'.
format(result_name, ave_psnr, ave_ssim))
def main():
"""
# ----------------------------------------------------------------------------------
# In real applications, you should set proper
# - "noise_level_img": from [3, 25], set 3 for clean image, try 15 for very noisy LR images
# - "k" (or "kernel_width"): blur kernel is very important!!! kernel_width from [0.6, 3.0]
# to get the best performance.
# ----------------------------------------------------------------------------------
"""
##############################################################################
testset_name = 'Set3C' # set test set, 'set5' | 'srbsd68'
noise_level_img = 3 # set noise level of image, from [3, 25], set 3 for clean image
model_name = 'drunet_color' # 'ircnn_color' # set denoiser, | 'drunet_color' | 'ircnn_gray' | 'drunet_gray' | 'ircnn_color'
sf = 2 # set scale factor, 1, 2, 3, 4
iter_num = 24 # set number of iterations, default: 24 for SISR
# --------------------------------
# set blur kernel
# --------------------------------
kernel_width_default_x1234 = [
0.6, 0.9, 1.7, 2.2
] # Gaussian kernel widths for x1, x2, x3, x4
noise_level_model = noise_level_img / 255. # noise level of model
kernel_width = kernel_width_default_x1234[sf - 1]
"""
# set your own kernel width !!!!!!!!!!
"""
# kernel_width = 1.0
k = utils_deblur.fspecial('gaussian', 25, kernel_width)
k = sr.shift_pixel(k, sf) # shift the kernel
k /= np.sum(k)
##############################################################################
show_img = False
util.surf(k) if show_img else None
x8 = True # default: False, x8 to boost performance
modelSigma1 = 49 # set sigma_1, default: 49
modelSigma2 = max(sf, noise_level_model * 255.)
classical_degradation = True # set classical degradation or bicubic degradation
task_current = 'sr' # 'sr' for super-resolution
n_channels = 1 if 'gray' in model_name else 3 # fixed
model_zoo = 'model_zoo' # fixed
testsets = 'testsets' # fixed
results = 'results' # fixed
result_name = testset_name + '_realapplications_' + task_current + '_' + model_name
model_path = os.path.join(model_zoo, model_name + '.pth')
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
torch.cuda.empty_cache()
# ----------------------------------------
# L_path, E_path, H_path
# ----------------------------------------
L_path = os.path.join(testsets,
testset_name) # L_path, for Low-quality images
E_path = os.path.join(results, result_name) # E_path, for Estimated images
util.mkdir(E_path)
logger_name = result_name
utils_logger.logger_info(logger_name,
log_path=os.path.join(E_path,
logger_name + '.log'))
logger = logging.getLogger(logger_name)
# ----------------------------------------
# load model
# ----------------------------------------
if 'drunet' in model_name:
from models.network_unet import UNetRes as net
model = net(in_nc=n_channels + 1,
out_nc=n_channels,
nc=[64, 128, 256, 512],
nb=4,
act_mode='R',
downsample_mode="strideconv",
upsample_mode="convtranspose")
model.load_state_dict(torch.load(model_path), strict=True)
model.eval()
for _, v in model.named_parameters():
v.requires_grad = False
model = model.to(device)
elif 'ircnn' in model_name:
from models.network_dncnn import IRCNN as net
model = net(in_nc=n_channels, out_nc=n_channels, nc=64)
model25 = torch.load(model_path)
former_idx = 0
logger.info('model_name:{}, image sigma:{:.3f}, model sigma:{:.3f}'.format(
model_name, noise_level_img, noise_level_model))
logger.info('Model path: {:s}'.format(model_path))
logger.info(L_path)
L_paths = util.get_image_paths(L_path)
for idx, img in enumerate(L_paths):
# --------------------------------
# (1) get img_L
# --------------------------------
logger.info('Model path: {:s} Image: {:s}'.format(model_path, img))
img_name, ext = os.path.splitext(os.path.basename(img))
img_L = util.imread_uint(img, n_channels=n_channels)
img_L = util.uint2single(img_L)
img_L = util.modcrop(img_L, 8) # modcrop
# --------------------------------
# (2) get rhos and sigmas
# --------------------------------
rhos, sigmas = pnp.get_rho_sigma(sigma=max(0.255 / 255.,
noise_level_model),
iter_num=iter_num,
modelSigma1=modelSigma1,
modelSigma2=modelSigma2,
w=1)
rhos, sigmas = torch.tensor(rhos).to(device), torch.tensor(sigmas).to(
device)
# --------------------------------
# (3) initialize x, and pre-calculation
# --------------------------------
x = cv2.resize(img_L, (img_L.shape[1] * sf, img_L.shape[0] * sf),
interpolation=cv2.INTER_CUBIC)
if np.ndim(x) == 2:
x = x[..., None]
if classical_degradation:
x = sr.shift_pixel(x, sf)
x = util.single2tensor4(x).to(device)
img_L_tensor, k_tensor = util.single2tensor4(
img_L), util.single2tensor4(np.expand_dims(k, 2))
[k_tensor, img_L_tensor] = util.todevice([k_tensor, img_L_tensor],
device)
FB, FBC, F2B, FBFy = sr.pre_calculate(img_L_tensor, k_tensor, sf)
# --------------------------------
# (4) main iterations
# --------------------------------
for i in range(iter_num):
print('Iter: {} / {}'.format(i, iter_num))
# --------------------------------
# step 1, FFT
# --------------------------------
tau = rhos[i].float().repeat(1, 1, 1, 1)
x = sr.data_solution(x, FB, FBC, F2B, FBFy, tau, sf)
if 'ircnn' in model_name:
current_idx = np.int(
np.ceil(sigmas[i].cpu().numpy() * 255. / 2.) - 1)
if current_idx != former_idx:
model.load_state_dict(model25[str(current_idx)],
strict=True)
model.eval()
for _, v in model.named_parameters():
v.requires_grad = False
model = model.to(device)
former_idx = current_idx
# --------------------------------
# step 2, denoiser
# --------------------------------
if x8:
x = util.augment_img_tensor4(x, i % 8)
if 'drunet' in model_name:
x = torch.cat(
(x, sigmas[i].repeat(1, 1, x.shape[2], x.shape[3])), dim=1)
x = utils_model.test_mode(model,
x,
mode=2,
refield=64,
min_size=256,
modulo=16)
elif 'ircnn' in model_name:
x = model(x)
if x8:
if i % 8 == 3 or i % 8 == 5:
x = util.augment_img_tensor4(x, 8 - i % 8)
else:
x = util.augment_img_tensor4(x, i % 8)
# --------------------------------
# (3) img_E
# --------------------------------
img_E = util.tensor2uint(x)
util.imsave(
img_E,
os.path.join(E_path, img_name + '_x' + str(sf) + '_' + model_name +
'.png'))
def main():
# ----------------------------------------
# load kernels
# ----------------------------------------
PSF_grid = np.load('./data/AC254-075-A-ML-Zemax(ZMX).npz')['PSF']
PSF_grid = PSF_grid.astype(np.float32)
gx, gy = PSF_grid.shape[:2]
k_tensor = []
for yy in range(gy):
for xx in range(gx):
PSF_grid[xx, yy] = PSF_grid[xx, yy] / np.sum(PSF_grid[xx, yy],
axis=(0, 1))
k_tensor.append(util.single2tensor4(PSF_grid[xx, yy]))
k_tensor = torch.cat(k_tensor, dim=0)
inv_weight = util_deblur.get_inv_spatial_weight(k_tensor)
# ----------------------------------------
# load model
# ----------------------------------------
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = net(n_iter=8,
h_nc=64,
in_nc=4,
out_nc=3,
nc=[64, 128, 256, 512],
nb=2,
act_mode="R",
downsample_mode='strideconv',
upsample_mode="convtranspose")
model.proj.load_state_dict(torch.load('./data/usrnet_pretrain.pth'),
strict=True)
model.train()
for _, v in model.named_parameters():
v.requires_grad = True
model = model.to(device)
# ----------------------------------------
# load training data
# ----------------------------------------
imgs = glob.glob('./DIV2K_train/*.png', recursive=True)
imgs.sort()
# ----------------------------------------
# positional lambda\mu for HQS
# ----------------------------------------
stage = 8
ab_buffer = np.ones((gx, gy, 2 * stage, 3), dtype=np.float32) * 0.1
#ab_buffer[:,:,0,:] = 0.01
ab = torch.tensor(ab_buffer, device=device, requires_grad=True)
# ----------------------------------------
# build optimizer
# ----------------------------------------
params = []
params += [{"params": [ab], "lr": 0.0005}]
for key, value in model.named_parameters():
params += [{"params": [value], "lr": 0.0001}]
optimizer = torch.optim.Adam(params, lr=0.0001, betas=(0.9, 0.999))
scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
step_size=1000,
gamma=0.9)
patch_size = [128, 128]
expand = PSF_grid.shape[2] // 2
patch_num = [2, 2]
global_iter = 0
running = True
while running:
#alpha.beta
img_idx = np.random.randint(len(imgs))
img = imgs[img_idx]
img_H = cv2.imread(img)
w, h = img_H.shape[:2]
#focus on the edges
mode = np.random.randint(5)
px_start = np.random.randint(0, gx - patch_num[0] + 1)
py_start = np.random.randint(0, gy - patch_num[1] + 1)
if mode == 0:
px_start = 0
if mode == 1:
px_start = gx - patch_num[0]
if mode == 2:
py_start = 0
if mode == 3:
py_start = gy - patch_num[1]
x_start = np.random.randint(
0, w - patch_size[0] * patch_num[0] - expand * 2 + 1)
y_start = np.random.randint(
0, h - patch_size[1] * patch_num[1] - expand * 2 + 1)
PSF_patch = PSF_grid[px_start:px_start + patch_num[0],
py_start:py_start + patch_num[1]]
patch_H = img_H[x_start:x_start+patch_size[0]*patch_num[0]+expand*2,\
y_start:y_start+patch_size[1]*patch_num[1]+expand*2]
patch_L = util_deblur.blockConv2d(patch_H, PSF_patch, expand)
block_expand = max(patch_size[0] // 8, expand)
patch_L_wrap = util_deblur.wrap_boundary_liu(
patch_L, (patch_size[0] * patch_num[0] + block_expand * 2,
patch_size[1] * patch_num[1] + block_expand * 2))
patch_L_wrap = np.hstack(
(patch_L_wrap[:, -block_expand:, :],
patch_L_wrap[:, :patch_size[1] * patch_num[1] + block_expand, :]))
patch_L_wrap = np.vstack(
(patch_L_wrap[-block_expand:, :, :],
patch_L_wrap[:patch_size[0] * patch_num[0] + block_expand, :, :]))
x = util.uint2single(patch_L_wrap)
x = util.single2tensor4(x)
x_gt = util.uint2single(patch_H[expand:-expand, expand:-expand])
x_gt = util.single2tensor4(x_gt)
inv_weight_patch = torch.ones_like(x_gt)
k_local = []
for h_ in range(patch_num[1]):
for w_ in range(patch_num[0]):
inv_weight_patch[0, 0,
w_ * patch_size[0]:(w_ + 1) * patch_size[0],
h_ * patch_size[1]:(h_ + 1) *
patch_size[1]] = inv_weight[w_ +
h_ * patch_num[0],
0]
inv_weight_patch[0, 1,
w_ * patch_size[0]:(w_ + 1) * patch_size[0],
h_ * patch_size[1]:(h_ + 1) *
patch_size[1]] = inv_weight[w_ +
h_ * patch_num[0],
1]
inv_weight_patch[0, 2,
w_ * patch_size[0]:(w_ + 1) * patch_size[0],
h_ * patch_size[1]:(h_ + 1) *
patch_size[1]] = inv_weight[w_ +
h_ * patch_num[0],
2]
k_local.append(k_tensor[w_ + h_ * patch_num[0]:w_ +
h_ * patch_num[0] + 1])
k = torch.cat(k_local, dim=0)
[x, x_gt, k, inv_weight_patch
] = [el.to(device) for el in [x, x_gt, k, inv_weight_patch]]
ab_patch = F.softplus(ab[px_start:px_start + patch_num[0],
py_start:py_start + patch_num[1]])
cd = []
for h_ in range(patch_num[1]):
for w_ in range(patch_num[0]):
cd.append(ab_patch[w_:w_ + 1, h_])
cd = torch.cat(cd, dim=0)
x_E = model.forward_patchwise(x, k, cd, patch_num, patch_size)
predict = x_E[...,block_expand:block_expand+patch_size[0]*patch_num[0],\
block_expand:block_expand+patch_size[1]*patch_num[1]]
loss = F.l1_loss(predict.div(inv_weight_patch),
x_gt.div(inv_weight_patch))
optimizer.zero_grad()
loss.backward()
optimizer.step()
scheduler.step()
print('iter:{},loss {}'.format(global_iter + 1, loss.item()))
patch_L = patch_L_wrap.astype(np.uint8)
patch_E = util.tensor2uint(x_E)[block_expand:-block_expand,
block_expand:-block_expand]
show = np.hstack((patch_H[expand:-expand, expand:-expand],
patch_L[block_expand:-block_expand,
block_expand:-block_expand], patch_E))
cv2.imshow('HL', show)
key = cv2.waitKey(1)
global_iter += 1
#change the save period
if global_iter % 100 == 0:
ab_numpy = ab.detach().cpu().numpy().flatten()
torch.save(
model.state_dict(),
'./ZEMAX_model/usrnet_ZEMAX_iter{}.pth'.format(global_iter))
np.savetxt('./ZEMAX_model/ab_ZEMAX_iter{}.txt'.format(global_iter),
ab_numpy)
if key == ord('q'):
running = False
break
ab_numpy = ab.detach().cpu().numpy().flatten()
torch.save(model.state_dict(), './ZEMAX_model/usrnet_ZEMAX.pth')
np.savetxt('./ZEMAX_model/ab_ZEMAX.txt', ab_numpy)
def main():
# ----------------------------------------
# Preparation
# ----------------------------------------
noise_level_img = 0 # default: 0, noise level for LR image
noise_level_model = noise_level_img # noise level for model
model_name = 'srmdnf_x4' # 'srmd_x2' | 'srmd_x3' | 'srmd_x4' | 'srmdnf_x2' | 'srmdnf_x3' | 'srmdnf_x4'
testset_name = 'set5' # test set, 'set5' | 'srbsd68'
sf = [int(s) for s in re.findall(r'\d+', model_name)][0] # scale factor
x8 = False # default: False, x8 to boost performance
need_degradation = True # default: True, use degradation model to generate LR image
show_img = False # default: False
srmd_pca_path = os.path.join('kernels', 'srmd_pca_matlab.mat')
task_current = 'sr' # 'dn' for denoising | 'sr' for super-resolution
n_channels = 3 # fixed
in_nc = 18 if 'nf' in model_name else 19
nc = 128 # fixed, number of channels
nb = 12 # fixed, number of conv layers
model_pool = 'model_zoo' # fixed
testsets = 'testsets' # fixed
results = 'results' # fixed
result_name = testset_name + '_' + model_name
border = sf if task_current == 'sr' else 0 # shave boader to calculate PSNR and SSIM
model_path = os.path.join(model_pool, model_name + '.pth')
# ----------------------------------------
# L_path, E_path, H_path
# ----------------------------------------
L_path = os.path.join(testsets,
testset_name) # L_path, for Low-quality images
H_path = L_path # H_path, for High-quality images
E_path = os.path.join(results, result_name) # E_path, for Estimated images
util.mkdir(E_path)
if H_path == L_path:
need_degradation = True
logger_name = result_name
utils_logger.logger_info(logger_name,
log_path=os.path.join(E_path,
logger_name + '.log'))
logger = logging.getLogger(logger_name)
need_H = True if H_path is not None else False
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# ----------------------------------------
# load model
# ----------------------------------------
from models.network_srmd import SRMD as net
model = net(in_nc=in_nc,
out_nc=n_channels,
nc=nc,
nb=nb,
upscale=sf,
act_mode='R',
upsample_mode='pixelshuffle')
model.load_state_dict(torch.load(model_path), strict=False)
model.eval()
for k, v in model.named_parameters():
v.requires_grad = False
model = model.to(device)
logger.info('Model path: {:s}'.format(model_path))
number_parameters = sum(map(lambda x: x.numel(), model.parameters()))
logger.info('Params number: {}'.format(number_parameters))
test_results = OrderedDict()
test_results['psnr'] = []
test_results['ssim'] = []
test_results['psnr_y'] = []
test_results['ssim_y'] = []
logger.info('model_name:{}, model sigma:{}, image sigma:{}'.format(
model_name, noise_level_img, noise_level_model))
logger.info(L_path)
L_paths = util.get_image_paths(L_path)
H_paths = util.get_image_paths(H_path) if need_H else None
# ----------------------------------------
# kernel and PCA reduced feature
# ----------------------------------------
# kernel = sr.anisotropic_Gaussian(ksize=15, theta=np.pi, l1=4, l2=4)
kernel = utils_deblur.fspecial('gaussian', 15,
0.01) # Gaussian kernel, delta kernel 0.01
P = loadmat(srmd_pca_path)['P']
degradation_vector = np.dot(P, np.reshape(kernel, (-1), order="F"))
if 'nf' not in model_name: # noise-free SR
degradation_vector = np.append(degradation_vector,
noise_level_model / 255.)
degradation_vector = torch.from_numpy(degradation_vector).view(
1, -1, 1, 1).float()
for idx, img in enumerate(L_paths):
# ------------------------------------
# (1) img_L
# ------------------------------------
img_name, ext = os.path.splitext(os.path.basename(img))
# logger.info('{:->4d}--> {:>10s}'.format(idx+1, img_name+ext))
img_L = util.imread_uint(img, n_channels=n_channels)
img_L = util.uint2single(img_L)
# degradation process, blur + bicubic downsampling + Gaussian noise
if need_degradation:
img_L = util.modcrop(img_L, sf)
img_L = sr.srmd_degradation(
img_L, kernel, sf
) # equivalent to bicubic degradation if kernel is a delta kernel
np.random.seed(seed=0) # for reproducibility
img_L += np.random.normal(0, noise_level_img / 255., img_L.shape)
util.imshow(util.single2uint(img_L),
title='LR image with noise level {}'.format(
noise_level_img)) if show_img else None
img_L = util.single2tensor4(img_L)
degradation_map = degradation_vector.repeat(1, 1, img_L.size(-2),
img_L.size(-1))
img_L = torch.cat((img_L, degradation_map), dim=1)
img_L = img_L.to(device)
# ------------------------------------
# (2) img_E
# ------------------------------------
if not x8:
img_E = model(img_L)
else:
img_E = utils_model.test_mode(model, img_L, mode=3, sf=sf)
img_E = util.tensor2uint(img_E)
if need_H:
# --------------------------------
# (3) img_H
# --------------------------------
img_H = util.imread_uint(H_paths[idx], n_channels=n_channels)
img_H = img_H.squeeze()
img_H = util.modcrop(img_H, sf)
# --------------------------------
# PSNR and SSIM
# --------------------------------
psnr = util.calculate_psnr(img_E, img_H, border=border)
ssim = util.calculate_ssim(img_E, img_H, border=border)
test_results['psnr'].append(psnr)
test_results['ssim'].append(ssim)
logger.info('{:s} - PSNR: {:.2f} dB; SSIM: {:.4f}.'.format(
img_name + ext, psnr, ssim))
util.imshow(np.concatenate([img_E, img_H], axis=1),
title='Recovered / Ground-truth') if show_img else None
if np.ndim(img_H) == 3: # RGB image
img_E_y = util.rgb2ycbcr(img_E, only_y=True)
img_H_y = util.rgb2ycbcr(img_H, only_y=True)
psnr_y = util.calculate_psnr(img_E_y, img_H_y, border=border)
ssim_y = util.calculate_ssim(img_E_y, img_H_y, border=border)
test_results['psnr_y'].append(psnr_y)
test_results['ssim_y'].append(ssim_y)
# ------------------------------------
# save results
# ------------------------------------
util.imsave(img_E, os.path.join(E_path, img_name + '.png'))
if need_H:
ave_psnr = sum(test_results['psnr']) / len(test_results['psnr'])
ave_ssim = sum(test_results['ssim']) / len(test_results['ssim'])
logger.info(
'Average PSNR/SSIM(RGB) - {} - x{} --PSNR: {:.2f} dB; SSIM: {:.4f}'
.format(result_name, sf, ave_psnr, ave_ssim))
if np.ndim(img_H) == 3:
ave_psnr_y = sum(test_results['psnr_y']) / len(
test_results['psnr_y'])
ave_ssim_y = sum(test_results['ssim_y']) / len(
test_results['ssim_y'])
logger.info(
'Average PSNR/SSIM( Y ) - {} - x{} - PSNR: {:.2f} dB; SSIM: {:.4f}'
.format(result_name, sf, ave_psnr_y, ave_ssim_y))
def main():
#0. global config
sf = 4
stage = 8
patch_size = [32, 32]
patch_num = [2, 2]
#1. local PSF
all_PSFs = load_kernels('./data')
#2. load model
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = net(n_iter=8,
h_nc=64,
in_nc=4,
out_nc=3,
nc=[64, 128, 256, 512],
nb=2,
sf=sf,
act_mode="R",
downsample_mode='strideconv',
upsample_mode="convtranspose")
model.load_state_dict(torch.load('./data/uabcnet_final.pth'), strict=True)
model.train()
for _, v in model.named_parameters():
v.requires_grad = True
model = model.to(device)
#3. set up discriminator
model_D = gan.PatchDiscriminator(5)
model_D = model_D.to(device)
gan_loss = gan.GANLoss(mode='lsgan')
gan_loss = gan_loss.to(device)
fake_images = ImagePool(16)
#positional lambda, mu for HQS.
ab_buffer = np.zeros((patch_num[0], patch_num[1], 2 * stage, 3))
ab_buffer[:, :, ::2, :] = 0.01
ab_buffer[:, :, 1::2, :] = 0.1
ab = torch.tensor(ab_buffer,
dtype=torch.float32,
device=device,
requires_grad=True)
params = []
params += [{"params": [ab], "lr": 5e-4}]
for key, value in model.named_parameters():
params += [{"params": [value], "lr": 1e-5}]
#
params_D = []
params_D += list(model_D.parameters())
optimizer = torch.optim.Adam(params, lr=1e-4, betas=(0.9, 0.999))
optimizer_D = torch.optim.Adam(params_D, lr=1e-4, betas=(0.9, 0.999))
scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
step_size=1000,
gamma=0.9)
#3.load training data
imgs_H = glob.glob('/home/xiu/databag/deblur/images/DIV2K_train/*.png',
recursive=True)
imgs_H.sort()
global_iter = 0
N_maxiter = 200000
PSF_grid = draw_random_kernel(all_PSFs)
for i in range(N_maxiter):
t0 = time.time()
img_idx = np.random.randint(len(imgs_H))
img_H = cv2.imread(imgs_H[img_idx])
#draw random kernel
patch_L, patch_H, patch_psf = draw_training_pair(
img_H, PSF_grid, sf, patch_num, patch_size)
t_data = time.time() - t0
x = util.uint2single(patch_L)
x = util.single2tensor4(x)
x_gt = util.uint2single(patch_H)
x_gt = util.single2tensor4(x_gt)
k_local = []
for h_ in range(patch_num[1]):
for w_ in range(patch_num[0]):
k_local.append(util.single2tensor4(patch_psf[w_, h_]))
k = torch.cat(k_local, dim=0)
[x, x_gt, k] = [el.to(device) for el in [x, x_gt, k]]
ab_patch = F.softplus(ab)
ab_patch_v = []
for h_ in range(patch_num[1]):
for w_ in range(patch_num[0]):
ab_patch_v.append(ab_patch[w_:w_ + 1, h_])
ab_patch_v = torch.cat(ab_patch_v, dim=0)
x_E = model.forward_patchwise_SR(x, k, ab_patch_v, patch_num,
[patch_size[0], patch_size[1]], sf)
loss_l1 = F.l1_loss(x_E, x_gt)
loss_gan = gan_loss(model_D(x_E), True)
loss = loss_l1 + loss_gan
optimizer.zero_grad()
loss.backward()
optimizer.step()
pred_real = model_D(x_gt)
loss_D_real = gan_loss(pred_real, True)
fake = fake_images.query(x_E)
pred_fake = model_D(fake.detach())
loss_D_fake = gan_loss(pred_fake, False)
loss_D = (loss_D_fake + loss_D_real) * 0.5
optimizer_D.zero_grad()
loss_D.backward()
optimizer_D.step()
scheduler.step()
t_iter = time.time() - t0 - t_data
print('[iter:{}] loss:{:.4f}, data_time:{:.2f}s, net_time:{:.2f}s'.
format(global_iter + 1, loss.item(), t_data, t_iter))
patch_L = cv2.resize(patch_L,
dsize=None,
fx=sf,
fy=sf,
interpolation=cv2.INTER_NEAREST)
patch_E = util.tensor2uint((x_E))
show = np.hstack((patch_H, patch_L, patch_E))
cv2.imshow('H,L,E', show)
key = cv2.waitKey(1)
global_iter += 1
if key == ord('q'):
break
ab_numpy = ab.detach().cpu().numpy().flatten()
torch.save(model.state_dict(), './data/uabcnet_finetune.pth')
np.savetxt('./data/ab_finetune.txt', ab_numpy)
def __getitem__(self, index):
# ------------------------------------
# get H image
# ------------------------------------
H_path = self.paths_H[index]
img_H = util.imread_uint(H_path, self.n_channels)
img_H = util.uint2single(img_H)
# ------------------------------------
# modcrop for SR
# ------------------------------------
img_H = util.modcrop(img_H, self.sf)
# ------------------------------------
# sythesize L image via matlab's bicubic
# ------------------------------------
H, W, _ = img_H.shape
img_L = util.imresize_np(img_H, 1 / self.sf, True)
if self.opt['phase'] == 'train':
"""
# --------------------------------
# get L/H patch pairs
# --------------------------------
"""
H, W, C = img_L.shape
# --------------------------------
# randomly crop L patch
# --------------------------------
rnd_h = random.randint(0, max(0, H - self.L_size))
rnd_w = random.randint(0, max(0, W - self.L_size))
img_L = img_L[rnd_h:rnd_h + self.L_size,
rnd_w:rnd_w + self.L_size, :]
# --------------------------------
# crop corresponding H patch
# --------------------------------
rnd_h_H, rnd_w_H = int(rnd_h * self.sf), int(rnd_w * self.sf)
img_H = img_H[rnd_h_H:rnd_h_H + self.patch_size,
rnd_w_H:rnd_w_H + self.patch_size, :]
# --------------------------------
# augmentation - flip and/or rotate
# --------------------------------
mode = np.random.randint(0, 8)
img_L, img_H = util.augment_img(
img_L, mode=mode), util.augment_img(img_H, mode=mode)
# --------------------------------
# get patch pairs
# --------------------------------
img_H, img_L = util.single2tensor3(img_H), util.single2tensor3(
img_L)
# --------------------------------
# select noise level and get Gaussian noise
# --------------------------------
if random.random() < 0.1:
noise_level = torch.zeros(1).float()
else:
noise_level = torch.FloatTensor([
np.random.uniform(self.sigma_min, self.sigma_max)
]) / 255.0
# noise_level = torch.rand(1)*50/255.0
# noise_level = torch.min(torch.from_numpy(np.float32([7*np.random.chisquare(2.5)/255.0])),torch.Tensor([50./255.]))
else:
img_H, img_L = util.single2tensor3(img_H), util.single2tensor3(
img_L)
noise_level = torch.FloatTensor([self.sigma_test])
# ------------------------------------
# add noise
# ------------------------------------
noise = torch.randn(img_L.size()).mul_(noise_level).float()
img_L.add_(noise)
# ------------------------------------
# get noise level map M
# ------------------------------------
M_vector = noise_level.unsqueeze(1).unsqueeze(1)
M = M_vector.repeat(1, img_L.size()[-2], img_L.size()[-1])
"""
# -------------------------------------
# concat L and noise level map M
# -------------------------------------
"""
img_L = torch.cat((img_L, M), 0)
L_path = H_path
return {'L': img_L, 'H': img_H, 'L_path': L_path, 'H_path': H_path}
def main():
#0. global config
#scale factor
sf = 4
stage = 5
patch_size = [32,32]
patch_num = [2,2]
#1. local PSF
#shape: gx,gy,kw,kw,3
all_PSFs = load_kernels('./data')
#2. local model
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = net(n_iter=5, h_nc=64, in_nc=4, out_nc=3, nc=[64, 128, 256, 512],
nb=2,sf=sf, act_mode="R", downsample_mode='strideconv', upsample_mode="convtranspose")
model.load_state_dict(torch.load('./logs/uabcnet_final.pth'),strict=True)
model.train()
for _, v in model.named_parameters():
v.requires_grad = True
model = model.to(device)
#positional lambda, mu for HQS, set as free trainable parameters here.
#ab_buffer = np.loadtxt('./data/ab.txt').reshape((patch_num[0],patch_num[1],2*stage,3)).astype(np.float32)
ab_pretrain = np.loadtxt('./logs/ab_pretrain.txt').reshape((1,1,2*stage,3)).astype(np.float32)
ab_buffer = np.ones((patch_num[0],patch_num[1],2*stage,3),dtype=np.float32)
for xx in range(patch_num[0]):
for yy in range(patch_num[1]):
ab_buffer[xx,yy] = ab_pretrain[0,0]
ab = torch.tensor(ab_buffer,device=device,requires_grad=True)
params = []
params += [{"params":[ab],"lr":0.0001}]
for key,value in model.named_parameters():
params += [{"params":[value],"lr":1e-6}]
optimizer = torch.optim.Adam(params,lr=0.0001,betas=(0.9,0.999))
scheduler = torch.optim.lr_scheduler.StepLR(optimizer,step_size=1000,gamma=0.9)
#3.load training data
imgs_H = glob.glob('./DIV2K_train/*.png',recursive=True)
imgs_H.sort()
global_iter = 0
all_PSNR = []
N_maxiter = 4000
PSF_grid = draw_random_kernel(all_PSFs,patch_num)
#def get_train_pairs()
for i in range(N_maxiter):
t0 = time.time()
#draw random image.
img_idx = np.random.randint(len(imgs_H))
img_H = cv2.imread(imgs_H[img_idx])
patch_L,patch_H,patch_psf = draw_training_pair(img_H,PSF_grid,sf,patch_num,patch_size)
#b. with img_L
#patch_L, patch_H, patch_psf,px_start, py_start,block_expand = draw_training_pair(img_H, PSF_grid, sf, patch_num, patch_size, img_L)
t_data = time.time()-t0
x = util.uint2single(patch_L)
x = util.single2tensor4(x)
x_gt = util.uint2single(patch_H)
x_gt = util.single2tensor4(x_gt)
k_local = []
for h_ in range(patch_num[1]):
for w_ in range(patch_num[0]):
k_local.append(util.single2tensor4(patch_psf[w_,h_]))
k = torch.cat(k_local,dim=0)
[x,x_gt,k] = [el.to(device) for el in [x,x_gt,k]]
ab_patch = F.softplus(ab)
ab_patch_v = []
for h_ in range(patch_num[1]):
for w_ in range(patch_num[0]):
ab_patch_v.append(ab_patch[w_:w_+1,h_])
ab_patch_v = torch.cat(ab_patch_v,dim=0)
x_E = model.forward_patchwise_SR(x,k,ab_patch_v,patch_num,[patch_size[0],patch_size[1]],sf)
loss = F.l1_loss(x_E,x_gt)
optimizer.zero_grad()
loss.backward()
optimizer.step()
scheduler.step()
t_iter = time.time() - t0 - t_data
print('[iter:{}] loss:{:.4f}, data_time:{:.2f}s, net_time:{:.2f}s'.format(global_iter+1,loss.item(),t_data,t_iter))
patch_L = cv2.resize(patch_L,dsize=None,fx=sf,fy=sf,interpolation=cv2.INTER_NEAREST)
#patch_L = patch_L[block_expand*sf:-block_expand*sf,block_expand*sf:-block_expand*sf]
patch_E = util.tensor2uint((x_E))
show = np.hstack((patch_H,patch_L,patch_E))
cv2.imshow('H,L,E',show)
key = cv2.waitKey(1)
global_iter+= 1
if i % 1000 ==0:
cv2.imwrite(os.path.join('./result', 'test' , 'resultE-{:04d}.png'.format(i + 1)), patch_E)
cv2.imwrite(os.path.join('./result', 'test', 'resultL-{:04d}.png'.format(i + 1)), patch_L)
cv2.imwrite(os.path.join('./result', 'test', 'resultH-{:04d}.png'.format(i + 1)), patch_H)
# if key==ord('q'):
# break
# if key==ord('s'):
# ab_numpy = ab.detach().cpu().numpy().flatten()
# np.savetxt('./data/ab.txt',ab_numpy)
ab_numpy = ab.detach().cpu().numpy().flatten()
torch.save(model.state_dict(),'./data/uabcnet_finetune.pth')
np.savetxt('./data/ab_finetune.txt',ab_numpy)
def main():
# ----------------------------------------
# Preparation
# ----------------------------------------
noise_level_img = 0/255.0 # set AWGN noise level for LR image, default: 0,
noise_level_model = noise_level_img # setnoise level of model, default 0
model_name = 'drunet_color' # set denoiser, | 'drunet_color' | 'ircnn_gray' | 'drunet_gray' | 'ircnn_color'
testset_name = 'srbsd68' # set test set, 'set5' | 'srbsd68'
x8 = True # default: False, x8 to boost performance
test_sf = [2] # set scale factor, default: [2, 3, 4], [2], [3], [4]
iter_num = 24 # set number of iterations, default: 24 for SISR
modelSigma1 = 49 # set sigma_1, default: 49
classical_degradation = True # set classical degradation or bicubic degradation
show_img = False # default: False
save_L = True # save LR image
save_E = True # save estimated image
save_LEH = False # save zoomed LR, E and H images
task_current = 'sr' # 'sr' for super-resolution
n_channels = 1 if 'gray' in model_name else 3 # fixed
model_zoo = 'model_zoo' # fixed
testsets = 'testsets' # fixed
results = 'results' # fixed
result_name = testset_name + '_' + task_current + '_' + model_name
model_path = os.path.join(model_zoo, model_name+'.pth')
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
torch.cuda.empty_cache()
# ----------------------------------------
# L_path, E_path, H_path
# ----------------------------------------
L_path = os.path.join(testsets, testset_name) # L_path, for Low-quality images
E_path = os.path.join(results, result_name) # E_path, for Estimated images
util.mkdir(E_path)
logger_name = result_name
utils_logger.logger_info(logger_name, log_path=os.path.join(E_path, logger_name+'.log'))
logger = logging.getLogger(logger_name)
# ----------------------------------------
# load model
# ----------------------------------------
if 'drunet' in model_name:
from models.network_unet import UNetRes as net
model = net(in_nc=n_channels+1, out_nc=n_channels, nc=[64, 128, 256, 512], nb=4, act_mode='R', downsample_mode="strideconv", upsample_mode="convtranspose")
model.load_state_dict(torch.load(model_path), strict=True)
model.eval()
for _, v in model.named_parameters():
v.requires_grad = False
model = model.to(device)
elif 'ircnn' in model_name:
from models.network_dncnn import IRCNN as net
model = net(in_nc=n_channels, out_nc=n_channels, nc=64)
model25 = torch.load(model_path)
former_idx = 0
logger.info('model_name:{}, image sigma:{:.3f}, model sigma:{:.3f}'.format(model_name, noise_level_img, noise_level_model))
logger.info('Model path: {:s}'.format(model_path))
logger.info(L_path)
L_paths = util.get_image_paths(L_path)
# --------------------------------
# load kernel
# --------------------------------
# kernels = hdf5storage.loadmat(os.path.join('kernels', 'Levin09.mat'))['kernels']
if classical_degradation:
kernels = hdf5storage.loadmat(os.path.join('kernels', 'kernels_12.mat'))['kernels']
else:
kernels = hdf5storage.loadmat(os.path.join('kernels', 'kernel_bicubicx234.mat'))['kernels']
test_results_ave = OrderedDict()
test_results_ave['psnr_sf_k'] = []
test_results_ave['psnr_y_sf_k'] = []
for sf in test_sf:
border = sf
modelSigma2 = max(sf, noise_level_model*255.)
k_num = 8 if classical_degradation else 1
for k_index in range(k_num):
logger.info('--------- sf:{:>1d} --k:{:>2d} ---------'.format(sf, k_index))
test_results = OrderedDict()
test_results['psnr'] = []
test_results['psnr_y'] = []
if not classical_degradation: # for bicubic degradation
k_index = sf-2
k = kernels[0, k_index].astype(np.float64)
util.surf(k) if show_img else None
for idx, img in enumerate(L_paths):
# --------------------------------
# (1) get img_L
# --------------------------------
img_name, ext = os.path.splitext(os.path.basename(img))
img_H = util.imread_uint(img, n_channels=n_channels)
img_H = util.modcrop(img_H, sf) # modcrop
if classical_degradation:
img_L = sr.classical_degradation(img_H, k, sf)
util.imshow(img_L) if show_img else None
img_L = util.uint2single(img_L)
else:
img_L = util.imresize_np(util.uint2single(img_H), 1/sf)
np.random.seed(seed=0) # for reproducibility
img_L += np.random.normal(0, noise_level_img, img_L.shape) # add AWGN
# --------------------------------
# (2) get rhos and sigmas
# --------------------------------
rhos, sigmas = pnp.get_rho_sigma(sigma=max(0.255/255., noise_level_model), iter_num=iter_num, modelSigma1=modelSigma1, modelSigma2=modelSigma2, w=1)
rhos, sigmas = torch.tensor(rhos).to(device), torch.tensor(sigmas).to(device)
# --------------------------------
# (3) initialize x, and pre-calculation
# --------------------------------
x = cv2.resize(img_L, (img_L.shape[1]*sf, img_L.shape[0]*sf), interpolation=cv2.INTER_CUBIC)
if np.ndim(x)==2:
x = x[..., None]
if classical_degradation:
x = sr.shift_pixel(x, sf)
x = util.single2tensor4(x).to(device)
img_L_tensor, k_tensor = util.single2tensor4(img_L), util.single2tensor4(np.expand_dims(k, 2))
[k_tensor, img_L_tensor] = util.todevice([k_tensor, img_L_tensor], device)
FB, FBC, F2B, FBFy = sr.pre_calculate(img_L_tensor, k_tensor, sf)
# --------------------------------
# (4) main iterations
# --------------------------------
for i in range(iter_num):
# --------------------------------
# step 1, FFT
# --------------------------------
tau = rhos[i].float().repeat(1, 1, 1, 1)
x = sr.data_solution(x.float(), FB, FBC, F2B, FBFy, tau, sf)
if 'ircnn' in model_name:
current_idx = np.int(np.ceil(sigmas[i].cpu().numpy()*255./2.)-1)
if current_idx != former_idx:
model.load_state_dict(model25[str(current_idx)], strict=True)
model.eval()
for _, v in model.named_parameters():
v.requires_grad = False
model = model.to(device)
former_idx = current_idx
# --------------------------------
# step 2, denoiser
# --------------------------------
if x8:
x = util.augment_img_tensor4(x, i % 8)
if 'drunet' in model_name:
x = torch.cat((x, sigmas[i].float().repeat(1, 1, x.shape[2], x.shape[3])), dim=1)
x = utils_model.test_mode(model, x, mode=2, refield=32, min_size=256, modulo=16)
elif 'ircnn' in model_name:
x = model(x)
if x8:
if i % 8 == 3 or i % 8 == 5:
x = util.augment_img_tensor4(x, 8 - i % 8)
else:
x = util.augment_img_tensor4(x, i % 8)
# --------------------------------
# (3) img_E
# --------------------------------
img_E = util.tensor2uint(x)
if save_E:
util.imsave(img_E, os.path.join(E_path, img_name+'_x'+str(sf)+'_k'+str(k_index)+'_'+model_name+'.png'))
if n_channels == 1:
img_H = img_H.squeeze()
# --------------------------------
# (4) img_LEH
# --------------------------------
img_L = util.single2uint(img_L).squeeze()
if save_LEH:
k_v = k/np.max(k)*1.0
if n_channels==1:
k_v = util.single2uint(k_v)
else:
k_v = util.single2uint(np.tile(k_v[..., np.newaxis], [1, 1, n_channels]))
k_v = cv2.resize(k_v, (3*k_v.shape[1], 3*k_v.shape[0]), interpolation=cv2.INTER_NEAREST)
img_I = cv2.resize(img_L, (sf*img_L.shape[1], sf*img_L.shape[0]), interpolation=cv2.INTER_NEAREST)
img_I[:k_v.shape[0], -k_v.shape[1]:, ...] = k_v
img_I[:img_L.shape[0], :img_L.shape[1], ...] = img_L
util.imshow(np.concatenate([img_I, img_E, img_H], axis=1), title='LR / Recovered / Ground-truth') if show_img else None
util.imsave(np.concatenate([img_I, img_E, img_H], axis=1), os.path.join(E_path, img_name+'_x'+str(sf)+'_k'+str(k_index)+'_LEH.png'))
if save_L:
util.imsave(img_L, os.path.join(E_path, img_name+'_x'+str(sf)+'_k'+str(k_index)+'_LR.png'))
psnr = util.calculate_psnr(img_E, img_H, border=border)
test_results['psnr'].append(psnr)
logger.info('{:->4d}--> {:>10s} -- sf:{:>1d} --k:{:>2d} PSNR: {:.2f}dB'.format(idx+1, img_name+ext, sf, k_index, psnr))
if n_channels == 3:
img_E_y = util.rgb2ycbcr(img_E, only_y=True)
img_H_y = util.rgb2ycbcr(img_H, only_y=True)
psnr_y = util.calculate_psnr(img_E_y, img_H_y, border=border)
test_results['psnr_y'].append(psnr_y)
# --------------------------------
# Average PSNR for all kernels
# --------------------------------
ave_psnr_k = sum(test_results['psnr']) / len(test_results['psnr'])
logger.info('------> Average PSNR(RGB) of ({}) scale factor: ({}), kernel: ({}) sigma: ({:.2f}): {:.2f} dB'.format(testset_name, sf, k_index, noise_level_model, ave_psnr_k))
test_results_ave['psnr_sf_k'].append(ave_psnr_k)
if n_channels == 3: # RGB image
ave_psnr_y_k = sum(test_results['psnr_y']) / len(test_results['psnr_y'])
logger.info('------> Average PSNR(Y) of ({}) scale factor: ({}), kernel: ({}) sigma: ({:.2f}): {:.2f} dB'.format(testset_name, sf, k_index, noise_level_model, ave_psnr_y_k))
test_results_ave['psnr_y_sf_k'].append(ave_psnr_y_k)
# ---------------------------------------
# Average PSNR for all sf and kernels
# ---------------------------------------
ave_psnr_sf_k = sum(test_results_ave['psnr_sf_k']) / len(test_results_ave['psnr_sf_k'])
logger.info('------> Average PSNR of ({}) {:.2f} dB'.format(testset_name, ave_psnr_sf_k))
if n_channels == 3:
ave_psnr_y_sf_k = sum(test_results_ave['psnr_y_sf_k']) / len(test_results_ave['psnr_y_sf_k'])
logger.info('------> Average PSNR of ({}) {:.2f} dB'.format(testset_name, ave_psnr_y_sf_k))
def __getitem__(self, index):
# -------------------
# get H image
# -------------------
H_path = self.paths_H[index]
img_H = util.imread_uint(H_path, self.n_channels)
L_path = H_path
if self.opt['phase'] == 'train':
# ---------------------------
# 1) scale factor, ensure each batch only involves one scale factor
# ---------------------------
if self.count % self.opt['dataloader_batch_size'] == 0:
# sf = random.choice([1,2,3,4])
self.sf = random.choice(self.scales)
# self.count = 0 # optional
self.count += 1
H, W, _ = img_H.shape
# ----------------------------
# randomly crop the patch
# ----------------------------
rnd_h = random.randint(0, max(0, H - self.patch_size))
rnd_w = random.randint(0, max(0, W - self.patch_size))
patch_H = img_H[rnd_h:rnd_h + self.patch_size,
rnd_w:rnd_w + self.patch_size, :]
# ---------------------------
# augmentation - flip, rotate
# ---------------------------
mode = np.random.randint(0, 8)
patch_H = util.augment_img(patch_H, mode=mode)
# ---------------------------
# 2) kernel
# ---------------------------
r_value = random.randint(0, 7)
if r_value > 3:
k = utils_deblur.blurkernel_synthesis(h=25) # motion blur
else:
sf_k = random.choice(self.scales)
k = utils_sisr.gen_kernel(scale_factor=np.array(
[sf_k, sf_k])) # Gaussian blur
mode_k = random.randint(0, 7)
k = util.augment_img(k, mode=mode_k)
# ---------------------------
# 3) noise level
# ---------------------------
if random.randint(0, 8) == 1:
noise_level = 0 / 255.0
else:
noise_level = np.random.randint(0, self.sigma_max) / 255.0
# ---------------------------
# Low-quality image
# ---------------------------
img_L = ndimage.filters.convolve(patch_H,
np.expand_dims(k, axis=2),
mode='wrap')
img_L = img_L[0::self.sf, 0::self.sf, ...]
# add Gaussian noise
img_L = util.uint2single(img_L) + np.random.normal(
0, noise_level, img_L.shape)
img_H = patch_H
else:
k = self.kernels[0, 0].astype(np.float64) # validation kernel
k /= np.sum(k)
noise_level = 0. / 255.0 # validation noise level
img_L = ndimage.filters.convolve(img_H,
np.expand_dims(k, axis=2),
mode='wrap') # blur
img_L = img_L[0::self.sf_validation, 0::self.sf_validation,
...] # downsampling
img_L = util.uint2single(img_L) + np.random.normal(
0, noise_level, img_L.shape)
k = util.single2tensor3(np.expand_dims(np.float32(k), axis=2))
img_H, img_L = util.uint2tensor3(img_H), util.single2tensor3(img_L)
noise_level = torch.FloatTensor([noise_level]).view([1, 1, 1])
return {
'L': img_L,
'H': img_H,
'k': k,
'sigma': noise_level,
'sf': self.sf,
'L_path': L_path,
'H_path': H_path
}
def main():
# ----------------------------------------
# Preparation
# ----------------------------------------
model_name = 'usrnet' # 'usrgan' | 'usrnet' | 'usrgan_tiny' | 'usrnet_tiny'
testset_name = 'set_real' # test set, 'set_real'
test_image = 'chip.png' # 'chip.png', 'comic.png'
#test_image = 'comic.png'
sf = 4 # scale factor, only from {1, 2, 3, 4}
show_img = False # default: False
save_E = True # save estimated image
save_LE = True # save zoomed LR, Estimated images
# ----------------------------------------
# set noise level and kernel
# ----------------------------------------
if 'chip' in test_image:
noise_level_img = 15 # noise level for LR image, 15 for chip
kernel_width_default_x1234 = [0.6, 0.9, 1.7, 2.2] # Gaussian kernel widths for x1, x2, x3, x4
else:
noise_level_img = 2 # noise level for LR image, 0.5~3 for clean images
kernel_width_default_x1234 = [0.4, 0.7, 1.5, 2.0] # default Gaussian kernel widths of clean/sharp images for x1, x2, x3, x4
noise_level_model = noise_level_img/255. # noise level of model
kernel_width = kernel_width_default_x1234[sf-1]
# set your own kernel width
# kernel_width = 2.2
k = utils_deblur.fspecial('gaussian', 25, kernel_width)
k = sr.shift_pixel(k, sf) # shift the kernel
k /= np.sum(k)
util.surf(k) if show_img else None
# scio.savemat('kernel_realapplication.mat', {'kernel':k})
# load approximated bicubic kernels
#kernels = hdf5storage.loadmat(os.path.join('kernels', 'kernel_bicubicx234.mat'))['kernels']
# kernels = loadmat(os.path.join('kernels', 'kernel_bicubicx234.mat'))['kernels']
# kernel = kernels[0, sf-2].astype(np.float64)
kernel = util.single2tensor4(k[..., np.newaxis])
n_channels = 1 if 'gray' in model_name else 3 # 3 for color image, 1 for grayscale image
model_pool = 'model_zoo' # fixed
testsets = 'testsets' # fixed
results = 'results' # fixed
result_name = testset_name + '_' + model_name
model_path = os.path.join(model_pool, model_name+'.pth')
# ----------------------------------------
# L_path, E_path
# ----------------------------------------
L_path = os.path.join(testsets, testset_name) # L_path, fixed, for Low-quality images
E_path = os.path.join(results, result_name) # E_path, fixed, for Estimated images
util.mkdir(E_path)
logger_name = result_name
utils_logger.logger_info(logger_name, log_path=os.path.join(E_path, logger_name+'.log'))
logger = logging.getLogger(logger_name)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# ----------------------------------------
# load model
# ----------------------------------------
if 'tiny' in model_name:
model = net(n_iter=6, h_nc=32, in_nc=4, out_nc=3, nc=[16, 32, 64, 64],
nb=2, act_mode="R", downsample_mode='strideconv', upsample_mode="convtranspose")
else:
model = net(n_iter=8, h_nc=64, in_nc=4, out_nc=3, nc=[64, 128, 256, 512],
nb=2, act_mode="R", downsample_mode='strideconv', upsample_mode="convtranspose")
model.load_state_dict(torch.load(model_path), strict=True)
model.eval()
for key, v in model.named_parameters():
v.requires_grad = False
number_parameters = sum(map(lambda x: x.numel(), model.parameters()))
logger.info('Params number: {}'.format(number_parameters))
model = model.to(device)
logger.info('Model path: {:s}'.format(model_path))
logger.info('model_name:{}, image sigma:{}'.format(model_name, noise_level_img))
logger.info(L_path)
img = os.path.join(L_path, test_image)
# ------------------------------------
# (1) img_L
# ------------------------------------
img_name, ext = os.path.splitext(os.path.basename(img))
img_L = util.imread_uint(img, n_channels=n_channels)
img_L = util.uint2single(img_L)
util.imshow(img_L) if show_img else None
w, h = img_L.shape[:2]
logger.info('{:>10s}--> ({:>4d}x{:<4d})'.format(img_name+ext, w, h))
# boundary handling
boarder = 8 # default setting for kernel size 25x25
img = cv2.resize(img_L, (sf*h, sf*w), interpolation=cv2.INTER_NEAREST)
img = utils_deblur.wrap_boundary_liu(img, [int(np.ceil(sf*w/boarder+2)*boarder), int(np.ceil(sf*h/boarder+2)*boarder)])
img_wrap = sr.downsample_np(img, sf, center=False)
img_wrap[:w, :h, :] = img_L
img_L = img_wrap
util.imshow(util.single2uint(img_L), title='LR image with noise level {}'.format(noise_level_img)) if show_img else None
img_L = util.single2tensor4(img_L)
img_L = img_L.to(device)
# ------------------------------------
# (2) img_E
# ------------------------------------
sigma = torch.tensor(noise_level_model).float().view([1, 1, 1, 1])
[img_L, kernel, sigma] = [el.to(device) for el in [img_L, kernel, sigma]]
img_E = model(img_L, kernel, sf, sigma)
img_E = util.tensor2uint(img_E)[:sf*w, :sf*h, ...]
if save_E:
util.imsave(img_E, os.path.join(E_path, img_name+'_x'+str(sf)+'_'+model_name+'.png'))
# --------------------------------
# (3) save img_LE
# --------------------------------
if save_LE:
k_v = k/np.max(k)*1.2
k_v = util.single2uint(np.tile(k_v[..., np.newaxis], [1, 1, 3]))
k_factor = 3
k_v = cv2.resize(k_v, (k_factor*k_v.shape[1], k_factor*k_v.shape[0]), interpolation=cv2.INTER_NEAREST)
img_L = util.tensor2uint(img_L)[:w, :h, ...]
img_I = cv2.resize(img_L, (sf*img_L.shape[1], sf*img_L.shape[0]), interpolation=cv2.INTER_NEAREST)
img_I[:k_v.shape[0], :k_v.shape[1], :] = k_v
util.imshow(np.concatenate([img_I, img_E], axis=1), title='LR / Recovered') if show_img else None
util.imsave(np.concatenate([img_I, img_E], axis=1), os.path.join(E_path, img_name+'_x'+str(sf)+'_'+model_name+'_LE.png'))
def __getitem__(self, index):
# ------------------------------------
# get H image
# ------------------------------------
H_path = self.paths_H[index]
img_H = util.imread_uint(H_path, self.n_channels)
L_path = H_path
if self.opt['phase'] == 'train':
"""
# --------------------------------
# get L/H patch pairs
# --------------------------------
"""
H, W, _ = img_H.shape
# --------------------------------
# randomly crop the patch
# --------------------------------
rnd_h = random.randint(0, max(0, H - self.patch_size))
rnd_w = random.randint(0, max(0, W - self.patch_size))
patch_H = img_H[rnd_h:rnd_h + self.patch_size,
rnd_w:rnd_w + self.patch_size, :]
# --------------------------------
# augmentation - flip, rotate
# --------------------------------
mode = np.random.randint(0, 8)
patch_H = util.augment_img(patch_H, mode=mode)
# --------------------------------
# HWC to CHW, numpy(uint) to tensor
# --------------------------------
img_H = util.uint2tensor3(patch_H)
img_L = img_H.clone()
# --------------------------------
# add noise
# --------------------------------
noise = torch.randn(img_L.size()).mul_(self.sigma / 255.0)
img_L.add_(noise)
else:
"""
# --------------------------------
# get L/H image pairs
# --------------------------------
"""
img_H = util.uint2single(img_H)
img_L = np.copy(img_H)
# --------------------------------
# add noise
# --------------------------------
np.random.seed(seed=0)
img_L += np.random.normal(0, self.sigma_test / 255.0, img_L.shape)
# --------------------------------
# HWC to CHW, numpy to tensor
# --------------------------------
img_L = util.single2tensor3(img_L)
img_H = util.single2tensor3(img_H)
return {'L': img_L, 'H': img_H, 'H_path': H_path, 'L_path': L_path}
def main():
# ----------------------------------------
# Preparation
# ----------------------------------------
model_name = 'usrnet' # 'usrgan' | 'usrnet' | 'usrgan_tiny' | 'usrnet_tiny'
testset_name = 'set5' # test set, 'set5' | 'srbsd68'
need_degradation = True # default: True
sf = 4 # scale factor, only from {2, 3, 4}
show_img = False # default: False
save_L = True # save LR image
save_E = True # save estimated image
# load approximated bicubic kernels
#kernels = hdf5storage.loadmat(os.path.join('kernels', 'kernels_bicubicx234.mat'))['kernels']
kernels = loadmat(os.path.join('kernels',
'kernels_bicubicx234.mat'))['kernels']
kernel = kernels[0, sf - 2].astype(np.float64)
kernel = util.single2tensor4(kernel[..., np.newaxis])
task_current = 'sr' # fixed, 'sr' for super-resolution
n_channels = 3 # fixed, 3 for color image
model_pool = 'model_zoo' # fixed
testsets = 'testsets' # fixed
results = 'results' # fixed
noise_level_img = 0 # fixed: 0, noise level for LR image
noise_level_model = noise_level_img # fixed, noise level of model, default 0
result_name = testset_name + '_' + model_name + '_bicubic'
border = sf if task_current == 'sr' else 0 # shave boader to calculate PSNR and SSIM
model_path = os.path.join(model_pool, model_name + '.pth')
# ----------------------------------------
# L_path, E_path, H_path
# ----------------------------------------
L_path = os.path.join(
testsets, testset_name) # L_path, fixed, for Low-quality images
H_path = L_path # H_path, 'None' | L_path, for High-quality images
E_path = os.path.join(results,
result_name) # E_path, fixed, for Estimated images
util.mkdir(E_path)
if H_path == L_path:
need_degradation = True
logger_name = result_name
utils_logger.logger_info(logger_name,
log_path=os.path.join(E_path,
logger_name + '.log'))
logger = logging.getLogger(logger_name)
need_H = True if H_path is not None else False
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# ----------------------------------------
# load model
# ----------------------------------------
from models.network_usrnet import USRNet as net # for pytorch version <= 1.7.1
# from models.network_usrnet_v1 import USRNet as net # for pytorch version >=1.8.1
if 'tiny' in model_name:
model = net(n_iter=6,
h_nc=32,
in_nc=4,
out_nc=3,
nc=[16, 32, 64, 64],
nb=2,
act_mode="R",
downsample_mode='strideconv',
upsample_mode="convtranspose")
else:
model = net(n_iter=8,
h_nc=64,
in_nc=4,
out_nc=3,
nc=[64, 128, 256, 512],
nb=2,
act_mode="R",
downsample_mode='strideconv',
upsample_mode="convtranspose")
model.load_state_dict(torch.load(model_path), strict=True)
model.eval()
for key, v in model.named_parameters():
v.requires_grad = False
number_parameters = sum(map(lambda x: x.numel(), model.parameters()))
logger.info('Params number: {}'.format(number_parameters))
model = model.to(device)
logger.info('Model path: {:s}'.format(model_path))
test_results = OrderedDict()
test_results['psnr'] = []
test_results['ssim'] = []
test_results['psnr_y'] = []
test_results['ssim_y'] = []
logger.info('model_name:{}, image sigma:{}'.format(model_name,
noise_level_img))
logger.info(L_path)
L_paths = util.get_image_paths(L_path)
H_paths = util.get_image_paths(H_path) if need_H else None
for idx, img in enumerate(L_paths):
# ------------------------------------
# (1) img_L
# ------------------------------------
img_name, ext = os.path.splitext(os.path.basename(img))
logger.info('{:->4d}--> {:>10s}'.format(idx + 1, img_name + ext))
img_L = util.imread_uint(img, n_channels=n_channels)
img_L = util.uint2single(img_L)
# degradation process, bicubic downsampling
if need_degradation:
img_L = util.modcrop(img_L, sf)
img_L = util.imresize_np(img_L, 1 / sf)
# img_L = util.uint2single(util.single2uint(img_L))
# np.random.seed(seed=0) # for reproducibility
# img_L += np.random.normal(0, noise_level_img/255., img_L.shape)
w, h = img_L.shape[:2]
if save_L:
util.imsave(
util.single2uint(img_L),
os.path.join(E_path, img_name + '_LR_x' + str(sf) + '.png'))
img = cv2.resize(img_L, (sf * h, sf * w),
interpolation=cv2.INTER_NEAREST)
img = utils_deblur.wrap_boundary_liu(img, [
int(np.ceil(sf * w / 8 + 2) * 8),
int(np.ceil(sf * h / 8 + 2) * 8)
])
img_wrap = sr.downsample_np(img, sf, center=False)
img_wrap[:w, :h, :] = img_L
img_L = img_wrap
util.imshow(util.single2uint(img_L),
title='LR image with noise level {}'.format(
noise_level_img)) if show_img else None
img_L = util.single2tensor4(img_L)
img_L = img_L.to(device)
# ------------------------------------
# (2) img_E
# ------------------------------------
sigma = torch.tensor(noise_level_model).float().view([1, 1, 1, 1])
[img_L, kernel,
sigma] = [el.to(device) for el in [img_L, kernel, sigma]]
img_E = model(img_L, kernel, sf, sigma)
img_E = util.tensor2uint(img_E)
img_E = img_E[:sf * w, :sf * h, :]
if need_H:
# --------------------------------
# (3) img_H
# --------------------------------
img_H = util.imread_uint(H_paths[idx], n_channels=n_channels)
img_H = img_H.squeeze()
img_H = util.modcrop(img_H, sf)
# --------------------------------
# PSNR and SSIM
# --------------------------------
psnr = util.calculate_psnr(img_E, img_H, border=border)
ssim = util.calculate_ssim(img_E, img_H, border=border)
test_results['psnr'].append(psnr)
test_results['ssim'].append(ssim)
logger.info('{:s} - PSNR: {:.2f} dB; SSIM: {:.4f}.'.format(
img_name + ext, psnr, ssim))
util.imshow(np.concatenate([img_E, img_H], axis=1),
title='Recovered / Ground-truth') if show_img else None
if np.ndim(img_H) == 3: # RGB image
img_E_y = util.rgb2ycbcr(img_E, only_y=True)
img_H_y = util.rgb2ycbcr(img_H, only_y=True)
psnr_y = util.calculate_psnr(img_E_y, img_H_y, border=border)
ssim_y = util.calculate_ssim(img_E_y, img_H_y, border=border)
test_results['psnr_y'].append(psnr_y)
test_results['ssim_y'].append(ssim_y)
# ------------------------------------
# save results
# ------------------------------------
if save_E:
util.imsave(
img_E,
os.path.join(
E_path,
img_name + '_x' + str(sf) + '_' + model_name + '.png'))
if need_H:
ave_psnr = sum(test_results['psnr']) / len(test_results['psnr'])
ave_ssim = sum(test_results['ssim']) / len(test_results['ssim'])
logger.info(
'Average PSNR/SSIM(RGB) - {} - x{} --PSNR: {:.2f} dB; SSIM: {:.4f}'
.format(result_name, sf, ave_psnr, ave_ssim))
if np.ndim(img_H) == 3:
ave_psnr_y = sum(test_results['psnr_y']) / len(
test_results['psnr_y'])
ave_ssim_y = sum(test_results['ssim_y']) / len(
test_results['ssim_y'])
logger.info(
'Average PSNR/SSIM( Y ) - {} - x{} - PSNR: {:.2f} dB; SSIM: {:.4f}'
.format(result_name, sf, ave_psnr_y, ave_ssim_y))
epochs = 50 model_path = os.path.join(model_pool, model_name+'.pth') test_path = os.path.join(test_im) train_path = os.path.join(train_im) # load model model = DnCNN(in_nc=n_channels, out_nc=n_channels, nc=64, nb=17, act_mode='R') model.load_state_dict(torch.load(model_path), strict=True) model = model.to(device) model.eval() # load test image x = util.imread_uint(test_path, n_channels=n_channels) orig_im = x.squeeze() x = util.uint2single(x) np.random.seed(seed=0) # for reproducibility y = x + np.random.normal(0, sigma/255., x.shape) # add gaussian noise y = util.single2tensor4(y) y = y.to(device) # denoise the image to compare PSNR before and after adaptation with torch.no_grad(): x_ = model(y) # compute PSNR denoised_im = util.tensor2uint(x_) prev_psnr = util.calculate_psnr(denoised_im, orig_im, border=0) # external adaptation
def main():
# ----------------------------------------
# Preparation
# ----------------------------------------
model_name = 'dncnn3' # 'dncnn3'- can be used for blind Gaussian denoising, JPEG deblocking (quality factor 5-100) and super-resolution (x234)
# important!
testset_name = 'bsd68' # test set, low-quality grayscale/color JPEG images
n_channels = 1 # set 1 for grayscale image, set 3 for color image
x8 = False # default: False, x8 to boost performance
testsets = 'testsets' # fixed
results = 'results' # fixed
result_name = testset_name + '_' + model_name # fixed
L_path = os.path.join(
testsets, testset_name
) # L_path, for Low-quality grayscale/Y-channel JPEG images
E_path = os.path.join(results, result_name) # E_path, for Estimated images
util.mkdir(E_path)
model_pool = 'model_zoo' # fixed
model_path = os.path.join(model_pool, model_name + '.pth')
logger_name = result_name
utils_logger.logger_info(logger_name,
log_path=os.path.join(E_path,
logger_name + '.log'))
logger = logging.getLogger(logger_name)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# ----------------------------------------
# load model
# ----------------------------------------
from models.network_dncnn import DnCNN as net
model = net(in_nc=1, out_nc=1, nc=64, nb=20, act_mode='R')
model.load_state_dict(torch.load(model_path), strict=True)
model.eval()
for k, v in model.named_parameters():
v.requires_grad = False
model = model.to(device)
logger.info('Model path: {:s}'.format(model_path))
number_parameters = sum(map(lambda x: x.numel(), model.parameters()))
logger.info('Params number: {}'.format(number_parameters))
logger.info(L_path)
L_paths = util.get_image_paths(L_path)
for idx, img in enumerate(L_paths):
# ------------------------------------
# (1) img_L
# ------------------------------------
img_name, ext = os.path.splitext(os.path.basename(img))
logger.info('{:->4d}--> {:>10s}'.format(idx + 1, img_name + ext))
img_L = util.imread_uint(img, n_channels=n_channels)
img_L = util.uint2single(img_L)
if n_channels == 3:
ycbcr = util.rgb2ycbcr(img_L, False)
img_L = ycbcr[..., 0:1]
img_L = util.single2tensor4(img_L)
img_L = img_L.to(device)
# ------------------------------------
# (2) img_E
# ------------------------------------
if not x8:
img_E = model(img_L)
else:
img_E = utils_model.test_mode(model, img_L, mode=3)
img_E = util.tensor2single(img_E)
if n_channels == 3:
ycbcr[..., 0] = img_E
img_E = util.ycbcr2rgb(ycbcr)
img_E = util.single2uint(img_E)
# ------------------------------------
# save results
# ------------------------------------
util.imsave(img_E, os.path.join(E_path, img_name + '.png'))
def main():
# ----------------------------------------
# load kernels
# ----------------------------------------
#PSF_grid = np.load('./data/AC254-075-A-ML-Zemax(ZMX).npz')['PSF']
PSF_grid = np.load('./data/Heide_PSF_plano_small.npz')['PSF']
PSF_grid = PSF_grid.astype(np.float32)
gx, gy = PSF_grid.shape[:2]
for xx in range(gx):
for yy in range(gy):
PSF_grid[xx, yy] = PSF_grid[xx, yy] / np.sum(PSF_grid[xx, yy],
axis=(0, 1))
# ----------------------------------------
# load model
# ----------------------------------------
stage = 8
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = net(n_iter=stage,
h_nc=64,
in_nc=4,
out_nc=3,
nc=[64, 128, 256, 512],
nb=2,
act_mode="R",
downsample_mode='strideconv',
upsample_mode="convtranspose")
model_code = 'iter800'
loaded_state = torch.load(
'/home/xiu/databag/deblur/models/plano/uabcnet_{}.pth'.format(
model_code))
#strip_state = strip_prefix_if_present(loaded_state,prefix="p.")
model.load_state_dict(loaded_state, strict=True)
model.eval()
for _, v in model.named_parameters():
v.requires_grad = False
model = model.to(device)
for img_id in range(1, 237):
#for img_id in range(1,12):
#img_L = cv2.imread('/home/xiu/workspace/UABC/ICCV2021/video1-3/res/2_{:03d}.bmp'.format(img_id))
#img_L = cv2.imread('/home/xiu/workspace/UABC/ICCV2021/video/{:08d}.bmp'.format(img_id))
#img_L = cv2.imread('/home/xiu/databag/deblur/ICCV2021/suo_image/{}/AC254-075-A-ML-Zemax(ZMX).bmp'.format(img_id))
#img_L = cv2.imread('/home/xiu/workspace/UABC/ICCV2021/ResolutionChart/Reso.bmp')
img_L = cv2.imread(
'/home/xiu/databag/deblur/ICCV2021/MPI_data/drain/blurry.jpg')
img_L = img_L.astype(np.float32)
img_L = img_L[38:-39, 74:-74]
img_L = cv2.resize(img_L, dsize=None, fx=0.5, fy=0.5)
#img_L = np.pad(img_L,((1,1),(61,62),(0,0)),mode='edge')
W, H = img_L.shape[:2]
print(gx, gy)
num_patch = [gx, gy]
#positional alpha-beta parameters for HQS
ab_numpy = np.loadtxt(
'/home/xiu/databag/deblur/models/plano/ab_{}.txt'.format(
model_code)).astype(np.float32).reshape(gx, gy, stage * 2, 3)
ab = torch.tensor(ab_numpy, device=device, requires_grad=False)
t0 = time.time()
px_start = 0
py_start = 0
PSF_patch = PSF_grid[px_start:px_start + num_patch[0],
py_start:py_start + num_patch[1]]
#block_expand = 1
patch_L = img_L[px_start * W // gx:(px_start + num_patch[0]) * W // gx,
py_start * H // gy:(py_start + num_patch[1]) * H //
gy, :]
p_W, p_H = patch_L.shape[:2]
expand = max(PSF_grid.shape[2] // 2, p_W // 16)
block_expand = expand
patch_L_wrap = util_deblur.wrap_boundary_liu(
patch_L, (p_W + block_expand * 2, p_H + block_expand * 2))
#centralize
patch_L_wrap = np.hstack((patch_L_wrap[:, -block_expand:, :],
patch_L_wrap[:, :p_H + block_expand, :]))
patch_L_wrap = np.vstack((patch_L_wrap[-block_expand:, :, :],
patch_L_wrap[:p_W + block_expand, :, :]))
x = util.uint2single(patch_L_wrap)
x = util.single2tensor4(x)
k_all = []
for h_ in range(num_patch[1]):
for w_ in range(num_patch[0]):
k_all.append(util.single2tensor4(PSF_patch[w_, h_]))
k = torch.cat(k_all, dim=0)
[x, k] = [el.to(device) for el in [x, k]]
ab_patch = F.softplus(ab[px_start:px_start + num_patch[0],
py_start:py_start + num_patch[1]])
cd = []
for h_ in range(num_patch[1]):
for w_ in range(num_patch[0]):
cd.append(ab_patch[w_:w_ + 1, h_])
cd = torch.cat(cd, dim=0)
x_E = model.forward_patchwise(x, k, cd, num_patch, [W // gx, H // gy])
x_E = x_E[..., block_expand:block_expand + p_W,
block_expand:block_expand + p_H]
patch_L = patch_L_wrap.astype(np.uint8)
patch_E = util.tensor2uint(x_E)
#patch_E_z = np.hstack((patch_E_all[::2]))
#patch_E_x = np.hstack((patch_E_all[1::2]))
#patch_E_show = np.vstack((patch_E_z,patch_E_x))
#if block_expand>0:
# show = np.hstack((patch_L[block_expand:-block_expand,block_expand:-block_expand],patch_E))
#else:
# show = np.hstack((patch_L,patch_E))
#cv2.imshow('stage',patch_E_show)
#cv2.imshow('HL',show)
#cv2.imshow('RGB',rgb)
#key = cv2.waitKey(-1)
#if key==ord('n'):
# break
t1 = time.time()
print(t1 - t0)
# print(i)
xk = patch_E
# #zk = zk.astype(np.uint8)
xk = xk.astype(np.uint8)
#cv2.imwrite('/home/xiu/workspace/UABC/ICCV2021/new_image/image/ours-{}.png'.format(img_id),xk)
#cv2.imwrite('/home/xiu/workspace/UABC/ICCV2021/video_deblur/{:08d}.png'.format(img_id),xk)
#cv2.imwrite('/home/xiu/workspace/UABC/ICCV2021/cap_result/1_{:03d}.png'.format(img_id),xk)
cv2.imshow('xx', xk)
cv2.imshow('img_L', patch_L.astype(np.uint8))
key = cv2.waitKey(-1)
if key == ord('q'):
break
def main():
# ----------------------------------------
# Preparation
# ----------------------------------------
noise_level_img = 15 # set AWGN noise level for noisy image
noise_level_model = noise_level_img # set noise level for model
model_name = 'drunet_gray' # set denoiser model, 'drunet_gray' | 'drunet_color'
testset_name = 'bsd68' # set test set, 'bsd68' | 'cbsd68' | 'set12'
x8 = False # default: False, x8 to boost performance
show_img = False # default: False
border = 0 # shave boader to calculate PSNR and SSIM
if 'color' in model_name:
n_channels = 3 # 3 for color image
else:
n_channels = 1 # 1 for grayscale image
model_pool = 'model_zoo' # fixed
testsets = 'testsets' # fixed
results = 'results' # fixed
task_current = 'dn' # 'dn' for denoising
result_name = testset_name + '_' + task_current + '_' + model_name
model_path = os.path.join(model_pool, model_name + '.pth')
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
torch.cuda.empty_cache()
# ----------------------------------------
# L_path, E_path, H_path
# ----------------------------------------
L_path = os.path.join(testsets,
testset_name) # L_path, for Low-quality images
E_path = os.path.join(results, result_name) # E_path, for Estimated images
util.mkdir(E_path)
logger_name = result_name
utils_logger.logger_info(logger_name,
log_path=os.path.join(E_path,
logger_name + '.log'))
logger = logging.getLogger(logger_name)
# ----------------------------------------
# load model
# ----------------------------------------
from models.network_unet import UNetRes as net
model = net(in_nc=n_channels + 1,
out_nc=n_channels,
nc=[64, 128, 256, 512],
nb=4,
act_mode='R',
downsample_mode="strideconv",
upsample_mode="convtranspose")
model.load_state_dict(torch.load(model_path), strict=True)
model.eval()
for k, v in model.named_parameters():
v.requires_grad = False
model = model.to(device)
logger.info('Model path: {:s}'.format(model_path))
number_parameters = sum(map(lambda x: x.numel(), model.parameters()))
logger.info('Params number: {}'.format(number_parameters))
test_results = OrderedDict()
test_results['psnr'] = []
test_results['ssim'] = []
logger.info('model_name:{}, model sigma:{}, image sigma:{}'.format(
model_name, noise_level_img, noise_level_model))
logger.info(L_path)
L_paths = util.get_image_paths(L_path)
for idx, img in enumerate(L_paths):
# ------------------------------------
# (1) img_L
# ------------------------------------
img_name, ext = os.path.splitext(os.path.basename(img))
# logger.info('{:->4d}--> {:>10s}'.format(idx+1, img_name+ext))
img_H = util.imread_uint(img, n_channels=n_channels)
img_L = util.uint2single(img_H)
# Add noise without clipping
np.random.seed(seed=0) # for reproducibility
img_L += np.random.normal(0, noise_level_img / 255., img_L.shape)
util.imshow(util.single2uint(img_L),
title='Noisy image with noise level {}'.format(
noise_level_img)) if show_img else None
img_L = util.single2tensor4(img_L)
img_L = torch.cat(
(img_L, torch.FloatTensor([noise_level_model / 255.]).repeat(
1, 1, img_L.shape[2], img_L.shape[3])),
dim=1)
img_L = img_L.to(device)
# ------------------------------------
# (2) img_E
# ------------------------------------
if not x8 and img_L.size(2) // 8 == 0 and img_L.size(3) // 8 == 0:
img_E = model(img_L)
elif not x8 and (img_L.size(2) // 8 != 0 or img_L.size(3) // 8 != 0):
img_E = utils_model.test_mode(model, img_L, refield=64, mode=5)
elif x8:
img_E = utils_model.test_mode(model, img_L, mode=3)
img_E = util.tensor2uint(img_E)
# --------------------------------
# PSNR and SSIM
# --------------------------------
if n_channels == 1:
img_H = img_H.squeeze()
psnr = util.calculate_psnr(img_E, img_H, border=border)
ssim = util.calculate_ssim(img_E, img_H, border=border)
test_results['psnr'].append(psnr)
test_results['ssim'].append(ssim)
logger.info('{:s} - PSNR: {:.2f} dB; SSIM: {:.4f}.'.format(
img_name + ext, psnr, ssim))
# ------------------------------------
# save results
# ------------------------------------
util.imsave(img_E, os.path.join(E_path, img_name + ext))
ave_psnr = sum(test_results['psnr']) / len(test_results['psnr'])
ave_ssim = sum(test_results['ssim']) / len(test_results['ssim'])
logger.info(
'Average PSNR/SSIM(RGB) - {} - PSNR: {:.2f} dB; SSIM: {:.4f}'.format(
result_name, ave_psnr, ave_ssim))
def main():
# ----------------------------------------
# load kernels
# ----------------------------------------
#PSF_grid = np.load('./data/Schuler_PSF01.npz')['PSF']
#PSF_grid = np.load('./data/Schuler_PSF_facade.npz')['PSF']
PSF_grid = np.load('./data/ZEMAX-AC254-075-A-new.npz')['PSF']
#PSF_grid = np.load('./data/Schuler_PSF03.npz')['PSF']
#PSF_grid = np.load('./data/PSF.npz')['PSF']
#print(PSF_grid.shape)
PSF_grid = PSF_grid.astype(np.float32)
gx,gy = PSF_grid.shape[:2]
for xx in range(gx):
for yy in range(gy):
PSF_grid[xx,yy] = PSF_grid[xx,yy]/np.sum(PSF_grid[xx,yy],axis=(0,1))
#PSF_grid = PSF_grid[:,1:-1,...]
# ----------------------------------------
# load model
# ----------------------------------------
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = net(n_iter=8, h_nc=64, in_nc=4, out_nc=3, nc=[64, 128, 256, 512],
nb=2, act_mode="R", downsample_mode='strideconv', upsample_mode="convtranspose")
loaded_state = torch.load('./usrnet_ZEMAX.pth')
#strip_state = strip_prefix_if_present(loaded_state,prefix="p.")
model.load_state_dict(loaded_state, strict=True)
model.eval()
#model.train()
for _, v in model.named_parameters():
v.requires_grad = False
# v.requires_grad = False
model = model.to(device)
for img_id in range(100):
img_L = cv2.imread('/home/xiu/workspace/UABC/ICCV2021/video/{:08d}.bmp'.format(img_id))
img_L = img_L.astype(np.float32)
img_E = np.zeros_like(img_L)
weight_E = np.zeros_like(img_L)
patch_size = 2*128
num_patch = 2
p_size = patch_size//num_patch
expand = PSF_grid.shape[2]
ab_numpy = np.loadtxt('ab_ZEMAX.txt').astype(np.float32).reshape(6,8,16,3)
ab = torch.tensor(ab_numpy,device=device,requires_grad=False)
#save img_L
t0 = time.time()
#while running:
for px_start in range(0,6-2+1,2):
for py_start in range(0,8-2+1,2):
PSF_patch = PSF_grid[px_start:px_start+num_patch,py_start:py_start+num_patch]
patch_L = img_L[px_start*p_size:(px_start+num_patch)*p_size,py_start*p_size:py_start*p_size+num_patch*p_size,:]
block_expand = expand
#block_expand = 1
if block_expand > 0:
patch_L_wrap = util_deblur.wrap_boundary_liu(patch_L,(patch_size+block_expand*2,patch_size+block_expand*2))
#centralize
patch_L_wrap = np.hstack((patch_L_wrap[:,-block_expand:,:],patch_L_wrap[:,:patch_size+block_expand,:]))
patch_L_wrap = np.vstack((patch_L_wrap[-block_expand:,:,:],patch_L_wrap[:patch_size+block_expand,:,:]))
else:
patch_L_wrap = patch_L
if block_expand>0:
x = util.uint2single(patch_L_wrap)
else:
x = util.uint2single(patch_L)
x = util.single2tensor4(x)
# x_blocky = torch.cat(torch.chunk(x,num_patch,dim=2),dim=0)
# x_blocky = torch.cat(torch.chunk(x_blocky,num_patch,dim=3),dim=0)
k_all = []
for w_ in range(num_patch):
for h_ in range(num_patch):
k_all.append(util.single2tensor4(PSF_patch[h_,w_]))
k = torch.cat(k_all,dim=0)
[x,k] = [el.to(device) for el in [x,k]]
cd = F.softplus(ab[px_start:px_start+num_patch,py_start:py_start+num_patch])
cd = cd.view(num_patch**2,2*8,3)
x_E = model.forward_patchwise(x,k,cd,[num_patch,num_patch],[patch_size//num_patch,patch_size//num_patch])
patch_L = patch_L_wrap.astype(np.uint8)
patch_E = util.tensor2uint(x_E)
patch_E_all = [util.tensor2uint(pp) for pp in x_E]
#patch_E_z = np.hstack((patch_E_all[::2]))
#patch_E_x = np.hstack((patch_E_all[1::2]))
#patch_E_show = np.vstack((patch_E_z,patch_E_x))
#if block_expand>0:
# show = np.hstack((patch_L[block_expand:-block_expand,block_expand:-block_expand],patch_E))
#else:
# show = np.hstack((patch_L,patch_E))
#get kernel
for i in range(8):
img_E_deconv[i][px_start*p_size:(px_start+num_patch)*p_size,py_start*p_size:py_start*p_size+num_patch*p_size,:] += patch_E_all[-2][expand:-expand,expand:-expand]
img_E_denoise[i][px_start*p_size:(px_start+num_patch)*p_size,py_start*p_size:py_start*p_size+num_patch*p_size,:] += patch_E_all[-1][expand:-expand,expand:-expand]
weight_E[px_start*p_size:(px_start+num_patch)*p_size,py_start*p_size:py_start*p_size+num_patch*p_size,:] += 1.0
#cv2.imshow('stage',patch_E_show)
#cv2.imshow('HL',show)
#cv2.imshow('RGB',rgb)
#key = cv2.waitKey(-1)
#if key==ord('n'):
# break
t1 = time.time()
print(t1-t0)
img_E = img_E/weight_E
img_E_deconv = [pp/weight_E for pp in img_E_deconv]
img_E_denoise = [pp/weight_E for pp in img_E_denoise]
# print(i)
xk = img_E_denoise[-1]
# #zk = zk.astype(np.uint8)
xk = xk.astype(np.uint8)
#cv2.imwrite('/home/xiu/workspace/UABC/ICCV2021/video_deblur/{:08d}.png'.format(img_id),xk)
cv2.imshow('xx',xk)
cv2.imshow('img_L',img_L.astype(np.uint8))
cv2.waitKey(-1)
def main():
# ----------------------------------------
# Preparation
# ----------------------------------------
noise_level_img = 7.65 / 255.0 # default: 0, noise level for LR image
noise_level_model = noise_level_img # noise level of model, default 0
model_name = 'drunet_gray' # 'drunet_gray' | 'drunet_color' | 'ircnn_gray' | 'ircnn_color'
testset_name = 'Set3C' # test set, 'set5' | 'srbsd68'
x8 = True # default: False, x8 to boost performance
iter_num = 8 # number of iterations
modelSigma1 = 49
modelSigma2 = noise_level_model * 255.
show_img = False # default: False
save_L = True # save LR image
save_E = True # save estimated image
save_LEH = False # save zoomed LR, E and H images
border = 0
# --------------------------------
# load kernel
# --------------------------------
kernels = hdf5storage.loadmat(os.path.join('kernels',
'Levin09.mat'))['kernels']
sf = 1
task_current = 'deblur' # 'deblur' for deblurring
n_channels = 3 if 'color' in model_name else 1 # fixed
model_zoo = 'model_zoo' # fixed
testsets = 'testsets' # fixed
results = 'results' # fixed
result_name = testset_name + '_' + task_current + '_' + model_name
model_path = os.path.join(model_zoo, model_name + '.pth')
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
torch.cuda.empty_cache()
# ----------------------------------------
# L_path, E_path, H_path
# ----------------------------------------
L_path = os.path.join(testsets,
testset_name) # L_path, for Low-quality images
E_path = os.path.join(results, result_name) # E_path, for Estimated images
util.mkdir(E_path)
logger_name = result_name
utils_logger.logger_info(logger_name,
log_path=os.path.join(E_path,
logger_name + '.log'))
logger = logging.getLogger(logger_name)
# ----------------------------------------
# load model
# ----------------------------------------
if 'drunet' in model_name:
from models.network_unet import UNetRes as net
model = net(in_nc=n_channels + 1,
out_nc=n_channels,
nc=[64, 128, 256, 512],
nb=4,
act_mode='R',
downsample_mode="strideconv",
upsample_mode="convtranspose")
model.load_state_dict(torch.load(model_path), strict=True)
model.eval()
for _, v in model.named_parameters():
v.requires_grad = False
model = model.to(device)
elif 'ircnn' in model_name:
from models.network_dncnn import IRCNN as net
model = net(in_nc=n_channels, out_nc=n_channels, nc=64)
model25 = torch.load(model_path)
former_idx = 0
logger.info('model_name:{}, image sigma:{:.3f}, model sigma:{:.3f}'.format(
model_name, noise_level_img, noise_level_model))
logger.info('Model path: {:s}'.format(model_path))
logger.info(L_path)
L_paths = util.get_image_paths(L_path)
test_results_ave = OrderedDict()
test_results_ave['psnr'] = [] # record average PSNR for each kernel
for k_index in range(kernels.shape[1]):
logger.info('-------k:{:>2d} ---------'.format(k_index))
test_results = OrderedDict()
test_results['psnr'] = []
k = kernels[0, k_index].astype(np.float64)
util.imshow(k) if show_img else None
for idx, img in enumerate(L_paths):
# --------------------------------
# (1) get img_L
# --------------------------------
img_name, ext = os.path.splitext(os.path.basename(img))
img_H = util.imread_uint(img, n_channels=n_channels)
img_H = util.modcrop(img_H, 8) # modcrop
img_L = ndimage.filters.convolve(img_H,
np.expand_dims(k, axis=2),
mode='wrap')
util.imshow(img_L) if show_img else None
img_L = util.uint2single(img_L)
np.random.seed(seed=0) # for reproducibility
img_L += np.random.normal(0, noise_level_img,
img_L.shape) # add AWGN
# --------------------------------
# (2) get rhos and sigmas
# --------------------------------
rhos, sigmas = pnp.get_rho_sigma(sigma=max(0.255 / 255.,
noise_level_model),
iter_num=iter_num,
modelSigma1=modelSigma1,
modelSigma2=modelSigma2,
w=1.0)
rhos, sigmas = torch.tensor(rhos).to(device), torch.tensor(
sigmas).to(device)
# --------------------------------
# (3) initialize x, and pre-calculation
# --------------------------------
x = util.single2tensor4(img_L).to(device)
img_L_tensor, k_tensor = util.single2tensor4(
img_L), util.single2tensor4(np.expand_dims(k, 2))
[k_tensor, img_L_tensor] = util.todevice([k_tensor, img_L_tensor],
device)
FB, FBC, F2B, FBFy = sr.pre_calculate(img_L_tensor, k_tensor, sf)
# --------------------------------
# (4) main iterations
# --------------------------------
for i in range(iter_num):
# --------------------------------
# step 1, FFT
# --------------------------------
tau = rhos[i].float().repeat(1, 1, 1, 1)
x = sr.data_solution(x, FB, FBC, F2B, FBFy, tau, sf)
if 'ircnn' in model_name:
current_idx = np.int(
np.ceil(sigmas[i].cpu().numpy() * 255. / 2.) - 1)
if current_idx != former_idx:
model.load_state_dict(model25[str(current_idx)],
strict=True)
model.eval()
for _, v in model.named_parameters():
v.requires_grad = False
model = model.to(device)
former_idx = current_idx
# --------------------------------
# step 2, denoiser
# --------------------------------
if x8:
x = util.augment_img_tensor4(x, i % 8)
if 'drunet' in model_name:
x = torch.cat((x, sigmas[i].float().repeat(
1, 1, x.shape[2], x.shape[3])),
dim=1)
x = utils_model.test_mode(model,
x,
mode=2,
refield=32,
min_size=256,
modulo=16)
elif 'ircnn' in model_name:
x = model(x)
if x8:
if i % 8 == 3 or i % 8 == 5:
x = util.augment_img_tensor4(x, 8 - i % 8)
else:
x = util.augment_img_tensor4(x, i % 8)
# --------------------------------
# (3) img_E
# --------------------------------
img_E = util.tensor2uint(x)
if n_channels == 1:
img_H = img_H.squeeze()
if save_E:
util.imsave(
img_E,
os.path.join(
E_path, img_name + '_k' + str(k_index) + '_' +
model_name + '.png'))
# --------------------------------
# (4) img_LEH
# --------------------------------
if save_LEH:
img_L = util.single2uint(img_L)
k_v = k / np.max(k) * 1.0
k_v = util.single2uint(np.tile(k_v[..., np.newaxis],
[1, 1, 3]))
k_v = cv2.resize(k_v, (3 * k_v.shape[1], 3 * k_v.shape[0]),
interpolation=cv2.INTER_NEAREST)
img_I = cv2.resize(img_L,
(sf * img_L.shape[1], sf * img_L.shape[0]),
interpolation=cv2.INTER_NEAREST)
img_I[:k_v.shape[0], -k_v.shape[1]:, :] = k_v
img_I[:img_L.shape[0], :img_L.shape[1], :] = img_L
util.imshow(np.concatenate([img_I, img_E, img_H], axis=1),
title='LR / Recovered / Ground-truth'
) if show_img else None
util.imsave(
np.concatenate([img_I, img_E, img_H], axis=1),
os.path.join(E_path,
img_name + '_k' + str(k_index) + '_LEH.png'))
if save_L:
util.imsave(
util.single2uint(img_L),
os.path.join(E_path,
img_name + '_k' + str(k_index) + '_LR.png'))
psnr = util.calculate_psnr(
img_E, img_H, border=border) # change with your own border
test_results['psnr'].append(psnr)
logger.info('{:->4d}--> {:>10s} --k:{:>2d} PSNR: {:.2f}dB'.format(
idx + 1, img_name + ext, k_index, psnr))
# --------------------------------
# Average PSNR
# --------------------------------
ave_psnr = sum(test_results['psnr']) / len(test_results['psnr'])
logger.info(
'------> Average PSNR of ({}), kernel: ({}) sigma: ({:.2f}): {:.2f} dB'
.format(testset_name, k_index, noise_level_model, ave_psnr))
test_results_ave['psnr'].append(ave_psnr)
def main():
# ----------------------------------------
# Preparation
# ----------------------------------------
noise_level_img = 0 # default: 0, noise level for LR image
noise_level_model = noise_level_img # noise level for model
model_name = 'dpsr_x4_gan' # 'dpsr_x2' | 'dpsr_x3' | 'dpsr_x4' | 'dpsr_x4_gan'
testset_name = 'set5' # test set, 'set5' | 'srbsd68'
need_degradation = True # default: True
x8 = False # default: False, x8 to boost performance
sf = [int(s) for s in re.findall(r'\d+', model_name)][0] # scale factor
show_img = False # default: False
task_current = 'sr' # 'dn' for denoising | 'sr' for super-resolution
n_channels = 3 # fixed
nc = 96 # fixed, number of channels
nb = 16 # fixed, number of conv layers
model_pool = 'model_zoo' # fixed
testsets = 'testsets' # fixed
results = 'results' # fixed
result_name = testset_name + '_' + model_name
border = sf if task_current == 'sr' else 0 # shave boader to calculate PSNR and SSIM
model_path = os.path.join(model_pool, model_name + '.pth')
# ----------------------------------------
# L_path, E_path, H_path
# ----------------------------------------
L_path = os.path.join(testsets,
testset_name) # L_path, for Low-quality images
H_path = L_path # H_path, for High-quality images
E_path = os.path.join(results, result_name) # E_path, for Estimated images
util.mkdir(E_path)
if H_path == L_path:
need_degradation = True
logger_name = result_name
utils_logger.logger_info(logger_name,
log_path=os.path.join(E_path,
logger_name + '.log'))
logger = logging.getLogger(logger_name)
need_H = True if H_path is not None else False
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# ----------------------------------------
# load model
# ----------------------------------------
from models.network_dpsr import MSRResNet_prior as net
model = net(in_nc=n_channels + 1,
out_nc=n_channels,
nc=nc,
nb=nb,
upscale=sf,
act_mode='R',
upsample_mode='pixelshuffle')
model.load_state_dict(torch.load(model_path), strict=False)
model.eval()
for k, v in model.named_parameters():
v.requires_grad = False
model = model.to(device)
logger.info('Model path: {:s}'.format(model_path))
number_parameters = sum(map(lambda x: x.numel(), model.parameters()))
logger.info('Params number: {}'.format(number_parameters))
test_results = OrderedDict()
test_results['psnr'] = []
test_results['ssim'] = []
test_results['psnr_y'] = []
test_results['ssim_y'] = []
logger.info('model_name:{}, model sigma:{}, image sigma:{}'.format(
model_name, noise_level_img, noise_level_model))
logger.info(L_path)
L_paths = util.get_image_paths(L_path)
H_paths = util.get_image_paths(H_path) if need_H else None
for idx, img in enumerate(L_paths):
# ------------------------------------
# (1) img_L
# ------------------------------------
img_name, ext = os.path.splitext(os.path.basename(img))
# logger.info('{:->4d}--> {:>10s}'.format(idx+1, img_name+ext))
img_L = util.imread_uint(img, n_channels=n_channels)
img_L = util.uint2single(img_L)
# degradation process, bicubic downsampling + Gaussian noise
if need_degradation:
img_L = util.modcrop(img_L, sf)
img_L = util.imresize_np(img_L, 1 / sf)
np.random.seed(seed=0) # for reproducibility
img_L += np.random.normal(0, noise_level_img / 255., img_L.shape)
util.imshow(util.single2uint(img_L),
title='LR image with noise level {}'.format(
noise_level_img)) if show_img else None
img_L = util.single2tensor4(img_L)
noise_level_map = torch.full((1, 1, img_L.size(2), img_L.size(3)),
noise_level_model / 255.).type_as(img_L)
img_L = torch.cat((img_L, noise_level_map), dim=1)
img_L = img_L.to(device)
# ------------------------------------
# (2) img_E
# ------------------------------------
if not x8:
img_E = model(img_L)
else:
img_E = utils_model.test_mode(model, img_L, mode=3, sf=sf)
img_E = util.tensor2uint(img_E)
if need_H:
# --------------------------------
# (3) img_H
# --------------------------------
img_H = util.imread_uint(H_paths[idx], n_channels=n_channels)
img_H = img_H.squeeze()
img_H = util.modcrop(img_H, sf)
# --------------------------------
# PSNR and SSIM
# --------------------------------
psnr = util.calculate_psnr(img_E, img_H, border=border)
ssim = util.calculate_ssim(img_E, img_H, border=border)
test_results['psnr'].append(psnr)
test_results['ssim'].append(ssim)
logger.info('{:s} - PSNR: {:.2f} dB; SSIM: {:.4f}.'.format(
img_name + ext, psnr, ssim))
util.imshow(np.concatenate([img_E, img_H], axis=1),
title='Recovered / Ground-truth') if show_img else None
if np.ndim(img_H) == 3: # RGB image
img_E_y = util.rgb2ycbcr(img_E, only_y=True)
img_H_y = util.rgb2ycbcr(img_H, only_y=True)
psnr_y = util.calculate_psnr(img_E_y, img_H_y, border=border)
ssim_y = util.calculate_ssim(img_E_y, img_H_y, border=border)
test_results['psnr_y'].append(psnr_y)
test_results['ssim_y'].append(ssim_y)
# ------------------------------------
# save results
# ------------------------------------
util.imsave(img_E, os.path.join(E_path, img_name + '.png'))
if need_H:
ave_psnr = sum(test_results['psnr']) / len(test_results['psnr'])
ave_ssim = sum(test_results['ssim']) / len(test_results['ssim'])
logger.info(
'Average PSNR/SSIM(RGB) - {} - x{} --PSNR: {:.2f} dB; SSIM: {:.4f}'
.format(result_name, sf, ave_psnr, ave_ssim))
if np.ndim(img_H) == 3:
ave_psnr_y = sum(test_results['psnr_y']) / len(
test_results['psnr_y'])
ave_ssim_y = sum(test_results['ssim_y']) / len(
test_results['ssim_y'])
logger.info(
'Average PSNR/SSIM( Y ) - {} - x{} - PSNR: {:.2f} dB; SSIM: {:.4f}'
.format(result_name, sf, ave_psnr_y, ave_ssim_y))
def main():
# ----------------------------------------
# load kernels
# ----------------------------------------
PSF_grid = np.load('./data/AC254-075-A-ML-Zemax(ZMX).npz')['PSF']
PSF_grid = PSF_grid.astype(np.float32)
gx, gy = PSF_grid.shape[:2]
for xx in range(gx):
for yy in range(gy):
PSF_grid[xx, yy] = PSF_grid[xx, yy] / np.sum(PSF_grid[xx, yy],
axis=(0, 1))
# ----------------------------------------
# load model
# ----------------------------------------
stage = 8
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = net(n_iter=stage,
h_nc=64,
in_nc=4,
out_nc=3,
nc=[64, 128, 256, 512],
nb=2,
act_mode="R",
downsample_mode='strideconv',
upsample_mode="convtranspose")
model_code = 'iter17000'
loaded_state = torch.load(
'/home/xiu/databag/deblur/models/ZEMAX/uabcnet_{}.pth'.format(
model_code))
model.load_state_dict(loaded_state, strict=True)
model.eval()
for _, v in model.named_parameters():
v.requires_grad = False
model = model.to(device)
img_names = glob.glob(
'/home/xiu/databag/deblur/ICCV2021/suo_image/*/AC254-075-A-ML-Zemax(ZMX).bmp'
)
img_names.sort()
for img_id, img_name in enumerate(img_names):
img_L = cv2.imread(img_name)
img_L = img_L.astype(np.float32)
W, H = img_L.shape[:2]
num_patch = [6, 8]
#positional alpha-beta parameters for HQS
ab_numpy = np.loadtxt(
'/home/xiu/databag/deblur/models/ZEMAX/ab_{}.txt'.format(
model_code)).astype(np.float32).reshape(gx, gy, stage * 2, 3)
ab = torch.tensor(ab_numpy, device=device, requires_grad=False)
#save img_L
t0 = time.time()
px_start = 0
py_start = 0
PSF_patch = PSF_grid[px_start:px_start + num_patch[0],
py_start:py_start + num_patch[1]]
#block_expand = 1
patch_L = img_L[px_start * W // gx:(px_start + num_patch[0]) * W // gx,
py_start * H // gy:(py_start + num_patch[1]) * H //
gy, :]
p_W, p_H = patch_L.shape[:2]
expand = max(PSF_grid.shape[2] // 2, p_W // 16)
block_expand = expand
patch_L_wrap = util_deblur.wrap_boundary_liu(
patch_L, (p_W + block_expand * 2, p_H + block_expand * 2))
#centralize
patch_L_wrap = np.hstack((patch_L_wrap[:, -block_expand:, :],
patch_L_wrap[:, :p_H + block_expand, :]))
patch_L_wrap = np.vstack((patch_L_wrap[-block_expand:, :, :],
patch_L_wrap[:p_W + block_expand, :, :]))
x = util.uint2single(patch_L_wrap)
x = util.single2tensor4(x)
k_all = []
for h_ in range(num_patch[1]):
for w_ in range(num_patch[0]):
k_all.append(util.single2tensor4(PSF_patch[w_, h_]))
k = torch.cat(k_all, dim=0)
[x, k] = [el.to(device) for el in [x, k]]
ab_patch = F.softplus(ab[px_start:px_start + num_patch[0],
py_start:py_start + num_patch[1]])
cd = []
for h_ in range(num_patch[1]):
for w_ in range(num_patch[0]):
cd.append(ab_patch[w_:w_ + 1, h_])
cd = torch.cat(cd, dim=0)
x_E = model.forward_patchwise(x, k, cd, num_patch, [W // gx, H // gy])
x_E = x_E[..., block_expand:block_expand + p_W,
block_expand:block_expand + p_H]
patch_L = patch_L_wrap.astype(np.uint8)
patch_E = util.tensor2uint(x_E)
t1 = time.time()
print('[{}/{}]: {} s per frame'.format(img_id, len(img_names),
t1 - t0))
xk = patch_E
xk = xk.astype(np.uint8)
cv2.imshow('res', xk)
cv2.imshow('input', patch_L.astype(np.uint8))
key = cv2.waitKey(-1)
if key == ord('q'):
break
