def main(): # Experiment specifications imagename = 'image_Lena512rgb.png' # Load noise-free image y = np.array(Image.open(imagename)) / 255 # Possible noise types to be generated 'gw', 'g1', 'g2', 'g3', 'g4', 'g1w', # 'g2w', 'g3w', 'g4w'. noise_type = 'g3' noise_var = 0.02 # Noise variance seed = 0 # seed for pseudorandom noise realization # Generate noise with given PSD noise, psd, kernel = get_experiment_noise(noise_type, noise_var, seed, y.shape) # N.B.: For the sake of simulating a more realistic acquisition scenario, # the generated noise is *not* circulant. Therefore there is a slight # discrepancy between PSD and the actual PSD computed from infinitely many # realizations of this noise with different seeds. # Generate noisy image corrupted by additive spatially correlated noise # with noise power spectrum PSD z = np.atleast_3d(y) + np.atleast_3d(noise) # Call BM3D With the default settings. y_est = bm3d_rgb(z, psd) # To include refiltering: # y_est = bm3d_rgb(z, psd, 'refilter'); # For other settings, use BM3DProfile. # profile = BM3DProfile(); # equivalent to profile = BM3DProfile('np'); # profile.gamma = 6; # redefine value of gamma parameter # y_est = bm3d_rgb(z, psd, profile); # Note: For white noise, you may instead of the PSD # also pass a standard deviation # y_est = bm3d_rgb(z, sqrt(noise_var)); # If the different channels have varying PSDs, you can supply a MxNx3 PSD or a list of 3 STDs: # y_est = bm3d_rgb(z, np.concatenate((psd1, psd2, psd3), 2)) # y_est = bm3d_rgb(z, [sigma1, sigma2, sigma3]) psnr = get_psnr(y, y_est) print("PSNR:", psnr) # PSNR ignoring 16-pixel wide borders (as used in the paper), due to refiltering potentially leaving artifacts # on the pixels near the boundary of the image when noise is not circulant psnr_cropped = get_cropped_psnr(y, y_est, [16, 16]) print("PSNR cropped:", psnr_cropped) # Ignore values outside range for display (or plt gives an error for multichannel input) y_est = np.minimum(np.maximum(y_est, 0), 1) z_rang = np.minimum(np.maximum(z, 0), 1) plt.title("y, z, y_est") plt.imshow(np.concatenate((y, np.squeeze(z_rang), y_est), axis=1)) plt.show()
def bm3d(img, sigma=None): if sigma == None: sigma = est_sigma(img) if sigma < 4: sigma **= 2 print("Used sigma = %.2f" % sigma) try: tic = time() dumb = bm3d_rgb(img, sigma) tock = time() print("Execute time = %.2f" % (tock - tic), "s") except: print("Something went wrong, trying different method.") tic = time() dumb = cv.xphoto.bm3dDenoising(img) tock = time() print("Execute time = %.2f" % (tock - tic), "s") dumb = np.clip(dumb, 0.0, 255.0).astype('uint8') return dumb
def generate_red_img(noisy_img, pred_psd, gt_img): noisy_img = np.array(noisy_img) noisy_img = noisy_img / 255.0 gt_img = np.array(gt_img) gt_img = gt_img / 255.0 profile = BM3DProfile() profile.bs_ht = random.choice([2, 4, 8]) profile.transform_2d_wiener_name = random.choice(['dct', 'dst']) profile.bs_wiener = random.choice([4, 5, 6, 7, 8, 9, 10, 11, 12]) cspace = random.choice(['opp', 'YCbCr']) cff = random.uniform(1, 15) red_img = bm3d_rgb(noisy_img, cff * pred_psd[0], profile, colorspace=cspace) red_img = np.minimum(np.maximum(red_img, 0), 1) psnr = get_psnr(gt_img, red_img) red_img = Image.fromarray(np.uint8(red_img * 255.0)) return red_img, cff, profile.bs_ht, cspace, profile.transform_2d_wiener_name, profile.bs_wiener, psnr
def generate_red_img(noisy_img, pred_psd, gt_img): noisy_img = np.array(noisy_img) noisy_img = noisy_img / 255.0 gt_img = np.array(gt_img) gt_img = gt_img / 255.0 profile = BM3DProfile() profile.bs_ht = 4 profile.transform_2d_wiener_name = 'dst' profile.bs_wiener = choose_neighborhood cspace = 'YCbCr' cff = choose_cff red_img = bm3d_rgb(noisy_img, cff * pred_psd[0], profile, colorspace=cspace) red_img = np.minimum(np.maximum(red_img, 0), 1) psnr = get_psnr(gt_img, red_img) red_img = Image.fromarray(np.uint8(red_img * 255.0)) return red_img, cff, profile.bs_ht, cspace, profile.transform_2d_wiener_name, profile.bs_wiener, psnr
def proxg(x, rho): return bm3d_rgb(x, bsigma)
""" Define proximal operator of (implicit, unknown) regularisation term for PPP problem. In this case we use BM3D :cite:`dabov-2008-image` as the denoiser, using the [code](https://pypi.org/project/bm3d) released with :cite:`makinen-2019-exact`. """ bsigma = 6.1e-2 # Denoiser parameter def proxg(x, rho): return bm3d_rgb(x, bsigma) """ Construct a baseline solution and initaliser for the PPP solution by BM3D denoising of a simple bilinear demosaicing solution. The `3 * nsigma` denoising parameter for BM3D is chosen empirically for best performance. """ imgb = bm3d_rgb(demosaic(sn), 3 * nsigma) """ Set algorithm options for PPP solver, including use of bilinear demosaiced solution as an initial solution. """ opt = PPP.Options({ 'Verbose': True, 'RelStopTol': 1e-3, 'MaxMainIter': 12, 'rho': 1.8e-1, 'Y0': imgb }) """ Create solver object and solve, returning the the demosaiced image ``imgp``. """
""" Define proximal operator of (implicit, unknown) regularisation term for PPP problem. In this case we use BM3D :cite:`dabov-2008-image` as the denoiser, using the [code](https://pypi.org/project/bm3d) released with :cite:`makinen-2019-exact`. """ bsigma = 6.1e-2 # Denoiser parameter def proxg(x, rho): return bm3d_rgb(x, bsigma) """ Construct a baseline solution and initaliser for the PPP solution by BM3D denoising of a simple bilinear demosaicing solution. The `3 * nsigma` denoising parameter for BM3D is chosen empirically for best performance. """ imgb = bm3d_rgb(bilinear_demosaic(sn), 3 * nsigma) """ Set algorithm options for PPP solver, including use of bilinear demosaiced solution as an initial solution. """ opt = PPP.Options({ 'Verbose': True, 'RelStopTol': 1e-3, 'MaxMainIter': 12, 'rho': 1.8e-1, 'Y0': imgb }) """ Create solver object and solve, returning the the demosaiced image ``imgp``. """
def main(t, sigma, args=None): # Experiment specifications #imagename = 'image_Lena512rgb.png' imagepath = args.path + '\\' imagename = imagepath + 'ref\\' + t + '_r.bmp' # Load noise-free image y = np.array(Image.open(imagename)) / 255 # Possible noise types to be generated 'gw', 'g1', 'g2', 'g3', 'g4', 'g1w', # 'g2w', 'g3w', 'g4w'. noise_type = 'gw' noise_var = (sigma / 255)**2 # Noise variance 25 std seed = 0 # seed for pseudorandom noise realization # Generate noise with given PSD noise, psd, kernel = get_experiment_noise(noise_type, noise_var, seed, y.shape) # N.B.: For the sake of simulating a more realistic acquisition scenario, # the generated noise is *not* circulant. Therefore there is a slight # discrepancy between PSD and the actual PSD computed from infinitely many # realizations of this noise with different seeds. # Generate noisy image corrupted by additive spatially correlated noise # with noise power spectrum PSD z = np.atleast_3d(y) + np.atleast_3d(noise) z_rang = np.minimum(np.maximum(z, 0), 1) noisyimagename = imagepath + 'noisy\\' + t + '_g.bmp' plt.imsave(noisyimagename, z_rang) z = np.array(Image.open(noisyimagename)) / 255 # Call BM3D With the default settings. y_est = bm3d_rgb(z, psd) # To include refiltering: # y_est = bm3d_rgb(z, psd, 'refilter'); # For other settings, use BM3DProfile. # profile = BM3DProfile(); # equivalent to profile = BM3DProfile('np'); # profile.gamma = 6; # redefine value of gamma parameter # y_est = bm3d_rgb(z, psd, profile); # Note: For white noise, you may instead of the PSD # also pass a standard deviation # y_est = bm3d_rgb(z, sqrt(noise_var)); # If the different channels have varying PSDs, you can supply a MxNx3 PSD or a list of 3 STDs: # y_est = bm3d_rgb(z, np.concatenate((psd1, psd2, psd3), 2)) # y_est = bm3d_rgb(z, [sigma1, sigma2, sigma3]) psnr = get_psnr(y, y_est) print("PSNR:", psnr) # PSNR ignoring 16-pixel wide borders (as used in the paper), due to refiltering potentially leaving artifacts # on the pixels near the boundary of the image when noise is not circulant psnr_cropped = get_cropped_psnr(y, y_est, [16, 16]) print("PSNR cropped:", psnr_cropped) # Ignore values outside range for display (or plt gives an error for multichannel input) y_est = np.minimum(np.maximum(y_est, 0), 1) z_rang = np.minimum(np.maximum(z, 0), 1) opath = 'bm3d_{0}.bmp'.format(t) plt.imsave(opath.format(t), y_est) y_est = np.array(Image.open(opath.format(t))) / 255 psnr = get_psnr(y, y_est) print("PSNR 2:", psnr) mse = ((y_est - y)**2).mean() * 255 print("MSE:", mse) #plt.imsave(imagepath+ 'noisy\\' + t + '_g.bmp', z_rang) # TEST CV2 PSNR try: from skimage.metrics import structural_similarity as compare_ssim except Exception: from skimage.measure import compare_ssim import cv2 opath = 'bm3d_{0}.bmp'.format(t) argref = imagename # path of noisy image d = cv2.imread(opath) tref = cv2.imread(argref) (score, diff) = compare_ssim(tref, d, full=True, multichannel=True) psnr2 = cv2.PSNR(tref, d) print('#######################') print('CV2 PSNR, SSIM: {:.2f}, {:.2f}'.format(psnr2, score)) print('#######################') print('') #plt.title("y, z, y_est") #plt.imshow(np.concatenate((y, np.squeeze(z_rang), y_est), axis=1)) #plt.show() return psnr, mse
def denoising(noise_im, clean_im, LR=1e-2, sigma=3, rho=1, eta=0.5, total_step=30, prob1_iter=500, noise_level=None, result_root=None, f=None): input_depth = 3 latent_dim = 3 en_net = UNet(input_depth, latent_dim).to(device) de_net = UNet(latent_dim, input_depth).to(device) parameters = [p for p in en_net.parameters()] + [p for p in de_net.parameters()] optimizer = torch.optim.Adam(parameters, lr=LR) l2_loss = torch.nn.MSELoss().cuda() i0 = np_to_torch(noise_im).to(device) noise_im_torch = np_to_torch(noise_im).to(device) i0_til_torch = np_to_torch(noise_im).to(device) Y = torch.zeros_like(noise_im_torch).to(device) diff_original_np = noise_im.astype(np.float32) - clean_im.astype(np.float32) diff_original_name = 'Original_dis.png' save_hist(diff_original_np, result_root+diff_original_name) best_psnr = 0 for i in range(total_step): ################################# sub-problem 1 ############################### for i_1 in range(prob1_iter): optimizer.zero_grad() mean = en_net(noise_im_torch) z = sample_z(mean) out = de_net(z) total_loss = 0.5 * l2_loss(out, noise_im_torch) total_loss += 0.5 * (1/sigma**2)*l2_loss(mean, i0) total_loss += (rho/2) * l2_loss(i0 + Y, i0_til_torch) total_loss.backward() optimizer.step() with torch.no_grad(): i0 = ((1/sigma**2)*mean.detach() + rho*(i0_til_torch - Y)) / ((1/sigma**2) + rho) with torch.no_grad(): ################################# sub-problem 2 ############################### i0_np = torch_to_np(i0) Y_np = torch_to_np(Y) sig = eval_sigma(i+1, noise_level) i0_til_np = bm3d.bm3d_rgb(i0_np.transpose(1, 2, 0) + Y_np.transpose(1, 2, 0), sig).transpose(2, 0, 1) i0_til_torch = np_to_torch(i0_til_np).to(device) ################################# sub-problem 3 ############################### Y = Y + eta * (i0 - i0_til_torch) ############################################################################### Y_name = 'Y_{:04d}'.format(i) + '.png' i0_name = 'i0_num_epoch_{:04d}'.format(i) + '.png' mean_name = 'Latent_im_num_epoch_{:04d}'.format(i) + '.png' out_name = 'res_of_dec_num_epoch_{:04d}'.format(i) + '.png' diff_name = 'Latent_dis_num_epoch_{:04d}'.format(i) + '.png' Y_np = torch_to_np(Y) Y_norm_np = np.sqrt((Y_np*Y_np).sum(0)) save_heatmap(Y_norm_np, result_root + Y_name) save_torch(mean, result_root + mean_name) save_torch(out, result_root + out_name) save_torch(i0, result_root + i0_name) mean_np = torch_to_np(mean) diff_np = mean_np - clean_im save_hist(diff_np, result_root + diff_name) i0_til_np = torch_to_np(i0_til_torch).clip(0, 255) psnr = compare_psnr(clean_im.transpose(1, 2, 0), i0_til_np.transpose(1, 2, 0), 255) ssim = compare_ssim(clean_im.transpose(1, 2, 0), i0_til_np.transpose(1, 2, 0), multichannel=True, data_range=255) i0_til_pil = np_to_pil(i0_til_np) i0_til_pil.save(os.path.join(result_root, '{}'.format(i) + '.png')) print('Iteration: {:02d}, VAE Loss: {:f}, PSNR: {:f}, SSIM: {:f}'.format(i, total_loss.item(), psnr, ssim), file=f, flush=True) if best_psnr < psnr: best_psnr = psnr best_ssim = ssim else: break return i0_til_np, best_psnr, best_ssim
def denoising_fixd_noise_level(noise_im, clean_im, LR=1e-2, sigma=5, rho=1, eta=0.5, alpha=0.5, total_step=19, prob1_iter=1000, noise_level=None, result_root=None, fo=None): sig = noise_level r_bm3d = bm3d.bm3d_rgb(noise_im.transpose(1, 2, 0), sig) r_bm3d = np.clip(r_bm3d, 0, 255) psnr_bm3d = compare_psnr(clean_im.transpose(1, 2, 0), r_bm3d, 255) ssim_bm3d = compare_ssim(r_bm3d, clean_im.transpose(1, 2, 0), multichannel=True, data_range=255) print('noise level {} '.format(noise_level), file=fo, flush=True) print('PSNR_BM3D: {}, SSIM_BM3D: {}'.format(psnr_bm3d, ssim_bm3d), file=fo, flush=True) r_bm3d = Image.fromarray(r_bm3d.astype(np.uint8)) r_bm3d.save(result_root + 'bm3d_result.png') input_depth = 3 latent_dim = 3 en_net = UNet(input_depth, latent_dim, need_sigmoid=True).cuda() de_net = UNet(latent_dim, input_depth, need_sigmoid=True).cuda() en_optimizer = torch.optim.Adam(en_net.parameters(), lr=LR) de_optimizer = torch.optim.Adam(de_net.parameters(), lr=LR) l2_loss = torch.nn.MSELoss().cuda() noise_im = noise_im / 255.0 i0 = torch.Tensor(noise_im[None, ...]).cuda() noise_im_torch = torch.Tensor(noise_im)[None, ...].cuda() Y = torch.zeros_like(noise_im_torch).cuda() i0_til_torch = torch.Tensor(noise_im[None, ...]).cuda() output = None for i in range(total_step): ############################### sub-problem 1 ################################# for i_1 in range(prob1_iter): mean_i = en_net(noise_im_torch) eps = mean_i.clone().normal_() out = de_net(mean_i + eps) total_loss = 0.5 * l2_loss(out, noise_im_torch) total_loss += 1 / (2 * sigma**2) * l2_loss(mean_i, i0) en_optimizer.zero_grad() de_optimizer.zero_grad() total_loss.backward() en_optimizer.step() de_optimizer.step() with torch.no_grad(): i0 = ((1 / sigma**2) * mean_i + rho * (i0_til_torch - Y) + alpha * noise_im_torch) / ( (1 / sigma**2) + rho + alpha) with torch.no_grad(): ############################### sub-problem 2 ################################# i0_np = i0.cpu().squeeze().detach().numpy() Y_np = Y.cpu().squeeze().detach().numpy() tmp = i0_np.transpose(1, 2, 0) + Y_np.transpose(1, 2, 0) tmp = np.clip(tmp, 0, 1) sig = noise_level i0_til_np = bm3d.bm3d_rgb(tmp * 255, sig) / 255 i0_til_torch = torch.Tensor( i0_til_np.transpose(2, 0, 1)[None, ...]).cuda() ############################### sub-problem 3 ################################# Y = Y + eta * (i0 - i0_til_torch) ############################################################################### i0_til_np = i0_til_torch.cpu().squeeze().detach().numpy() i0_til_np = np.clip(i0_til_np, 0, 1) output = i0_til_np psnr_gt = compare_psnr(clean_im.transpose(1, 2, 0), 255 * i0_til_np.transpose(1, 2, 0), 255) ssim_gt = compare_ssim(255 * i0_til_np.transpose(1, 2, 0), clean_im.transpose(1, 2, 0), multichannel=True, data_range=255) if not i % 5: denoise_obj_pil = Image.fromarray((tmp * 255).astype(np.uint8)) Y_np = Y.cpu().squeeze().detach().numpy() i0_np = np.clip(i0_np, 0, 1) i0_pil = Image.fromarray( np.uint8(255 * i0_np.transpose(1, 2, 0))) i0_til_np = i0_til_np.transpose(1, 2, 0) i0_til_pil = Image.fromarray( (255 * i0_til_np).astype(np.uint8)) mean_i_np = mean_i.cpu().squeeze().detach().numpy().clip(0, 1) mean_i_pil = Image.fromarray( (255 * mean_i_np.transpose(1, 2, 0)).astype(np.uint8)) out_np = out.cpu().squeeze().detach().numpy().clip(0, 1) out_pil = Image.fromarray( (255 * out_np.transpose(1, 2, 0)).astype(np.uint8)) denoise_obj_name = 'denoise_obj_{:04d}'.format(i) + '.png' Y_name = 'Y_{:04d}'.format(i) + '.png' i0_name = 'i0_num_epoch_{:04d}'.format(i) + '.png' result_name = 'num_epoch_{:04d}'.format(i) + '.png' mean_i_name = 'Latent_im_num_epoch_{:04d}'.format(i) + '.png' out_name = 'res_of_dec_num_epoch_{:04d}'.format(i) + '.png' denoise_obj_pil.save(result_root + denoise_obj_name) Y_norm_np = np.sqrt((Y_np * Y_np).sum(0)) save_heatmap(Y_norm_np, result_root + Y_name) i0_pil.save(result_root + i0_name) i0_til_pil.save(result_root + result_name) mean_i_pil.save(result_root + mean_i_name) out_pil.save(result_root + out_name) print('Iteration %02d Loss %f PSNR_gt: %f, SSIM_gt: %f' % (i, total_loss.item(), psnr_gt, ssim_gt), file=fo, flush=True) psnr = psnr_gt ssim = ssim_gt output_pil = Image.fromarray( (255 * output.transpose(1, 2, 0)).astype(np.uint8)) output_pil.save(result_root + 'ours_result.png') return psnr, ssim, psnr_bm3d, ssim_bm3d