From 00f9d3ee80ab37a313c4707d30b0b9d228b43eed Mon Sep 17 00:00:00 2001
From: Lorenzo Bertoni <34957815+bertoni9@users.noreply.github.com>
Date: Tue, 9 Feb 2021 12:12:44 +0100
Subject: [PATCH] add focal length argument (#6)

---
 monstereo/network/process.py | 7 +++----
 monstereo/predict.py         | 2 +-
 monstereo/run.py             | 3 +++
 3 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/monstereo/network/process.py b/monstereo/network/process.py
index a79e6b9..5cb83c3 100644
--- a/monstereo/network/process.py
+++ b/monstereo/network/process.py
@@ -17,7 +17,6 @@ z_min = 4
 z_max = 60
 D_MIN = BF / z_max
 D_MAX = BF / z_min
-FL = 5.7  # nuScenes focal length (mm)
 Sx = 7.2  # nuScenes sensor size (mm)
 Sy = 5.4  # nuScenes sensor size (mm)
 
@@ -67,7 +66,7 @@ def preprocess_monoloco(keypoints, kk, zero_center=False):
     return kps_out
 
 
-def factory_for_gt(im_size, name=None, path_gt=None, verbose=True):
+def factory_for_gt(im_size, focal_length=5.7, name=None, path_gt=None, verbose=True):
     """Look for ground-truth annotations file and define calibration matrix based on image size """
 
     try:
@@ -91,8 +90,8 @@ def factory_for_gt(im_size, name=None, path_gt=None, verbose=True):
             kk = [[718.3351, 0., 600.3891], [0., 718.3351, 181.5122], [0., 0., 1.]]  # Kitti calibration
         else:  # nuScenes camera parameters
             kk = [
-                [im_size[0]*FL/Sx, 0., im_size[0]/2],
-                [0., im_size[1]*FL/Sy, im_size[1]/2],
+                [im_size[0]*focal_length/Sx, 0., im_size[0]/2],
+                [0., im_size[1]*focal_length/Sy, im_size[1]/2],
                 [0., 0., 1.]]
         if verbose:
             logger.info("Using a standard calibration matrix...")
diff --git a/monstereo/predict.py b/monstereo/predict.py
index 5b4f443..ba0f3c3 100644
--- a/monstereo/predict.py
+++ b/monstereo/predict.py
@@ -130,7 +130,7 @@ def predict(args):
 
             im_name = os.path.basename(meta['file_name'])
             im_size = (cpu_image.size[0], cpu_image.size[1])  # Original
-            kk, dic_gt = factory_for_gt(im_size, name=im_name, path_gt=args.path_gt)
+            kk, dic_gt = factory_for_gt(im_size, focal_length=args.focal, name=im_name, path_gt=args.path_gt)
 
             # Preprocess pifpaf outputs and run monoloco
             boxes, keypoints = preprocess_pifpaf(pifpaf_outs['left'], im_size, enlarge_boxes=False)
diff --git a/monstereo/run.py b/monstereo/run.py
index 8957c33..c9c2bae 100644
--- a/monstereo/run.py
+++ b/monstereo/run.py
@@ -39,6 +39,9 @@ def cli():
     predict_parser.add_argument('--dpi', help='image resolution',  type=int, default=150)
     predict_parser.add_argument('--long-edge', default=None, type=int,
                                 help='rescale the long side of the image (aspect ratio maintained)')
+    predict_parser.add_argument('--focal',
+                                help='focal length in mm for a sensor size of 7.2x5.4 mm. Default nuScenes sensor',
+                                type=float, default=5.7)
 
     # Pifpaf parsers
     decoder.cli(predict_parser)