Merge remote-tracking branch 'origin/compat_pycolmap'

kapture_localization/colmap/pycolmap_command.py

@@ -0,0 +1,101 @@
+from packaging import version
+try:
+    import pycolmap
+    version_number = pycolmap.__version__
+    if version.parse(version_number) < version.parse("0.5.0"):
+        compatibility_mode = True  # compatibility with older versions
+    else:
+        compatibility_mode = False
+    has_pycolmap = True
+except Exception as e:  # ModuleNotFoundError:
+    print(e)
+    has_pycolmap = False
+    compatibility_mode = False
+import numpy as np
+import quaternion
+
+
+def absolute_pose_estimation(points2D, points3D, camera_dict, max_error, min_inlier_ratio, min_num_iterations,
+                             max_num_iterations, confidence):
+    assert has_pycolmap
+    if compatibility_mode:
+        return pycolmap.absolute_pose_estimation(points2D, points3D, camera_dict, max_error,
+                                                 min_inlier_ratio, min_num_iterations, max_num_iterations, confidence)
+    else:
+        pycolmap_camera = pycolmap.Camera(
+            model=camera_dict['model'], width=camera_dict['width'], height=camera_dict['height'],
+            params=camera_dict['params'])
+
+        pycolmap_estimation_options = dict(ransac=dict(max_error=max_error, min_inlier_ratio=min_inlier_ratio,
+                                           min_num_trials=min_num_iterations, max_num_trials=max_num_iterations,
+                                           confidence=confidence))
+        ret = pycolmap.absolute_pose_estimation(points2D, points3D, pycolmap_camera,
+                                                estimation_options=pycolmap_estimation_options)
+        if ret is None:
+            ret = {'success': False}
+        else:
+            ret['success'] = True
+            retmat = ret['cam_from_world'].matrix
+            ret['qvec'] = quaternion.from_rotation_matrix(retmat[:3, :3])
+            ret['tvec'] = retmat[:3, 3]
+        return ret
+
+
+def rig_absolute_pose_estimation(points2D, points3D, cameras_dict, qvec, tvec, max_error,
+                                 min_inlier_ratio, min_num_iterations, max_num_iterations,
+                                 confidence):
+    assert has_pycolmap
+    if compatibility_mode:
+        return pycolmap.rig_absolute_pose_estimation(points2D, points3D, cameras_dict, qvec, tvec, max_error,
+                                                     min_inlier_ratio, min_num_iterations, max_num_iterations,
+                                                     confidence)
+    else:
+        pycolmap_cameras = []
+        camera_idxs = []
+        cams_from_rig = []
+        for idx, (camera_dict, qvec_idx, tvec_idx) in enumerate(zip(cameras_dict, qvec, tvec)):
+            pycolmap_cameras.append(pycolmap.Camera(
+                                    model=camera_dict['model'], width=camera_dict['width'],
+                                    height=camera_dict['height'], params=camera_dict['params']))
+            camera_idxs.append(idx)
+            cam_from_rig = np.eye(4)
+            cam_from_rig[:3, :3] = quaternion.as_rotation_matrix(quaternion.from_float_array(qvec_idx))
+            cam_from_rig[:3, 3] = tvec_idx
+            cams_from_rig.append(pycolmap.Rigid3d(cam_from_rig[:3, :]))
+        pycolmap_estimation_options = dict(ransac=dict(max_error=max_error, min_inlier_ratio=min_inlier_ratio,
+                                           min_num_trials=min_num_iterations, max_num_trials=max_num_iterations,
+                                           confidence=confidence))
+        ret = pycolmap.absolute_pose_estimation(points2D, points3D, pycolmap_cameras, camera_idxs, cams_from_rig,
+                                                estimation_options=pycolmap_estimation_options)
+        if ret is None:
+            ret = {'success': False}
+        else:
+            ret['success'] = True
+            retmat = ret['rig_from_world'].matrix
+            ret['qvec'] = quaternion.from_rotation_matrix(retmat[:3, :3])
+            ret['tvec'] = retmat[:3, 3]
+        return ret
+
+
+def pose_refinement(tvec, qvec, points2D, points3D, inlier_mask, camera_dict):
+    assert has_pycolmap
+    if compatibility_mode:
+        return pycolmap.pose_refinement(tvec, qvec, points2D, points3D, inlier_mask, camera_dict)
+    else:
+        # pycolmap.pose_refinement(cam_from_world, points2D, points3D, inlier_mask, camera)
+        pycolmap_camera = pycolmap.Camera(model=camera_dict['model'],
+                                          width=camera_dict['width'], height=camera_dict['height'],
+                                          params=camera_dict['params'])
+        cam_from_world = np.eye(4)
+        cam_from_world[:3, :3] = quaternion.as_rotation_matrix(quaternion.from_float_array(qvec))
+        cam_from_world[:3, 3] = tvec
+        cam_from_world_pycolmap = pycolmap.Rigid3d(cam_from_world[:3, :])
+        ret = pycolmap.pose_refinement(cam_from_world_pycolmap, points2D, points3D, inlier_mask, pycolmap_camera)
+        if ret is None:
+            ret = {'success': False}
+        else:
+            ret['success'] = True
+            retmat = ret['cam_from_world'].matrix
+            ret['qvec'] = quaternion.from_rotation_matrix(retmat[:3, :3])
+            ret['tvec'] = retmat[:3, 3]
+        return ret

tools/kapture_pycolmap_localize.py

@@ -11,12 +11,6 @@
 from typing import Optional
 from tqdm import tqdm
 import numpy as np
-try:
-    import pycolmap
-    has_pycolmap = True
-except ModuleNotFoundError:
-    has_pycolmap = False
-
 
 import path_to_kapture_localization  # noqa: F401
 from kapture_localization.utils.logging import getLogger, save_to_json
@@ -25,6 +19,7 @@
 from kapture_localization.localization.RerankCorrespondencesStrategy import RerankCorrespondencesStrategy
 from kapture_localization.localization.reprojection_error import compute_reprojection_error
 from kapture_localization.utils.cv_camera_matrix import get_camera_matrix_from_kapture
+from kapture_localization.colmap.pycolmap_command import absolute_pose_estimation, has_pycolmap
 
 import kapture_localization.utils.path_to_kapture  # noqa: F401
 import kapture
@@ -231,8 +226,8 @@ def pycolmap_localize_from_loaded_data(kapture_data: kapture.Kapture,
         # compute absolute pose
         # inlier_threshold - RANSAC inlier threshold in pixels
         # answer - dictionary containing the RANSAC output
-        ret = pycolmap.absolute_pose_estimation(points2D, points3D, cfg, max_error,
-                                                min_inlier_ratio, min_num_iterations, max_num_iterations, confidence)
+        ret = absolute_pose_estimation(points2D, points3D, cfg, max_error,
+                                       min_inlier_ratio, min_num_iterations, max_num_iterations, confidence)
         # add pose to output kapture
         if ret['success'] and ret['num_inliers'] > 0:
             pose = kapture.PoseTransform(ret['qvec'], ret['tvec'])

tools/kapture_pycolmap_localsfm.py

@@ -10,11 +10,6 @@
 import os
 from tqdm import tqdm
 from typing import Optional
-try:
-    import pycolmap
-    has_pycolmap = True
-except ModuleNotFoundError:
-    has_pycolmap = False
 
 import datetime
 import numpy as np
@@ -26,6 +21,7 @@
 from kapture_localization.utils.cv_camera_matrix import get_camera_matrix_from_kapture
 from kapture_localization.triangulation.integration import aggregate_matches, triangulate_all_points
 from kapture_localization.triangulation.integration import convert_correspondences
+from kapture_localization.colmap.pycolmap_command import absolute_pose_estimation, has_pycolmap
 
 import kapture_localization.utils.path_to_kapture  # noqa: F401
 import kapture
@@ -253,9 +249,9 @@ def kapture_pycolmap_localsfm_from_loaded_data(kapture_data: kapture.Kapture,
             # compute absolute pose
             # inlier_threshold - RANSAC inlier threshold in pixels
             # answer - dictionary containing the RANSAC output
-            ret = pycolmap.absolute_pose_estimation(points2D, points3D, cfg, max_error,
-                                                    min_inlier_ratio, min_num_iterations,
-                                                    max_num_iterations, confidence)
+            ret = absolute_pose_estimation(points2D, points3D, cfg, max_error,
+                                           min_inlier_ratio, min_num_iterations,
+                                           max_num_iterations, confidence)
 
             # add pose to output kapture
             if ret['success'] and ret['num_inliers'] > 0:

tools/kapture_pycolmap_rig_localize.py

@@ -11,11 +11,6 @@
 from typing import Optional, List
 import numpy as np
 from tqdm import tqdm
-try:
-    import pycolmap
-    has_pycolmap = True
-except ModuleNotFoundError:
-    has_pycolmap = False
 
 import path_to_kapture_localization  # noqa: F401
 from kapture_localization.utils.logging import getLogger, save_to_json
@@ -25,6 +20,7 @@
 from kapture_localization.localization.reprojection_error import compute_reprojection_error
 from kapture_localization.utils.cv_camera_matrix import get_camera_matrix_from_kapture
 from kapture_localization.utils.rigs_extension import get_top_level_rig_ids, get_all_cameras_from_rig_ids
+from kapture_localization.colmap.pycolmap_command import rig_absolute_pose_estimation, has_pycolmap
 
 import kapture_localization.utils.path_to_kapture  # noqa: F401
 import kapture
@@ -265,9 +261,9 @@ def pycolmap_rig_localize_from_loaded_data(kapture_data: kapture.Kapture,
             # compute absolute pose
             # inlier_threshold - RANSAC inlier threshold in pixels
             # answer - dictionary containing the RANSAC output
-            ret = pycolmap.rig_absolute_pose_estimation(points2D, points3D, cameras_dict, qvec, tvec, max_error,
-                                                        min_inlier_ratio, min_num_iterations, max_num_iterations,
-                                                        confidence)
+            ret = rig_absolute_pose_estimation(points2D, points3D, cameras_dict, qvec, tvec, max_error,
+                                               min_inlier_ratio, min_num_iterations, max_num_iterations,
+                                               confidence)
 
             # add pose to output kapture
             if ret['success'] and ret['num_inliers'] > 0:

tools/kapture_pycolmap_rig_localsfm.py

@@ -10,11 +10,6 @@
 import os
 from tqdm import tqdm
 from typing import Optional, List
-try:
-    import pycolmap
-    has_pycolmap = True
-except ModuleNotFoundError:
-    has_pycolmap = False
 
 import datetime
 import numpy as np
@@ -27,6 +22,7 @@
 from kapture_localization.utils.rigs_extension import get_top_level_rig_ids, get_all_cameras_from_rig_ids
 from kapture_localization.triangulation.integration import aggregate_matches, triangulate_all_points
 from kapture_localization.triangulation.integration import convert_correspondences
+from kapture_localization.colmap.pycolmap_command import rig_absolute_pose_estimation, has_pycolmap
 
 import kapture_localization.utils.path_to_kapture  # noqa: F401
 import kapture
@@ -296,9 +292,9 @@ def kapture_pycolmap_rig_localsfm_from_loaded_data(
                 # compute absolute pose
                 # inlier_threshold - RANSAC inlier threshold in pixels
                 # answer - dictionary containing the RANSAC output
-                ret = pycolmap.rig_absolute_pose_estimation(points2D, points3D, cameras_dict, qvec, tvec, max_error,
-                                                            min_inlier_ratio, min_num_iterations, max_num_iterations,
-                                                            confidence)
+                ret = rig_absolute_pose_estimation(points2D, points3D, cameras_dict, qvec, tvec, max_error,
+                                                   min_inlier_ratio, min_num_iterations, max_num_iterations,
+                                                   confidence)
 
                 # add pose to output kapture
                 if ret['success'] and ret['num_inliers'] > 0:

tools/kapture_pyransaclib_localsfm.py

@@ -15,11 +15,6 @@
     has_pyransaclib = True
 except ModuleNotFoundError:
     has_pyransaclib = False
-try:
-    import pycolmap
-    has_pycolmap = True
-except ModuleNotFoundError:
-    has_pycolmap = False
 
 import datetime
 import cv2
@@ -32,6 +27,7 @@
 from kapture_localization.utils.cv_camera_matrix import get_camera_matrix_from_kapture
 from kapture_localization.triangulation.integration import aggregate_matches, triangulate_all_points
 from kapture_localization.triangulation.integration import convert_correspondences
+from kapture_localization.colmap.pycolmap_command import pose_refinement, has_pycolmap
 
 import kapture_localization.utils.path_to_kapture  # noqa: F401
 import kapture
@@ -296,7 +292,7 @@ def kapture_pyransaclib_localsfm_from_loaded_data(kapture_data: kapture.Kapture,
                         'height': int(height),
                         'params': params
                     }
-                    ret_refine = pycolmap.pose_refinement(pose.t_raw, pose.r_raw, points2D, points3D, inlier_mask, cfg)
+                    ret_refine = pose_refinement(pose.t_raw, pose.r_raw, points2D, points3D, inlier_mask, cfg)
                     if ret_refine['success']:
                         pose = kapture.PoseTransform(ret_refine['qvec'], ret_refine['tvec'])
                         logger.debug(f'{image_name} refinement success, new pose: {pose}')