def validate_and_log(model_temp, dataset_val, valnoisestd, temp_psz, writer, \
epoch, lr, logger, trainimg):
"""Validation step after the epoch finished
"""
t1 = time.time()
psnr_val = 0
with torch.no_grad():
for seq_val in dataset_val:
noise = torch.FloatTensor(seq_val.size()).normal_(mean=0,
std=valnoisestd)
seqn_val = seq_val + noise
seqn_val = seqn_val.cuda()
sigma_noise = torch.cuda.FloatTensor([valnoisestd])
out_val = denoise_seq_fastdvdnet(seq=seqn_val, \
noise_std=sigma_noise, \
temp_psz=temp_psz,\
model_temporal=model_temp)
psnr_val += batch_psnr(out_val.cpu(), seq_val.squeeze_(), 1.)
psnr_val /= len(dataset_val)
t2 = time.time()
print("\n[epoch %d] PSNR_val: %.4f, on %.2f sec" %
(epoch + 1, psnr_val, (t2 - t1)))
writer.add_scalar('PSNR on validation data', psnr_val, epoch)
writer.add_scalar('Learning rate', lr, epoch)
# Log val images
try:
idx = 0
if epoch == 0:
# Log training images
_, _, Ht, Wt = trainimg.size()
img = tutils.make_grid(trainimg.view(-1, 3, Ht, Wt), \
nrow=8, normalize=True, scale_each=True)
writer.add_image('Training patches', img, epoch)
# Log validation images
img = tutils.make_grid(seq_val.data[idx].clamp(0., 1.),\
nrow=2, normalize=False, scale_each=False)
imgn = tutils.make_grid(seqn_val.data[idx].clamp(0., 1.),\
nrow=2, normalize=False, scale_each=False)
writer.add_image('Clean validation image {}'.format(idx), img,
epoch)
writer.add_image('Noisy validation image {}'.format(idx), imgn,
epoch)
# Log validation results
irecon = tutils.make_grid(out_val.data[idx].clamp(0., 1.),\
nrow=2, normalize=False, scale_each=False)
writer.add_image('Reconstructed validation image {}'.format(idx),
irecon, epoch)
except Exception as e:
logger.error(
"validate_and_log_temporal(): Couldn't log results, {}".format(e))
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_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)
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)
# return seq, expanded_h, expanded_w
if len(seq_list) == NUM_IN_FR_EXT:
print("Infer batch ...")
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=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))