def inpainting_lista(missing_pixels, dicsize, n_nonzero_coefs, blksize,
overlap, maxiter, grayimg):
img_ori, masked_img, patchvec, maskvec = generate_inpainting_data(
missing_pixels=missing_pixels,
blksize=blksize,
overlap=overlap,
grayimg=grayimg)
# settings
h, w = img_ori.shape[0], img_ori.shape[1]
m = patchvec.shape[0]
n = dicsize
s = n_nonzero_coefs / n
print('[INFO] Dictionary size ({}, {}).'.format(m, n))
print('[INFO] Sparsity {}.'.format(n_nonzero_coefs))
print('[INFO] OMP max iter {}.'.format(maxiter))
D = get_init_dict(dicsize=dicsize, blksize=blksize)
Y = patchvec
# training
model = train_lista(D, Y, maskvec, s, img_ori, blksize, overlap)
# testing
with torch.no_grad():
Y_pt = torch.from_numpy(Y.T / 255.0).cuda()
coef_pt = model(Y_pt)
coef = coef_pt.cpu().detach().numpy().T
# testing
recovered_img = patch2im_for_inpainting(patch_vecs=255.0 *
np.matmul(D, coef),
mask_vecs=maskvec,
imgsize=img_ori.shape,
blksize=blksize,
overlap=overlap)
# save results
if not os.path.exists('./save_lista'):
os.mkdir('./save_lista')
np.save('./save_lista/img_ori', img_ori)
np.save('./save_lista/masked_img', masked_img)
np.save('./save_lista/recovered_img', recovered_img)
return img_ori, masked_img, recovered_img
def denoising_omp(dicsize, sparsity, blksize, overlap, maxiter):
np.random.seed(0)
img_ori, masked_img, mask, patchvec, maskvec = generate_denoising_data(blksize, overlap)
# settings
m = patchvec.shape[0]
n = dicsize
print('[INFO] Dictionary size ({}, {}).'.format(m, n))
print('[INFO] Sparsity {}.'.format(sparsity))
print('[INFO] De-noising max iter {}.'.format(maxiter))
load_results = False
if load_results:
print('[INFO] Load OMP results.')
dict = np.load('./dict.npy')
coef = np.load('./coef.npy')
# reconstruct image
print('[INFO] Reconstruction ...')
recovered_img = patch2im_for_denoising(image=img_ori, patch_vecs=np.matmul(dict, coef), imgsize=img_ori.shape,
blksize=blksize, overlap=overlap)
return img_ori, masked_img, recovered_img
else:
# OMP for de-noising
print('[INFO] Run OMP for denoising')
recovered_img = None
for iter in range(maxiter):
D = get_init_dict(dicsize=dicsize, blksize=blksize)
dict, coef = denoise_omp_once(patchvec, D, sparsity)
np.save('./dict', dict)
np.save('./coef', coef)
# reconstruct image
print('[INFO] Reconstruction ...')
recovered_img = patch2im_for_denoising(image=img_ori, patch_vecs=np.matmul(dict, coef), imgsize=img_ori.shape,
blksize=blksize, overlap=overlap)
return img_ori, masked_img, recovered_img
def inpainting_ksvd(missing_pixels, dicsize, n_nonzero_coefs, blksize, overlap, maxiter, grayimg=False):
img_ori, masked_img, patchvec, maskvec = generate_inpainting_data(missing_pixels=missing_pixels,
blksize=blksize, overlap=overlap, grayimg=grayimg)
# img_ori, masked_img, patchvec, maskvec = load_inpainting_image(blksize=blksize, overlap=overlap)
# settings
h, w = img_ori.shape[0], img_ori.shape[1]
m = patchvec.shape[0]
n = dicsize
print('[INFO] Dictionary size ({}, {}).'.format(m, n))
print('[INFO] Sparsity {}.'.format(n_nonzero_coefs))
print('[INFO] OMP max iter {}.'.format(maxiter))
load_results = True
if load_results:
print('[INFO] Load OMP results.')
if grayimg:
dict = np.load('./D.npy')
coef = np.load('./coef.npy')
recovered_img = patch2im_for_inpainting(patch_vecs=np.matmul(dict, coef), mask_vecs=maskvec,
imgsize=(h, w), blksize=blksize, overlap=overlap)
else:
dictR = np.load('./DR.npy')
dictG = np.load('./DG.npy')
dictB = np.load('./DB.npy')
coefR = np.load('./coefR.npy')
coefG = np.load('./coefG.npy')
coefB = np.load('./coefB.npy')
recovered_imgR = patch2im_for_inpainting(patch_vecs=np.matmul(dictR, coefR), mask_vecs=maskvec,
imgsize=(h, w), blksize=blksize, overlap=overlap)
recovered_imgG = patch2im_for_inpainting(patch_vecs=np.matmul(dictG, coefG), mask_vecs=maskvec,
imgsize=(h, w), blksize=blksize, overlap=overlap)
recovered_imgB = patch2im_for_inpainting(patch_vecs=np.matmul(dictB, coefB), mask_vecs=maskvec,
imgsize=(h, w), blksize=blksize, overlap=overlap)
recovered_img = np.stack((recovered_imgR, recovered_imgG, recovered_imgB), axis=2)
return img_ori, masked_img, recovered_img
else:
# OMP for inpainting
print('[INFO] Run OMP for image inpainting')
if grayimg:
D = get_init_dict(dicsize=dicsize, blksize=blksize)
Y = patchvec
newD, coef = inpainting_ksvd_once(Y, D, maskvec, n_nonzero_coefs=50, maxiter=maxiter)
print('[INFO] Saving dictionary ...')
np.save('./D', newD)
print('[INFO] Saving coefficients ...')
np.save('./coef', coef)
# recover data
print('[INFO] Reconstruction ...')
recovered_img = patch2im_for_inpainting(patch_vecs=np.matmul(newD, coef), mask_vecs=maskvec, imgsize=img_ori.shape,
blksize=blksize, overlap=overlap)
else:
D = get_init_dict(dicsize=dicsize, blksize=blksize)
RGBname = ['R', 'G', 'B']
recovered_img = np.zeros_like(img_ori)
# do R-G-B
for channel in range(3):
Y = patchvec[:, :, channel]
newD, coef = inpainting_ksvd_once(Y, D, maskvec, n_nonzero_coefs=50, maxiter=maxiter)
print('[INFO] Saving dictionary ...')
np.save('./D{}'.format(RGBname[channel]), newD)
print('[INFO] Saving coefficients ...')
np.save('./coef{}'.format(RGBname[channel]), coef)
# recover data
print('[INFO] Reconstruction ...')
recovered_img_c = patch2im_for_inpainting(patch_vecs=np.matmul(newD, coef), mask_vecs=maskvec, imgsize=img_ori.shape,
blksize=blksize, overlap=overlap)
recovered_img[:, :, channel] = recovered_img_c
return img_ori, masked_img, recovered_img