def save_out_seq(seqnoisy, seqclean, save_dir, sigmaval, suffix, save_noisy, start_count=0): """Saves the denoised and noisy sequences under save_dir """ seq_len = seqnoisy.size()[0] for idx in range(seq_len): # Build Outname fext = OUTIMGEXT noisy_name = os.path.join(save_dir, ('n{}_{}').format(sigmaval, idx+start_count) + fext) if len(suffix) == 0: out_name = os.path.join(save_dir, ('n{}_FastDVDnet_{}').format(sigmaval, idx+start_count) + fext) else: out_name = os.path.join(save_dir, ('n{}_FastDVDnet_{}_{}').format(sigmaval, suffix, idx+start_count) + fext) # Save result if save_noisy: noisyimg = variable_to_cv2_image(seqnoisy[idx].clamp(0., 1.)) cv2.imwrite(noisy_name, noisyimg) outimg = variable_to_cv2_image(seqclean[idx].unsqueeze(dim=0)) if os.path.isfile(out_name): return # Don't overwrite files. cv2.imwrite(out_name, outimg)
def test(self, input_image, output_image): ''' Test image ''' img = extract_n_normalize_image(input_image) x_image = np.reshape(np.array([img]), (1, self.image_size, self.image_size, 1)) sess, _ = self.init_session() y_image = sess.run(self.Y, feed_dict={self.X: x_image}) encoded_image = y_image.reshape((self.image_size, self.image_size)) outimg = variable_to_cv2_image(encoded_image) cv2.imwrite(output_image, outimg)
def noise(imorig): imorig = np.expand_dims(imorig, 0) # Handle odd sizes expanded_h = False expanded_w = False sh_im = imorig.shape if sh_im[2] % 2 == 1: expanded_h = True imorig = np.concatenate((imorig, \ imorig[:, :, -1, :][:, :, np.newaxis, :]), axis=2) # if sh_im[3] % 2 == 1: expanded_w = True imorig = np.concatenate((imorig, \ imorig[:, :, :, -1][:, :, :, np.newaxis]), axis=3) imorig = normalize(imorig) imorig = torch.Tensor(imorig) # Add noise imnoisy = imorig.clone() # with torch.no_grad(): # PyTorch v0.4.0 imorig, imnoisy = Variable(imorig.type(dtype), volatile=True), \ Variable(imnoisy.type(dtype), volatile=True) nsigma = Variable( \ torch.FloatTensor([25/255.]).type(dtype), \ volatile=True) # Measure runtime # Estimate noise and subtract it to the input image im_noise_estim = model(imnoisy, nsigma) outim = torch.clamp(imnoisy - im_noise_estim, 0., 1.) if expanded_h: imorig = imorig[:, :, :-1, :] outim = outim[:, :, :-1, :] imnoisy = imnoisy[:, :, :-1, :] if expanded_w: imorig = imorig[:, :, :, :-1] outim = outim[:, :, :, :-1] imnoisy = imnoisy[:, :, :, :-1] # Save images outimg = variable_to_cv2_image(outim) return outimg
def test_ffdnet(**args): r"""Denoises an input image with FFDNet """ # Init logger logger = init_logger_ipol() # Check if input exists and if it is RGB try: rgb_den = is_rgb(args['input']) except: raise Exception('Could not open the input image') # Open image as a CxHxW torch.Tensor if rgb_den: in_ch = 3 model_fn = 'models/net_rgb.pth' imorig = cv2.imread(args['input']) # from HxWxC to CxHxW, RGB image imorig = (cv2.cvtColor(imorig, cv2.COLOR_BGR2RGB)).transpose(2, 0, 1) else: # from HxWxC to CxHxW grayscale image (C=1) in_ch = 1 model_fn = 'logs/net.pth' imorig = cv2.imread(args['input'], cv2.IMREAD_GRAYSCALE) imorig_copy = imorig.copy() imorig = np.expand_dims(imorig, 0) imorig = np.expand_dims(imorig, 0) # Handle odd sizes expanded_h = False expanded_w = False sh_im = imorig.shape if sh_im[2] % 2 == 1: expanded_h = True imorig = np.concatenate((imorig, \ imorig[:, :, -1, :][:, :, np.newaxis, :]), axis=2) if sh_im[3] % 2 == 1: expanded_w = True imorig = np.concatenate((imorig, \ imorig[:, :, :, -1][:, :, :, np.newaxis]), axis=3) imorig = normalize(imorig) imorig = torch.Tensor(imorig) # Absolute path to model file model_fn = os.path.join(os.path.abspath(os.path.dirname(__file__)), \ model_fn) # Create model print('Loading model ...\n') net = FFDNet(num_input_channels=in_ch) # Load saved weights if args['cuda']: state_dict = torch.load(model_fn) device_ids = [0] model = nn.DataParallel(net, device_ids=device_ids).cuda() else: state_dict = torch.load(model_fn, map_location='cpu') # CPU mode: remove the DataParallel wrapper state_dict = remove_dataparallel_wrapper(state_dict) model = net model.load_state_dict(state_dict) # Sets the model in evaluation mode (e.g. it removes BN) model.eval() # Sets data type according to CPU or GPU modes if args['cuda']: dtype = torch.cuda.FloatTensor else: dtype = torch.FloatTensor # Add noise if args['add_noise']: noise = torch.FloatTensor(imorig.size()).\ normal_(mean=0, std=args['noise_sigma']) imnoisy = imorig + noise else: imnoisy = imorig.clone() # Test mode with torch.no_grad(): # PyTorch v0.4.0 imorig, imnoisy = Variable(imorig.type(dtype)), \ Variable(imnoisy.type(dtype)) nsigma = Variable(torch.FloatTensor([args['noise_sigma']]).type(dtype)) # Measure runtime start_t = time.time() # Estimate noise and subtract it to the input image im_noise_estim = model(imnoisy, nsigma) outim = torch.clamp(imnoisy - im_noise_estim, 0., 1.) stop_t = time.time() if expanded_h: imorig = imorig[:, :, :-1, :] outim = outim[:, :, :-1, :] imnoisy = imnoisy[:, :, :-1, :] if expanded_w: imorig = imorig[:, :, :, :-1] outim = outim[:, :, :, :-1] imnoisy = imnoisy[:, :, :, :-1] # Compute PSNR and log it if rgb_den: print("### RGB denoising ###") else: print("### Grayscale denoising ###") if args['add_noise']: psnr = batch_psnr(outim, imorig, 1.) psnr_noisy = batch_psnr(imnoisy, imorig, 1.) print("----------PSNR noisy {0:0.2f}dB".format(psnr_noisy)) print("----------PSNR denoised {0:0.2f}dB".format(psnr)) else: logger.info("\tNo noise was added, cannot compute PSNR") print("----------Runtime {0:0.4f}s".format(stop_t - start_t)) # Compute difference diffout = 2 * (outim - imorig) + .5 diffnoise = 2 * (imnoisy - imorig) + .5 # Save images if not args['dont_save_results']: noisyimg = variable_to_cv2_image(imnoisy) outimg = variable_to_cv2_image(outim) cv2.imwrite( "bfffd/noisy-" + str(int(args['noise_sigma'] * 255)) + '-' + args['input'], noisyimg) cv2.imwrite( "bfffd/ffdnet-" + str(int(args['noise_sigma'] * 255)) + '-' + args['input'], outimg) if args['add_noise']: cv2.imwrite("noisy_diff.png", variable_to_cv2_image(diffnoise)) cv2.imwrite("ffdnet_diff.png", variable_to_cv2_image(diffout)) (score, diff) = compare_ssim(noisyimg, imorig_copy, full=True) (score2, diff) = compare_ssim(outimg, imorig_copy, full=True) print("----------Noisy ssim: {0:0.4f}".format(score)) print("----------Denoisy ssim: {0:0.4f}".format(score2))
def denoise_seq_dvdnet(seq, noise_std, temp_psz, model_temporal, model_spatial, mc_algo): r"""Denoises a sequence of frames with DVDnet. Args: seq: Tensor. [numframes, 1, C, H, W] array containing the noisy input frames noise_std: Tensor. Standard deviation of the added noise temp_psz: size of the temporal patch model_temp: instance of the PyTorch model of the temporal denoiser spatial_temp: instance of the PyTorch model of the spatial denoiser mc_algo: motion compensation algorithm to apply """ # init arrays to handle contiguous frames and related patches numframes, _, C, H, W = seq.shape ctrlfr_idx = int((temp_psz - 1) // 2) inframes = list() inframes_wrpd = np.empty((temp_psz, H, W, C)) denframes = torch.empty((numframes, C, H, W)).to(seq.device) # build noise map from noise std---assuming Gaussian noise noise_map = noise_std.expand((1, C, H, W)) for fridx in range(numframes): # load input frames # denoise each frame with spatial denoiser when appending if not inframes: # if list not yet created, fill it with temp_patchsz frames for idx in range(temp_psz): relidx = max(0, idx - ctrlfr_idx) inframes.append( spatial_denoise(model_spatial, seq[relidx], noise_map)) else: del inframes[0] relidx = min(numframes - 1, fridx + ctrlfr_idx) inframes.append( spatial_denoise(model_spatial, seq[relidx], noise_map)) # save converted central frame # OpenCV images are HxWxC uint8 images inframes_wrpd[ctrlfr_idx] = variable_to_cv2_image( inframes[ctrlfr_idx], conv_rgb_to_bgr=False) # register frames w.r.t central frame # need to convert them to OpenCV images first for idx in range(temp_psz): if not idx == ctrlfr_idx: img_to_warp = variable_to_cv2_image(inframes[idx], conv_rgb_to_bgr=False) inframes_wrpd[idx] = align_frames(img_to_warp, \ inframes_wrpd[ctrlfr_idx], \ mc_alg=mc_algo) # denoise with temporal model # temp_pszxHxWxC to temp_pszxCxHxW inframes_t = normalize(inframes_wrpd.transpose(0, 3, 1, 2)) inframes_t = torch.from_numpy(inframes_t).contiguous().view( (1, temp_psz * C, H, W)).to(seq.device) # append result to output list denframes[fridx] = temporal_denoise(model_temporal, inframes_t, noise_map) # free memory up del inframes del inframes_wrpd del inframes_t torch.cuda.empty_cache() # convert to appropiate type and return return denframes
def test(args): # Image image = cv2.imread(args.test_path) if image is None: raise Exception(f'File {args.test_path} not found or error') is_gray = utils.is_image_gray(image) image = read_image(args.test_path, is_gray) print("{} image shape: {}".format("Gray" if is_gray else "RGB", image.shape)) # Expand odd shape to even expend_W = False expend_H = False if image.shape[1] % 2 != 0: expend_W = True image = np.concatenate((image, image[:, -1, :][:, np.newaxis, :]), axis=1) if image.shape[2] % 2 != 0: expend_H = True image = np.concatenate((image, image[:, :, -1][:, :, np.newaxis]), axis=2) # Noise image = torch.FloatTensor([image]) # 1 * C(1 / 3) * W * H if args.add_noise: image = utils.add_batch_noise(image, args.noise_sigma) noise_sigma = torch.FloatTensor([args.noise_sigma]) # Model & GPU model = FFDNet(is_gray=is_gray) if args.cuda: image = image.cuda() noise_sigma = noise_sigma.cuda() model = model.cuda() # Dict model_path = args.model_path + ('net_gray.pth' if is_gray else 'net_rgb.pth') print(f"> Loading model param in {model_path}...") state_dict = torch.load(model_path) model.load_state_dict(state_dict) model.eval() print('\n') # Test with torch.no_grad(): start_time = time.time() image_pred = model(image, noise_sigma) stop_time = time.time() print("Test time: {0:.4f}s".format(stop_time - start_time)) # PSNR psnr = utils.batch_psnr(img=image_pred, imclean=image, data_range=1) print("PSNR denoised {0:.2f}dB".format(psnr)) # UnExpand odd if expend_W: image_pred = image_pred[:, :, :-1, :] if expend_H: image_pred = image_pred[:, :, :, :-1] # Save cv2.imwrite("ffdnet.png", utils.variable_to_cv2_image(image_pred)) if args.add_noise: cv2.imwrite("noisy.png", utils.variable_to_cv2_image(image))
denframes = denoise_seq_fastdvdnet(seq=seq, noise_std=noisestd, temp_psz=NUM_IN_FR_EXT, model_temporal=model_temp) # Compute PSNR and log it stop_time = time.time() psnr = batch_psnr(denframes, seq, 1.) psnr_noisy = batch_psnr(seqn.squeeze(), seq, 1.) loadtime = (seq_time - start_time) runtime = (stop_time - seq_time) seq_length = seq.size()[0] print("\tDenoised {} frames in {:.3f}s, loaded seq in {:.3f}s". format(seq_length, runtime, loadtime)) print("\tPSNR noisy {:.4f}dB, PSNR result {:.4f}dB".format( psnr_noisy, psnr)) # Save outputs seq_len = len(seq_list) for idx in range(seq_len): out_name = os.path.join(args.save_path, seq_outnames[idx]) print("Saving %s" % out_name) outimg = variable_to_cv2_image( denframes[idx].unsqueeze(dim=0)) cv2.imwrite(out_name, outimg) seq_list = [] seq_outnames = []