def test_dvdnet(**args):
"""Denoises all sequences present in a given folder. Sequences must be stored as numbered
image sequences. The different sequences must be stored in subfolders under the "test_path" folder.
Inputs:
args (dict) fields:
"model_spatial_file": path to model of the pretrained spatial denoiser
"model_temp_file": path to model of the pretrained temporal denoiser
"test_path": path to sequence to denoise
"suffix": suffix to add to output name
"max_num_fr_per_seq": max number of frames to load per sequence
"noise_sigma": noise level used on test set
"dont_save_results: if True, don't save output images
"no_gpu": if True, run model on CPU
"save_path": where to save outputs as png
"""
start_time = time.time()
# If save_path does not exist, create it
if not os.path.exists(args['save_path']):
os.makedirs(args['save_path'])
logger = init_logger_test(args['save_path'])
# Sets data type according to CPU or GPU modes
if args['cuda']:
device = torch.device('cuda')
else:
device = torch.device('cpu')
# Create models
model_spa = DVDnet_spatial()
model_temp = DVDnet_temporal(num_input_frames=NUM_IN_FRAMES)
# Load saved weights
state_spatial_dict = torch.load(args['model_spatial_file'])
state_temp_dict = torch.load(args['model_temp_file'])
if args['cuda']:
device_ids = [0]
model_spa = nn.DataParallel(model_spa, device_ids=device_ids).cuda()
model_temp = nn.DataParallel(model_temp, device_ids=device_ids).cuda()
else:
# CPU mode: remove the DataParallel wrapper
state_spatial_dict = remove_dataparallel_wrapper(state_spatial_dict)
state_temp_dict = remove_dataparallel_wrapper(state_temp_dict)
model_spa.load_state_dict(state_spatial_dict)
model_temp.load_state_dict(state_temp_dict)
# Sets the model in evaluation mode (e.g. it removes BN)
model_spa.eval()
model_temp.eval()
with torch.no_grad():
# process data
seq, _, _ = open_sequence(args['test_path'],\
False,\
expand_if_needed=False,\
max_num_fr=args['max_num_fr_per_seq'])
seq = torch.from_numpy(seq[:, np.newaxis, :, :, :]).to(device)
seqload_time = time.time()
# Add noise
noise = torch.empty_like(seq).normal_(mean=0, std=args['noise_sigma']).to(device)
seqn = seq + noise
noisestd = torch.FloatTensor([args['noise_sigma']]).to(device)
denframes = denoise_seq_dvdnet(seq=seqn,\
noise_std=noisestd,\
temp_psz=NUM_IN_FRAMES,\
model_temporal=model_temp,\
model_spatial=model_spa,\
mc_algo=MC_ALGO)
den_time = time.time()
# Compute PSNR and log it
psnr = batch_psnr(denframes, seq.squeeze(), 1.)
psnr_noisy = batch_psnr(seqn.squeeze(), seq.squeeze(), 1.)
print("\tPSNR on {} : {}\n".format(os.path.split(args['test_path'])[-1], psnr))
print("\tDenoising time: {:.2f}s".format(den_time - seqload_time))
print("\tSequence loaded in : {:.2f}s".format(seqload_time - start_time))
print("\tTotal time: {:.2f}s\n".format(den_time - start_time))
logger.info("%s, %s, PSNR noisy %fdB, PSNR %f dB" % \
(args['test_path'], args['suffix'], psnr_noisy, psnr))
# Save outputs
if not args['dont_save_results']:
# Save sequence
save_out_seq(seqn, denframes, args['save_path'], int(args['noise_sigma']*255), \
args['suffix'], args['save_noisy'])
# close logger
close_logger(logger)
def test_fastdvdnet(**args):
"""Denoises all sequences present in a given folder. Sequences must be stored as numbered
image sequences. The different sequences must be stored in subfolders under the "test_path" folder.
Inputs:
args (dict) fields:
"model_file": path to model
"test_path": path to sequence to denoise
"suffix": suffix to add to output name
"max_num_fr_per_seq": max number of frames to load per sequence
"noise_sigma": noise level used on test set
"dont_save_results: if True, don't save output images
"no_gpu": if True, run model on CPU
"save_path": where to save outputs as png
"gray": if True, perform denoising of grayscale images instead of RGB
"""
# Start time
start_time = time.time()
# If save_path does not exist, create it
logger = init_logger_test(os.path.dirname(args['save_path']))
# Sets data type according to CPU or GPU modes
if args['cuda']:
device = torch.device('cuda')
else:
device = torch.device('cpu')
# Create models
print('Loading models ...')
model_temp = FastDVDnet(num_input_frames=NUM_IN_FR_EXT)
# Load saved weights
state_temp_dict = torch.load(args['model_file'])
if args['cuda']:
device_ids = [0]
model_temp = nn.DataParallel(model_temp, device_ids=device_ids).cuda()
else:
# CPU mode: remove the DataParallel wrapper
state_temp_dict = remove_dataparallel_wrapper(state_temp_dict)
model_temp.load_state_dict(state_temp_dict)
# Sets the model in evaluation mode (e.g. it removes BN)
model_temp.eval()
with torch.no_grad():
# process data
seq = open_sequence(args['test_path'], args['first'], args['last'],
args['already_norm'])
seq = torch.from_numpy(seq).to(device)
seq_time = time.time()
# Add noise
if not args['already_noisy']:
noise = torch.empty_like(seq).normal_(
mean=0, std=args['noise_sigma']).to(device)
seqn = seq + noise
else:
seqn = seq
noisestd = torch.FloatTensor([args['noise_sigma']]).to(device)
denframes = denoise_seq_fastdvdnet(seq=seqn,
noise_std=noisestd,
temp_psz=NUM_IN_FR_EXT,
model_temporal=model_temp)
# Compute PSNR and log it
stop_time = time.time()
psnr = compute_psnr(denframes.cpu().numpy(), seq.cpu().numpy(), 1.)
psnr_noisy = compute_psnr(seqn.cpu().numpy().squeeze(),
seq.cpu().numpy(), 1.)
loadtime = (seq_time - start_time)
runtime = (stop_time - seq_time)
seq_length = seq.size()[0]
logger.info("Finished denoising {}".format(args['test_path']))
logger.info(
"\tDenoised {} frames in {:.3f}s, loaded seq in {:.3f}s".format(
seq_length, runtime, loadtime))
logger.info("\tPSNR noisy {:.4f}dB, PSNR result {:.4f}dB".format(
psnr_noisy, psnr))
# Save outputs
if not args['dont_save_results']:
# Save sequence
save_out_seq(denframes.cpu(), args['save_path'], args['first'])
# close logger
close_logger(logger)
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 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')
# Measure runtime
start_t = time.time()
# Open image as a CxHxW torch.Tensor
if rgb_den:
in_ch = 3
model_fn = 'net_rgb.pth'
imorig = Image.open(args['input'])
imorig = np.array(imorig, dtype=np.float32).transpose(2, 0, 1)
else:
# from HxWxC to CxHxW grayscale image (C=1)
in_ch = 1
model_fn = 'models/net_gray.pth'
# imorig = cv2.imread(args['input'], cv2.IMREAD_GRAYSCALE)
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,1,2,3]
#model = nn.DataParallel(net, device_ids=device_ids).cuda()
#state_dict = remove_dataparallel_wrapper(state_dict)
model = net
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
# Test mode
with torch.no_grad(): # PyTorch v0.4.0
imorig = Variable(imorig.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(imorig, nsigma)
stop_t = time.time()
# log time
if rgb_den:
print("### RGB denoising ###")
else:
print("### Grayscale denoising ###")
print("\tRuntime {0:0.4f}s".format(stop_t - start_t))
# Save noises
noise = variable_to_numpy(imorig.to(3) - im_noise_estim).transpose(1, 2, 0)
filename = args['input'].split('/')[-1].split('.')[0]
if args['save_path']:
sio.savemat(
'./output_noise/' + args['save_path'] + '/' + filename + '.mat',
{'Noisex': noise})
else:
sio.savemat('./output_noise/' + filename + '.mat', {'Noisex': noise})
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 cuda:
state_dict = torch.load(model_fn)
device_ids = [0]
model = nn.DataParallel(net, device_ids=device_ids).cuda()
# Sets data type according to CPU or GPU modes
dtype = torch.cuda.FloatTensor
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
dtype = torch.FloatTensor
model.load_state_dict(state_dict)
# If use the mid-saved model, uncomment the following code
# model.load_state_dict(state_dict['state_dict'])
# Sets the model in evaluation mode (e.g. it removes BN)
model.eval()
# Load the fingerprint generated with a set of flat images
data = loadmat('mat' + os.sep + 'FlatFingerprint1.mat')
Fingerprint = data['Fingerprint']
[M, N] = [Fingerprint.shape[0], Fingerprint.shape[1]]
up = int(M / 2 - B / 2)
down = int(M / 2 + B / 2)
left = int(N / 2 - B / 2)
right = int(N / 2 + B / 2)
def test_fastdvdnet(**args):
"""Denoises all sequences present in a given folder. Sequences must be stored as numbered
image sequences. The different sequences must be stored in subfolders under the "test_path" folder.
Inputs:
args (dict) fields:
"model_file": path to model
"test_path": path to sequence to denoise
"suffix": suffix to add to output name
"max_num_fr_per_seq": max number of frames to load per sequence
"dont_save_results: if True, don't save output images
"no_gpu": if True, run model on CPU
"save_path": where to save outputs as png
"gray": if True, perform denoising of grayscale images instead of RGB
"""
# Start time
start_time = time.time()
# If save_path does not exist, create it
if not os.path.exists(args['save_path']):
os.makedirs(args['save_path'])
logger = init_logger_test(args['save_path'])
# Sets data type according to CPU or GPU modes
if args['cuda']:
device = args['device_id'][0]
else:
device = torch.device('cpu')
# Create models
print('Loading models ...')
model_temp = FastDVDnet(num_input_frames=NUM_IN_FR_EXT)
# Load saved weights
state_temp_dict = torch.load(args['model_file'])
if args['cuda']:
device_ids = args['device_id']
model_temp = nn.DataParallel(model_temp,
device_ids=device_ids).cuda(device)
else:
# CPU mode: remove the DataParallel wrapper
state_temp_dict = remove_dataparallel_wrapper(state_temp_dict)
model_temp.load_state_dict(state_temp_dict)
# Sets the model in evaluation mode (e.g. it removes BN)
model_temp.eval()
gt = None
with torch.no_grad():
# process data
seq, _, _ = open_sequence(args['test_path'],
args['gray'],
expand_if_needed=False,
max_num_fr=args['max_num_fr_per_seq'])
seq = torch.from_numpy(seq).to(device)
seq_time = time.time()
denframes = denoise_seq_fastdvdnet(seq=seq,
temp_psz=NUM_IN_FR_EXT,
model_temporal=model_temp)
if args['gt_path'] is not None:
gt, _, _ = open_sequence(args['gt_path'],
args['gray'],
expand_if_needed=False,
max_num_fr=args['max_num_fr_per_seq'])
gt = torch.from_numpy(gt).to(device)
# Compute PSNR and log it
stop_time = time.time()
if gt is None:
psnr = 0
psnr_noisy = 0
else:
psnr = batch_psnr(denframes, gt, 1.)
psnr_noisy = batch_psnr(seq.squeeze(), gt, 1.)
loadtime = (seq_time - start_time)
runtime = (stop_time - seq_time)
seq_length = seq.size()[0]
logger.info("Finished denoising {}".format(args['test_path']))
logger.info(
"\tDenoised {} frames in {:.3f}s, loaded seq in {:.3f}s".format(
seq_length, runtime, loadtime))
logger.info("\tPSNR noisy {:.4f}dB, PSNR result {:.4f}dB".format(
psnr_noisy, psnr))
# Save outputs
if not args['dont_save_results']:
# Save sequence
save_out_seq(seq, denframes, args['save_path'], 0, args['suffix'],
args['save_noisy'])
# close logger
close_logger(logger)
def test_fastdvdnet(**args):
"""Denoises all sequences present in a given folder. Sequences must be stored as numbered
image sequences. The different sequences must be stored in subfolders under the "test_path" folder.
Inputs:
args (dict) fields:
"model_file": path to model
"test_path": path to sequence to denoise
"suffix": suffix to add to output name
"max_num_fr_per_seq": max number of frames to load per sequence
"noise_sigma": noise level used on test set
"dont_save_results: if True, don't save output images
"no_gpu": if True, run model on CPU
"save_path": where to save outputs as png
"gray": if True, perform denoising of grayscale images instead of RGB
"""
# If save_path does not exist, create it
if not os.path.exists(args['save_path']):
os.makedirs(args['save_path'])
logger = init_logger_test(args['save_path'])
# Sets data type according to CPU or GPU modes
if args['cuda']:
device = torch.device('cuda')
else:
device = torch.device('cpu')
# Create models
print('Loading models ...')
model_temp = FastDVDnet(num_input_frames=NUM_IN_FR_EXT)
# Load saved weights
state_temp_dict = torch.load(args['model_file'])
if args['cuda']:
device_ids = [0]
model_temp = nn.DataParallel(model_temp, device_ids=device_ids).cuda()
else:
# CPU mode: remove the DataParallel wrapper
state_temp_dict = remove_dataparallel_wrapper(state_temp_dict)
model_temp.load_state_dict(state_temp_dict)
# Sets the model in evaluation mode (e.g. it removes BN)
model_temp.eval()
processed_count = 0
# To avoid out of memory issues, we only process one folder at a time.
for tmp_folder in get_next_folder(args['test_path'], args['max_num_fr_per_seq']):
folder = tmp_folder.name
# Start time
print("Processing {}".format(os.listdir(tmp_folder.name)))
logger.info("Processing {}".format(os.listdir(folder)))
start_time = time.time()
with torch.no_grad():
# process data
seq, _, _ = open_sequence(folder,
args['gray'],
expand_if_needed=False,
max_num_fr=args['max_num_fr_per_seq'])
seq = torch.from_numpy(seq).to(device)
seq_time = time.time()
# Add noise
noise = torch.empty_like(seq).normal_(
mean=0, std=args['noise_sigma']).to(device)
seqn = seq + noise
noisestd = torch.FloatTensor([args['noise_sigma']]).to(device)
denframes = denoise_seq_fastdvdnet(seq=seqn,
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]
logger.info("Finished denoising {}".format(args['test_path']))
logger.info("\tDenoised {} frames in {:.3f}s, loaded seq in {:.3f}s".
format(seq_length, runtime, loadtime))
logger.info(
"\tPSNR noisy {:.4f}dB, PSNR result {:.4f}dB".format(psnr_noisy, psnr))
# Save outputs
if not args['dont_save_results']:
# Save sequence
save_out_seq(seqn, denframes, args['save_path'],
int(args['noise_sigma']*255), args['suffix'], args['save_noisy'], processed_count)
# subtract half stride because of the half-steps get_next_folder takes.
processed_count+=seqn.size()[0]
# close logger
close_logger(logger)