From 55870cb559f41e629b5323d8591a06773025c318 Mon Sep 17 00:00:00 2001 From: Michael Firman Date: Mon, 17 May 2021 21:43:44 +0100 Subject: [PATCH] Added option to test_simple.py to directly predict depth. Also dataloader comments. --- datasets/kitti_dataset.py | 6 +++++- test_simple.py | 25 ++++++++++++++++++++----- 2 files changed, 25 insertions(+), 6 deletions(-) diff --git a/datasets/kitti_dataset.py b/datasets/kitti_dataset.py index 81ef28823..120ca35fc 100644 --- a/datasets/kitti_dataset.py +++ b/datasets/kitti_dataset.py @@ -21,7 +21,11 @@ class KITTIDataset(MonoDataset): def __init__(self, *args, **kwargs): super(KITTIDataset, self).__init__(*args, **kwargs) - # NOTE: Make sure your intrinsics matrix is *normalized* by the original image size + # NOTE: Make sure your intrinsics matrix is *normalized* by the original image size. + # To normalize you need to scale the first row by 1 / image_width and the second row + # by 1 / image_height. Monodepth2 assumes a principal point to be exactly centered. + # If your principal point is far from the center you might need to disable the horizontal + # flip augmentation. self.K = np.array([[0.58, 0, 0.5, 0], [0, 1.92, 0.5, 0], [0, 0, 1, 0], diff --git a/test_simple.py b/test_simple.py index f4dd77435..6489e1b68 100644 --- a/test_simple.py +++ b/test_simple.py @@ -21,6 +21,7 @@ import networks from layers import disp_to_depth from utils import download_model_if_doesnt_exist +from evaluate_depth import STEREO_SCALE_FACTOR def parse_args(): @@ -46,6 +47,10 @@ def parse_args(): parser.add_argument("--no_cuda", help='if set, disables CUDA', action='store_true') + parser.add_argument("--pred_depth", + help='if set, predicts metric depth instead of disparity. (This only ' + 'makes sense for stereo-trained models).', + action='store_true') return parser.parse_args() @@ -61,6 +66,10 @@ def test_simple(args): else: device = torch.device("cpu") + if args.pred_depth and "stereo" not in args.model_name: + print("Warning: The --pred_depth flag only makes sense for stereo-trained models. " + "For mono-trained models, output depths will not in metric space.") + download_model_if_doesnt_exist(args.model_name) model_path = os.path.join("models", args.model_name) print("-> Loading model from ", model_path) @@ -129,9 +138,13 @@ def test_simple(args): # Saving numpy file output_name = os.path.splitext(os.path.basename(image_path))[0] - name_dest_npy = os.path.join(output_directory, "{}_disp.npy".format(output_name)) - scaled_disp, _ = disp_to_depth(disp, 0.1, 100) - np.save(name_dest_npy, scaled_disp.cpu().numpy()) + scaled_disp, depth = disp_to_depth(disp, 0.1, 100) + if args.pred_depth: + name_dest_npy = os.path.join(output_directory, "{}_depth.npy".format(output_name)) + np.save(name_dest_npy, STEREO_SCALE_FACTOR * depth.cpu().numpy()) + else: + name_dest_npy = os.path.join(output_directory, "{}_disp.npy".format(output_name)) + np.save(name_dest_npy, scaled_disp.cpu().numpy()) # Saving colormapped depth image disp_resized_np = disp_resized.squeeze().cpu().numpy() @@ -144,8 +157,10 @@ def test_simple(args): name_dest_im = os.path.join(output_directory, "{}_disp.jpeg".format(output_name)) im.save(name_dest_im) - print(" Processed {:d} of {:d} images - saved prediction to {}".format( - idx + 1, len(paths), name_dest_im)) + print(" Processed {:d} of {:d} images - saved predictions to:".format( + idx + 1, len(paths))) + print(" - {}".format(name_dest_im)) + print(" - {}".format(name_dest_npy)) print('-> Done!')