From e6e62aae365b45e6beaf38fca9f95018a8408108 Mon Sep 17 00:00:00 2001
From: Piero Toffanin <pt@masseranolabs.com>
Date: Tue, 25 Jun 2019 11:04:54 -0400
Subject: [PATCH] Image size fallback fix, osfm extract camera

Former-commit-id: 7b8ca293fb473ecbc21007955e6f70c5c03a2fcb
---
 opendm/get_image_size.py |  2 +-
 opendm/osfm.py           | 39 +++++++++++++++++++++++++++++++++++----
 stages/run_opensfm.py    |  1 +
 3 files changed, 37 insertions(+), 5 deletions(-)

diff --git a/opendm/get_image_size.py b/opendm/get_image_size.py
index 0ff2eede..94b60a52 100644
--- a/opendm/get_image_size.py
+++ b/opendm/get_image_size.py
@@ -14,7 +14,7 @@ def get_image_size(file_path, fallback_on_error=True):
             width, height = img.size
     except Exception as e:
         if fallback_on_error:
-            log.warning("Cannot read %s with PIL, fallback to cv2: %s" % (file_path, str(e)))
+            log.ODM_WARNING("Cannot read %s with PIL, fallback to cv2: %s" % (file_path, str(e)))
             img = cv2.imread(file_path)
             width = img.shape[1]
             height = img.shape[0]
diff --git a/opendm/osfm.py b/opendm/osfm.py
index f06620d0..3f08c406 100644
--- a/opendm/osfm.py
+++ b/opendm/osfm.py
@@ -2,7 +2,7 @@
 OpenSfM related utils
 """
 
-import os, shutil, sys
+import os, shutil, sys, json
 import yaml
 from opendm import io
 from opendm import log
@@ -85,11 +85,12 @@ class OSFMContext:
                 log.ODM_DEBUG("Copied image_groups.txt to OpenSfM directory")
                 io.copy(image_groups_file, os.path.join(self.opensfm_project_path, "image_groups.txt"))
         
-            # check for camera_models.json
-            camera_models_file = os.path.join(args.project_path, "camera_models.json")
+            # check for cameras.json
+            # TODO: use config.cameras
+            camera_models_file = os.path.join(args.project_path, "cameras.json")
             has_camera_calibration = io.file_exists(camera_models_file)
             if has_camera_calibration:
-                log.ODM_DEBUG("Copied camera_models.json to OpenSfM directory (camera_models_overrides.json)")
+                log.ODM_DEBUG("Copied cameras.json to OpenSfM directory (camera_models_overrides.json)")
                 io.copy(camera_models_file, os.path.join(self.opensfm_project_path, "camera_models_overrides.json"))
 
             # create config file for OpenSfM
@@ -192,6 +193,36 @@ class OSFMContext:
     def path(self, *paths):
         return os.path.join(self.opensfm_project_path, *paths)
 
+    def extract_cameras(self, output, rerun=False):
+        reconstruction_file = self.path("reconstruction.json")
+        if not os.path.exists(output) or rerun:
+            if os.path.exists(reconstruction_file):
+                result = {}
+                with open(reconstruction_file, 'r') as fin:
+                    json_f = json.loads(fin.read())
+                    for recon in json_f:
+                        if 'cameras' in recon:
+                            for camera_id in recon['cameras']:
+                                # Strip "v2" from OpenSfM camera IDs
+                                new_camera_id = camera_id
+                                if new_camera_id.startswith("v2 "):
+                                    new_camera_id = new_camera_id[3:]
+
+                                result[new_camera_id] = recon['cameras'][camera_id]
+                                
+                                # Remove "_prior" keys
+                                keys = list(result[new_camera_id].keys())
+                                for k in keys:
+                                    if k.endswith('_prior'):
+                                        result[new_camera_id].pop(k)
+
+                with open(output, 'w') as fout:
+                    fout.write(json.dumps(result))
+            else:
+                log.ODM_WARNING("Cannot export cameras to %s. reconstruction.json does not exist." % output)
+        else:
+            log.ODM_INFO("Already extracted cameras")
+
     def update_config(self, cfg_dict):
         cfg_file = self.get_config_file_path()
         log.ODM_DEBUG("Updating %s" % cfg_file)
diff --git a/stages/run_opensfm.py b/stages/run_opensfm.py
index 78ccbab2..626e981d 100644
--- a/stages/run_opensfm.py
+++ b/stages/run_opensfm.py
@@ -27,6 +27,7 @@ class ODMOpenSfMStage(types.ODM_Stage):
         octx.feature_matching(self.rerun())
         self.update_progress(30)
         octx.reconstruct(self.rerun())
+        octx.extract_cameras(tree.path("cameras.json"), self.rerun())
         self.update_progress(70)
 
         # If we find a special flag file for split/merge we stop right here