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 = []
