[WIP] Bring some more methods from upstream and other minor fixes.

lamar/tasks/feature_extraction.py

@@ -147,4 +147,25 @@ class RetrievalFeatureExtraction(FeatureExtraction):
                 'preprocessing': {'resize_max': 640},
             }
         },
+        'cosplace': {
+            'name': 'cosplace',
+            'hloc': {
+                'model': {'name': 'cosplace'},
+                'preprocessing': {'resize_max': 640},
+            }
+        },
+        'openibl': {
+            'name': 'openibl',
+            'hloc': {
+                'model': {'name': 'openibl'},
+                'preprocessing': {'resize_max': 640},
+            }
+        },
+        'salad': {
+            'name': 'salad',
+            'hloc': {
+                'model': {'name': 'salad'},
+                'preprocessing': {'resize_max': 640},
+            }
+        }
     }

lamar/tasks/feature_matching.py

@@ -33,6 +33,15 @@ class FeatureMatching:
                 },
             },
         },
+        'lightglue': {
+            'name': 'lightglue',
+            'hloc': {
+                'model': {
+                    'name': 'lightglue',
+                    'features': 'superpoint',
+                },
+            },
+        },
         'mnn': {
             'name': 'mnn',
             'hloc': {

lamar/tasks/mapping.py

@@ -1,6 +1,7 @@
 import logging
 from pathlib import Path
 from copy import deepcopy
+from threading import Lock
 import numpy as np
 
 import pycolmap
@@ -94,6 +95,19 @@ def __init__(self, config, outputs, capture, session_id,
             self.name2key[image.name]: image.image_id
             for image in self.reconstruction.images.values()
         }
+        # We cache the 3D points for mapping images to avoid parsing them from
+        # the reconstruction each time. If we count 5K keypoints per image, this
+        # leads to ~160KB / image, so even for 25K mapping images, this would
+        # only be ~4GB of RAM. For very large datasets, we might want to look
+        # into a LRU cache.
+        self.points3d_cache = {}
+        # We use a lock per image to make sure other threads using different 
+        # images can run in parallel. These locks block for the expensive
+        # reconstruction parsing.
+        self.image_locks = {}
+        # We use a lock for the image locks to make sure that we do not
+        # create multiple locks for the same image.
+        self.lock = Lock()
 
     def run(self, capture):
         run_capture_to_empty_colmap.run(capture, [self.session_id], self.paths.sfm_empty)
@@ -106,20 +120,38 @@ def run(self, capture):
             self.matching.paths.matches,
         )
 
-    def get_points3D(self, key, point2D_indices):
-        image = self.reconstruction.images[self.key2imageid[key]]
-        valid = []
-        xyz = []
-        ids = []
-        if len(image.points2D) > 0:
-            for idx in point2D_indices:
-                p = image.points2D[idx]
-                valid.append(p.has_point3D())
-                if valid[-1]:
-                    ids.append(p.point3D_id)
-                    xyz.append(self.reconstruction.points3D[ids[-1]].xyz)
-        return np.array(valid, bool), xyz, ids
+    def _get_points3D_from_cache(self, key, point2D_indices):
+        ids, xyz = self.points3d_cache[key]
+        if len(ids) == 0:
+            # Not registered.
+            return np.array([], bool), [], []
+        valid = ids[point2D_indices] != -1
+        return valid, xyz[point2D_indices][valid], ids[point2D_indices][valid]
+
 
+    def get_points3D(self, key, point2D_indices):
+        if key not in self.image_locks:
+            with self.lock:
+                # Key might have been added while we were waiting for the lock.
+                if key not in self.image_locks:
+                    self.image_locks[key] = Lock()
+        if key not in self.points3d_cache:
+            with self.image_locks[key]:              
+                # Key might have been added while we were waiting for the lock.
+                if key not in self.points3d_cache:
+                    image = self.reconstruction.images[self.key2imageid[key]]
+                    ids = []
+                    xyz = []
+                    for p2d in image.points2D:
+                        if p2d.has_point3D():
+                            ids.append(p2d.point3D_id)
+                            xyz.append(self.reconstruction.points3D[ids[-1]].xyz)
+                        else:
+                            ids.append(-1)
+                            xyz.append([np.nan, np.nan, np.nan])
+                    self.points3d_cache[key] = (np.array(ids), np.array(xyz))
+        return self._get_points3D_from_cache(key, point2D_indices)
+
 
 class MeshLifting(Mapping):
     method = {

lamar/tasks/pose_estimation.py

@@ -51,7 +51,7 @@ class PoseEstimation:
     method2class = {}
     method = None
     evaluation = {
-        'Rt_thresholds': [(1, 0.1), (5, 1.)],
+        'Rt_thresholds': [(1, .1), (5, .5), (5, 1), (10, .25), (10, .5), (10, 1), (10, 2.5), (10, 5), (10, 10)],
     }
 
     def __init_subclass__(cls):
@@ -71,7 +71,7 @@ def __init__(self, config, outputs, capture, query_id,
                  matching: FeatureMatching,
                  mapping: Mapping,
                  query_keys: list = None,
-                 parallel: bool = True,
+                 parallel: bool = False,
                  return_covariance: bool = False,
                  override_workdir_root: Path = None):
         if extraction.config['name'] != mapping.extraction.config['name']:

scantools/proc/overlap.py

@@ -100,7 +100,7 @@ def trajectory_overlap(self, keys_q: List, session_q: Session,
         if is_self:
             keys_r, session_r, poses_r = keys_q, session_q, poses_q
 
-        overlap_matrix = np.full((len(keys_q), len(keys_r)), -1, float)
+        overlap_matrix = np.full((len(keys_q), len(keys_r)), -1, dtype=np.float16)
         overlap_matrix[discard] = 0
 
         # cache the image poses as they might be compositions of rig poses
@@ -195,8 +195,8 @@ def compute_overlaps_for_sequence(capture: Capture, id_q: str, id_ref: str,
         valid_image_indices_list = [np.array(list(range(len(keys_ref))))]
         selected_keys_ref_list = [keys_ref]
 
-    ov_q2r = np.zeros([len(keys_q), len(keys_ref)])
-    ov_r2q = np.zeros([len(keys_ref), len(keys_q)])
+    ov_q2r = np.zeros([len(keys_q), len(keys_ref)], dtype=np.float16)
+    ov_r2q = np.zeros([len(keys_ref), len(keys_q)], dtype=np.float16)
     for sub_info in zip(sub_mesh_id_list, valid_image_indices_list, selected_keys_ref_list):
         sub_mesh_id, valid_image_indices, selected_keys_ref = sub_info
         sub_mesh_path = capture.proc_path(id_ref) / session_ref.proc.meshes[sub_mesh_id]