def main(args, config):
cuda = cnn_utils.check_cuda(config)
model = cnn_utils.load_model_and_weights(args, config)
if cuda:
model = model.cuda()
file_path = os.path.join(config['PATH']['hdf5_dir'],
config['PATH']['hdf5_name'])
with h5py.File(file_path, mode='r', libver='latest') as hdf5_file:
overall_psnr_accum = (0, 0, 0)
overall_ssim_accum = (0, 0, 0)
for sample_num in range(args.nSamples):
p1, s1 = do_one_demo(args, config, hdf5_file, model, sample_num,
cuda)
overall_psnr_accum = welford.update(overall_psnr_accum, p1)
overall_ssim_accum = welford.update(overall_ssim_accum, s1)
if args.nSamples > 1:
psnr_mean, psnr_var, _ = welford.finalize(overall_psnr_accum)
ssim_mean, ssim_var, _ = welford.finalize(overall_ssim_accum)
print("\nOverall cnn psnr average {:5f}, stddev {:5f}".format(
psnr_mean, math.sqrt(psnr_var)))
print("Overall cnn ssim average {:5f}, stddev {:5f}".format(
ssim_mean, math.sqrt(ssim_var)))
#Ground truth possible
"""
def main(args, config, sample_index):
cuda = cnn_utils.check_cuda(config)
model = cnn_utils.load_model_and_weights(args, config)
if cuda:
model = model.cuda()
model.eval()
# Create output directory
base_dir = os.path.join(config['PATH']['output_dir'], 'warped')
if not os.path.isdir(base_dir):
pathlib.Path(base_dir).mkdir(parents=True, exist_ok=True)
save_dir = get_sub_dir_for_saving(base_dir)
start_time = time.time()
file_path = os.path.join(config['PATH']['hdf5_dir'],
config['PATH']['hdf5_name'])
with h5py.File(file_path, mode='r', libver='latest') as hdf5_file:
depth_grp = hdf5_file['val']['disparity']
SNUM = sample_index
depth_images = torch.squeeze(
torch.tensor(depth_grp['images'][SNUM], dtype=torch.float32))
colour_grp = hdf5_file['val']['colour']
colour_images = torch.tensor(colour_grp['images'][SNUM],
dtype=torch.float32)
sample = {
'depth': depth_images,
'colour': colour_images,
'grid_size': depth_images.shape[0]
}
warped = data_transform.center_normalise(sample)
im_input = warped['inputs'].unsqueeze_(0)
if cuda:
im_input = im_input.cuda()
output = model(im_input)
time_taken = time.time() - start_time
print("Time taken was {:4f}s".format(time_taken))
grid_size = 64
psnr_accumulator = (0, 0, 0)
ssim_accumulator = (0, 0, 0)
print("Saving output to", save_dir)
output = torch.squeeze(denormalise_lf(output))
cpu_output = np.around(output.cpu().detach().numpy()).astype(np.uint8)
if (not args.no_eval) or args.get_diff:
ground_truth = np.around(
denormalise_lf(colour_images).numpy()).astype(np.uint8)
grid_len = int(math.sqrt(grid_size))
for i in range(grid_size):
row, col = i // grid_len, i % grid_len
file_name = 'Colour{}{}.png'.format(row, col)
save_location = os.path.join(save_dir, file_name)
if i == 0:
print("Saving images of size ", cpu_output[i].shape)
image_warping.save_array_as_image(cpu_output[i], save_location)
if args.get_diff:
colour = ground_truth[i]
diff = image_warping.get_diff_image(colour, cpu_output[i])
#diff = get_diff_image_floatint(res, colour)
file_name = 'Diff{}{}.png'.format(row, col)
save_location = os.path.join(save_dir, file_name)
image_warping.save_array_as_image(diff, save_location)
if not args.no_eval:
img = ground_truth[i]
file_name = 'GT_Colour{}{}.png'.format(row, col)
save_location = os.path.join(save_dir, file_name)
image_warping.save_array_as_image(img, save_location)
psnr = evaluate.my_psnr(cpu_output[i], img)
ssim = evaluate.ssim(cpu_output[i], img)
psnr_accumulator = welford.update(psnr_accumulator, psnr)
ssim_accumulator = welford.update(ssim_accumulator, ssim)
psnr_mean, psnr_var, _ = welford.finalize(psnr_accumulator)
ssim_mean, ssim_var, _ = welford.finalize(ssim_accumulator)
print("For cnn, psnr average {:5f}, stddev {:5f}".format(
psnr_mean, math.sqrt(psnr_var)))
print("For cnn, ssim average {:5f}, stddev {:5f}".format(
ssim_mean, math.sqrt(ssim_var)))
def main(args, config, writer):
best_loss = math.inf
best_model, best_epoch = None, None
cuda = cnn_utils.check_cuda(config)
# Attempts to otimise - see
# https://discuss.pytorch.org/t/what-does-torch-backends-cudnn-benchmark-do
torch.backends.cudnn.benchmark = True
data_loaders = create_dataloaders(args, config)
model, criterion, optimizer, lr_scheduler = setup_model(args)
if cuda: # GPU support
model = model.cuda()
# The below is only needed if loss fn has params
#criterion = criterion.cuda()
if args.checkpoint: # Resume from a checkpoint
best_loss = cnn_utils.load_from_checkpoint(model, optimizer, args,
config)
if args.pretrained: # Direct copy weights from another model
cnn_utils.load_weights(model, args, config, frozen=args.frozen)
# Perform training and testing
print("Beginning training loop")
for epoch in range(args.start_epoch, args.start_epoch + args.nEpochs):
epoch_loss = 0
epoch_loss_all = train(model=model,
dset_loaders=data_loaders,
optimizer=optimizer,
lr_scheduler=lr_scheduler,
criterion=criterion,
epoch=epoch,
cuda=cuda,
clip=args.clip,
writer=writer)
for val in epoch_loss_all:
epoch_loss += val / len(epoch_loss_all)
if epoch_loss < best_loss:
best_loss = epoch_loss
best_epoch = epoch
best_model = copy.deepcopy(model)
# Update the scheduler - restarting
if lr_scheduler.last_epoch == lr_scheduler.T_max:
for group in optimizer.param_groups:
group['lr'] = args.lr
lr_scheduler = CosineAnnealingLR(optimizer,
T_max=lr_scheduler.T_max * 2)
# cnn_utils.log_all_layer_weights(model, writer, epoch)
if epoch % 1 == 0 and epoch != 0:
cnn_utils.save_checkpoint(model, epoch, optimizer, best_loss,
config['PATH']['checkpoint_dir'],
args.tag + "{}.pth".format(epoch))
if args.prompt:
if not helpers.prompt_user(CONTINUE_MESSAGE):
print("Ending training")
break
print("Best loss was {:.5f} at epoch {}".format(best_loss, best_epoch))
save = True
if args.prompt:
if not helpers.prompt_user(SAVE_MESSAGE):
print("Not saving the model")
save = False
# Save the best model
if save:
cnn_utils.save_checkpoint(
best_model, best_epoch, optimizer, best_loss,
config['PATH']['model_dir'],
args.tag + "_best_at{}.pth".format(best_epoch))
parent_dir = os.path.abspath(os.pardir)
scalar_dir = os.path.join(parent_dir, "logs", args.tag)
# if not os.path.isdir(scalar_dir):
# pathlib.Path(scalar_dir).mkdir(parents=True, exist_ok=True)
# writer.export_scalars_to_json(
# os.path.join(scalar_dir, "all_scalars.json"))
writer.close()
