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 log_train_psnr(result, imsource, loss, writer, epoch, idx, num_minibatches, training_params):
'''Logs trai loss.
'''
#Compute pnsr of the whole batch
psnr_train = batch_psnr(torch.clamp(result, 0., 1.), imsource, 1.)
# Log the scalar values
writer.add_scalar('loss', loss.item(), training_params['step'])
# writer.add_scalar('PSNR on training data', psnr_train, \
# training_params['step'])
print("[epoch {}][{}/{}] loss: {:1.4f} PSNR_train: {:1.4f}".\
format(epoch+1, idx+1, num_minibatches, loss.item(), psnr_train))
def main(args):
r"""Performs the main training loop
"""
# Load dataset
print('> Loading dataset ...')
dataset_train = Dataset(train=True, gray_mode=args.gray, shuffle=True)
dataset_val = Dataset(train=False, gray_mode=args.gray, shuffle=False)
loader_train = DataLoader(dataset=dataset_train, num_workers=6, \
batch_size=args.batch_size, shuffle=True)
print("\t# of training samples: %d\n" % int(len(dataset_train)))
# Init loggers
if not os.path.exists(args.log_dir):
os.makedirs(args.log_dir)
writer = SummaryWriter(args.log_dir)
logger = init_logger(args)
# Create model
if not args.gray:
in_ch = 3
else:
in_ch = 1
net = FFDNet(num_input_channels=in_ch)
# Initialize model with He init
net.apply(weights_init_kaiming)
# Define loss
criterion = nn.MSELoss(size_average=False)
# Move to GPU
device_ids = [0]
model = nn.DataParallel(net, device_ids=device_ids).cuda()
criterion.cuda()
# Optimizer
optimizer = optim.Adam(model.parameters(), lr=args.lr)
# Resume training or start anew
if args.resume_training:
resumef = os.path.join(args.log_dir, 'ckpt.pth')
if os.path.isfile(resumef):
checkpoint = torch.load(resumef)
print("> Resuming previous training")
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
new_epoch = args.epochs
new_milestone = args.milestone
current_lr = args.lr
args = checkpoint['args']
training_params = checkpoint['training_params']
start_epoch = training_params['start_epoch']
args.epochs = new_epoch
args.milestone = new_milestone
args.lr = current_lr
print("=> loaded checkpoint '{}' (epoch {})"\
.format(resumef, start_epoch))
print("=> loaded parameters :")
print("==> checkpoint['optimizer']['param_groups']")
print("\t{}".format(checkpoint['optimizer']['param_groups']))
print("==> checkpoint['training_params']")
for k in checkpoint['training_params']:
print("\t{}, {}".format(k, checkpoint['training_params'][k]))
argpri = vars(checkpoint['args'])
print("==> checkpoint['args']")
for k in argpri:
print("\t{}, {}".format(k, argpri[k]))
args.resume_training = False
else:
raise Exception("Couldn't resume training with checkpoint {}".\
format(resumef))
else:
start_epoch = 0
training_params = {}
training_params['step'] = 0
training_params['current_lr'] = 0
training_params['no_orthog'] = args.no_orthog
# Training
for epoch in range(start_epoch, args.epochs):
# Learning rate value scheduling according to args.milestone
if epoch > args.milestone[1]:
current_lr = args.lr / 1000.
training_params['no_orthog'] = True
elif epoch > args.milestone[0]:
current_lr = args.lr / 10.
else:
current_lr = args.lr
# set learning rate in optimizer
for param_group in optimizer.param_groups:
param_group["lr"] = current_lr
print('learning rate %f' % current_lr)
# train
for i, data in enumerate(loader_train, 0):
# Pre-training step
model.train()
model.zero_grad()
optimizer.zero_grad()
# inputs: noise and noisy image
img_train = data
noise = torch.zeros(img_train.size())
stdn = np.random.uniform(args.noiseIntL[0], args.noiseIntL[1], \
size=noise.size()[0])
for nx in range(noise.size()[0]):
sizen = noise[0, :, :, :].size()
noise[nx, :, :, :] = torch.FloatTensor(sizen).\
normal_(mean=0, std=stdn[nx])
imgn_train = img_train + noise
# Create input Variables
img_train = Variable(img_train.cuda(), volatile=True)
imgn_train = Variable(imgn_train.cuda(), volatile=True)
noise = Variable(noise.cuda())
stdn_var = Variable(torch.cuda.FloatTensor(stdn), volatile=True)
# Evaluate model and optimize it
out_train = model(imgn_train, stdn_var)
loss = criterion(out_train, noise) / (imgn_train.size()[0] * 2)
loss.backward()
optimizer.step()
# Results
model.eval()
out_train = torch.clamp(imgn_train - model(imgn_train, stdn_var),
0., 1.)
psnr_train = batch_psnr(out_train, img_train, 1.)
# PyTorch v0.4.0: loss.data[0] --> loss.item()
if training_params['step'] % args.save_every == 0:
# Apply regularization by orthogonalizing filters
if not training_params['no_orthog']:
model.apply(svd_orthogonalization)
# Log the scalar values
writer.add_scalar('loss', loss.data[0],
training_params['step'])
writer.add_scalar('PSNR on training data', psnr_train, \
training_params['step'])
print("[epoch %d][%d/%d] loss: %.4f PSNR_train: %.4f" %\
(epoch+1, i+1, len(loader_train), loss.data[0], psnr_train))
training_params['step'] += 1
# The end of each epoch
model.eval()
# Validation
psnr_val = 0
for valimg in dataset_val:
img_val = torch.unsqueeze(valimg, 0)
noise = torch.FloatTensor(img_val.size()).\
normal_(mean=0, std=args.val_noiseL)
imgn_val = img_val + noise
img_val, imgn_val = Variable(img_val.cuda()), Variable(
imgn_val.cuda())
sigma_noise = Variable(torch.cuda.FloatTensor([args.val_noiseL]))
out_val = torch.clamp(imgn_val - model(imgn_val, sigma_noise), 0.,
1.)
psnr_val += batch_psnr(out_val, img_val, 1.)
psnr_val /= len(dataset_val)
print("\n[epoch %d] PSNR_val: %.4f" % (epoch + 1, psnr_val))
writer.add_scalar('PSNR on validation data', psnr_val, epoch)
writer.add_scalar('Learning rate', current_lr, epoch)
# Log val images
try:
if epoch == 0:
# Log graph of the model
writer.add_graph(
model,
(imgn_val, sigma_noise),
)
# Log validation images
for idx in range(2):
imclean = utils.make_grid(img_val.data[idx].clamp(0., 1.), \
nrow=2, normalize=False, scale_each=False)
imnsy = utils.make_grid(imgn_val.data[idx].clamp(0., 1.), \
nrow=2, normalize=False, scale_each=False)
writer.add_image('Clean validation image {}'.format(idx),
imclean, epoch)
writer.add_image('Noisy validation image {}'.format(idx),
imnsy, epoch)
for idx in range(2):
imrecons = utils.make_grid(out_val.data[idx].clamp(0., 1.), \
nrow=2, normalize=False, scale_each=False)
writer.add_image('Reconstructed validation image {}'.format(idx), \
imrecons, epoch)
# Log training images
imclean = utils.make_grid(img_train.data, nrow=8, normalize=True, \
scale_each=True)
writer.add_image('Training patches', imclean, epoch)
except Exception as e:
logger.error("Couldn't log results: {}".format(e))
# save model and checkpoint
training_params['start_epoch'] = epoch + 1
torch.save(model.state_dict(), os.path.join(args.log_dir, 'net.pth'))
save_dict = { \
'state_dict': model.state_dict(), \
'optimizer' : optimizer.state_dict(), \
'training_params': training_params, \
'args': args\
}
torch.save(save_dict, os.path.join(args.log_dir, 'ckpt.pth'))
if epoch % args.save_every_epochs == 0:
torch.save(save_dict, os.path.join(args.log_dir, \
'ckpt_e{}.pth'.format(epoch+1)))
del save_dict
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_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_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
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'], map_location=device)
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['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'])
# close logger
close_logger(logger)
def train(args: opts.TrainingOptions) -> None:
# Create the loader:
training_data = datsetprocess.AHDRDataset(
scene_directory=args.training_data)
loader = data.DataLoader(training_data,
batch_size=args.batch_size,
shuffle=True,
num_workers=1)
# Create the model:
model = models.AHDRNet().cuda()
criterion = nn.L1Loss().to(args.device.value)
optimizer = optim.Adam(model.parameters(), lr=args.learn_rate)
loss_list = []
current_epoch = 0
# Load pre-train model:
if os.path.exists(args.checkpoint):
checkpoint_file = os.path.realpath(args.checkpoint)
basename, _ = os.path.splitext(os.path.basename(checkpoint_file))
current_epoch = int(basename) + 1
state = torch.load(args.checkpoint)
loss_list = state['loss_list']
model.load_state_dict(state['model'])
# Train:
progress_bar = tqdm.tqdm(
iterable=range(current_epoch, args.max_epoch),
desc='STEP ?/? | LOSS ?.?????? | PSNR ?.????',
total=args.max_epoch,
initial=current_epoch,
)
for epoch in progress_bar:
losses = []
for step, sample in enumerate(loader):
(batch_x1, batch_x2, batch_x3, batch_x4) = (
sample['input1'],
sample['input2'],
sample['input3'],
sample['label'],
)
(batch_x1, batch_x2, batch_x3, batch_x4) = (
autograd.Variable(batch_x1).cuda(),
autograd.Variable(batch_x2).cuda(),
autograd.Variable(batch_x3).cuda(),
autograd.Variable(batch_x4).cuda(),
)
# Forward and compute loss:
pre = model(batch_x1, batch_x2, batch_x3)
loss = criterion(pre, batch_x4)
psnr = utils.batch_psnr(torch.clamp(pre, 0., 1.), batch_x4, 1.0)
losses.append(loss.item())
progress_bar.set_description(
'STEP {}/{} | LOSS {:0.6f} | PSNR {:0.6f}'.format(
step + 1, len(loader), losses[-1], psnr))
# Update the parameters:
optimizer.zero_grad()
loss.backward()
optimizer.step()
loss_list.append(np.mean(losses))
# Save the training model.
if epoch > 0 and epoch % args.checkpoint_interval == 0:
# Save progress:
save_file = '{:06d}.pkl'.format(epoch)
save_path = os.path.join(args.checkpoint_directory, save_file)
torch.save({
'model': model.state_dict(),
'loss_list': loss_list,
}, save_path)
# Create a symlink for easy resume:
latest = os.path.join(args.checkpoint_directory, 'latest.pkl')
if os.path.exists(latest):
os.unlink(latest)
os.symlink(save_file, latest)
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))
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))
# 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)
