[WIP] Add support for dense matching + dense SfM

lamar/run.py

@@ -85,28 +85,28 @@ def run(outputs: Path,
             capture, query_id, query_list, sequence_length_seconds)
         image_keys = keys_from_chunks(query_chunks)
 
-    extraction_map = FeatureExtraction(outputs, capture, ref_id, configs['extraction'])
     pairs_map = PairSelection(outputs, capture, ref_id, ref_id, configs['pairs_map'])
+    # extraction_map = FeatureExtraction(outputs, capture, ref_id, configs['extraction'])
     matching_map = FeatureMatching(
-        outputs, capture, ref_id, ref_id, configs['matching'], pairs_map, extraction_map)
+        outputs, capture, ref_id, ref_id, {'extraction': configs['extraction'], 'matching': configs['matching']}, pairs_map)
 
     mapping = Mapping(
-        configs['mapping'], outputs, capture, ref_id, extraction_map, matching_map)
+        configs['mapping'], outputs, capture, ref_id, matching_map.extraction_ref, matching_map)
 
-    extraction_query = FeatureExtraction(
-        outputs, capture, query_id, configs['extraction'], image_keys)
+    # extraction_query = FeatureExtraction(
+    #     outputs, capture, query_id, configs['extraction'], image_keys)
 
     if is_sequential:
         query_list, query_chunks = avoid_duplicate_keys_in_chunks(
             session_q, query_list, query_chunks)
         T_c2w_gt = session_q.proc.alignment_trajectories
-        chunk_alignment = ChunkAlignment(
-            configs, outputs, capture, query_id, extraction_query, mapping, query_chunks,
-            sequence_length_seconds)
-        if T_c2w_gt:
-            results = chunk_alignment.evaluate(T_c2w_gt, query_list)
-        else:
-            results = str(chunk_alignment.paths.poses)
+        # chunk_alignment = ChunkAlignment(
+        #     configs, outputs, capture, query_id, extraction_query, mapping, query_chunks,
+        #     sequence_length_seconds)
+        # if T_c2w_gt:
+        #     results = chunk_alignment.evaluate(T_c2w_gt, query_list)
+        # else:
+        #     results = str(chunk_alignment.paths.poses)
     else:
         T_c2w_gt = session_q.proc.alignment_trajectories
         if T_c2w_gt and is_rig and not do_rig:
@@ -115,11 +115,11 @@ def run(outputs: Path,
             outputs, capture, query_id, ref_id, configs['pairs_loc'], query_list,
             query_poses=T_c2w_gt)
         matching_query = FeatureMatching(
-            outputs, capture, query_id, ref_id, configs['matching_query'],
-            pairs_loc, extraction_query, extraction_map)
+            outputs, capture, query_id, ref_id, {'extraction': configs['extraction'], 'matching': configs['matching_query']},
+            pairs_loc)
         pose_estimation = PoseEstimation(
             configs['poses'], outputs, capture, query_id,
-            extraction_query, matching_query, mapping, query_list)
+            matching_query.extraction, matching_query, mapping, query_list)
         if T_c2w_gt:
             results = pose_estimation.evaluate(T_c2w_gt)
         else:

lamar/tasks/feature_extraction.py

@@ -20,6 +20,9 @@ def __init__(self, root, config, session_id):
 
 class FeatureExtraction:
     methods = {
+        'anypoint': {
+            'name': 'anypoint'
+        },
         'superpoint': {
             'name': 'superpoint',
             'hloc': {
@@ -119,14 +122,17 @@ def __init__(self, outputs, capture, session_id, config, query_keys=None, overwr
         logger.info('Extraction local features %s for session %s.', config['name'], session_id)
         _, names, image_root = list_images_for_session(capture, session_id, query_keys)
         names = np.unique(names)
-        extract_features.main(
-            config['hloc'],
-            image_root,
-            feature_path=self.paths.features,
-            image_list=names,
-            as_half=True,
-            overwrite=overwrite,
-        )
+        self.image_root = image_root
+        self.names = names
+        if 'hloc' in config:
+            extract_features.main(
+                config['hloc'],
+                image_root,
+                feature_path=self.paths.features,
+                image_list=names,
+                as_half=True,
+                overwrite=overwrite,
+            )
 
         write_config(config, self.paths.config)
 
@@ -147,4 +153,11 @@ class RetrievalFeatureExtraction(FeatureExtraction):
                 'preprocessing': {'resize_max': 640},
             }
         },
+        'salad': {
+            'name': 'salad',
+            'hloc': {
+                'model': {'name': 'salad'},
+                'preprocessing': {'resize_max': 640},
+            }
+        }
     }

lamar/tasks/feature_matching.py

@@ -2,7 +2,7 @@
 from typing import Optional
 from copy import deepcopy
 
-from hloc import match_features
+from hloc import match_features, match_dense
 
 from .feature_extraction import FeatureExtraction
 from .pair_selection import PairSelection
@@ -23,6 +23,22 @@ def __init__(self, root, config, query_id, ref_id):
 
 class FeatureMatching:
     methods = {
+        'loftr': {
+            'name': 'loftr',
+            'hloc': {
+                'model': {
+                    'name': 'loftr',
+                    'weights': 'outdoor'
+                },
+                'preprocessing': {
+                    'grayscale': True,
+                    'resize_max': 1024,
+                    'dfactor': 8
+                },
+                'max_error': 2,  # max error for assigned keypoints (in px)
+                'cell_size': 8,  # size of quantization patch (max 1 kp/patch)
+            }
+        },
         'superglue': {
             'name': 'superglue',
             'hloc': {
@@ -74,43 +90,57 @@ class FeatureMatching:
 
     def __init__(self, outputs, capture, query_id, ref_id, config,
                  pair_selection: PairSelection,
-                 extraction: FeatureExtraction,
-                 extraction_ref: Optional[FeatureExtraction] = None,
                  overwrite=False):
-        extraction_ref = extraction_ref or extraction
-        if extraction.config['name'] != extraction_ref.config['name']:
-            raise ValueError('Matching two different features: '
-                             f'{extraction.config} vs {extraction_ref.config}')
-        assert query_id == extraction.session_id
-        assert query_id == pair_selection.query_id
-        assert ref_id == extraction_ref.session_id
-        assert ref_id == pair_selection.ref_id
-
-        self.config = config = {
-            **deepcopy(config),
-            'features': extraction.config,  # detect upstream changes
-            # do not include the pairs so the same file can be reused
-        }
+        self.config = config = deepcopy(config)
         self.query_id = query_id
         self.ref_id = ref_id
-        self.extraction = extraction
-        self.extraction_ref = extraction_ref
         self.pair_selection = pair_selection
-        self.paths = FeatureMatchingPaths(outputs, config, query_id, ref_id)
+        self.extraction = FeatureExtraction(outputs, capture, query_id, config['extraction'])
+        self.extraction_ref = FeatureExtraction(outputs, capture, ref_id, config['extraction'])
+        self.paths = FeatureMatchingPaths(outputs, {**config['matching'], 'features': self.extraction.config}, query_id, ref_id)
         self.paths.workdir.mkdir(parents=True, exist_ok=True)
 
-        logger.info('Matching local features with %s for sessions (%s, %s).',
-                    config['name'], query_id, ref_id)
-        if not same_configs(config, self.paths.config):
-            logger.warning('Existing matches will be overwritten.')
-            overwrite = True
-        match_features.main(
-            config['hloc'],
-            pair_selection.paths.pairs_hloc,
-            extraction.paths.features,
-            matches=self.paths.matches,
-            features_ref=extraction_ref.paths.features,
-            overwrite=overwrite,
-        )
+        if 'hloc' in config['extraction']:
+            logger.info('Matching local features with %s for sessions (%s, %s).',
+                        config['matching']['name'], query_id, ref_id)
+            if not same_configs(config, self.paths.config):
+                logger.warning('Existing matches will be overwritten.')
+                overwrite = True
+            match_features.main(
+                config['matching']['hloc'],
+                pair_selection.paths.pairs_hloc,
+                self.extraction.paths.features,
+                matches=self.paths.matches,
+                features_ref=self.extraction_ref.paths.features,
+                overwrite=overwrite,
+            )
+            write_config(config, self.paths.config)
+        else:
+            logger.info('Matching dense features with %s for sessions (%s, %s).',
+                        config['matching']['name'], query_id, ref_id)
+            if not same_configs(config, self.paths.config):
+                logger.warning('Existing matches will be overwritten.')
+                overwrite = True
+            if query_id == ref_id:
+                match_dense.main(
+                    config['matching']['hloc'],
+                    pair_selection.paths.pairs_hloc,
+                    self.extraction.image_root,
+                    matches=self.paths.matches,
+                    features=self.extraction.paths.features,
+                    max_kps=8192,
+                    overwrite=overwrite,
+                )
+            else:
+                match_dense.main(
+                    config['matching']['hloc'],
+                    pair_selection.paths.pairs_hloc,
+                    self.extraction.image_root,
+                    matches=self.paths.matches,
+                    features=self.extraction.paths.features,
+                    features_ref=self.extraction_ref.paths.features,
+                    max_kps=None,
+                    overwrite=overwrite,
+                )
+            write_config(config, self.paths.config)
 
-        write_config(config, self.paths.config)

lamar/tasks/mapping.py

@@ -24,7 +24,7 @@ def __init__(self, root, config, session_id):
         self.root = root
         self.workdir = root / 'mapping' / session_id / config['name'] / config['features']['name']
         if config['name'] == 'triangulation':
-            self.workdir /= Path(config['pairs']['name'], config['matches']['name'])
+            self.workdir /= Path(config['pairs']['name'], config['matches']['matching']['name'])
             self.sfm_empty = self.workdir / 'sfm_empty'
         self.sfm = self.workdir / 'sfm'
         self.config = self.workdir / 'configuration.json'
@@ -94,6 +94,7 @@ def __init__(self, config, outputs, capture, session_id,
             self.name2key[image.name]: image.image_id
             for image in self.reconstruction.images.values()
         }
+        self.points3d_cache = {}
 
     def run(self, capture):
         run_capture_to_empty_colmap.run(capture, [self.session_id], self.paths.sfm_empty)
@@ -107,18 +108,24 @@ def run(self, capture):
         )
 
     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)
+        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)
-        return np.array(valid, bool), xyz, ids
+                else:
+                    ids.append(-1)
+                    xyz.append([np.nan, np.nan, np.nan])
+            self.points3d_cache[key] = (np.array(ids), np.array(xyz))
+        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]
 
 
 class MeshLifting(Mapping):

lamar/tasks/pose_estimation.py

@@ -31,7 +31,7 @@ def __init__(self, root, config, query_id, ref_id, override_workdir_root=None):
             root = override_workdir_root
         self.workdir = (
             root / 'pose_estimation' / query_id / ref_id
-            / config['features']['name'] / config['matches']['name']
+            / config['features']['name'] / config['matches']['matching']['name']
             / config['pairs']['name'] / config['mapping']['name'] / config['name']
         )
         self.poses = self.workdir / 'poses.txt'
@@ -51,7 +51,7 @@ class PoseEstimation:
     method2class = {}
     method = None
     evaluation = {
-        'Rt_thresholds': [(1, 0.1), (5, 1.)],
+        'Rt_thresholds': [(5, 0.5)],
     }
 
     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']:

lamar/utils/localization.py

@@ -66,13 +66,22 @@ def estimate_camera_pose(query: str, camera: Camera,
                          recover_matches: Callable,
                          pnp_error_multiplier: float,
                          return_covariance: bool) -> Pose:
+    # import time
+    # t0 = time.time()
     matches_2d3d = recover_matches(query, ref_key_names)
+    # t1 = time.time()
+    # print('recover_matches', t1 - t0)
     keypoint_noise = matches_2d3d['keypoint_noise']
 
+    if keypoint_noise is None:
+        keypoint_noise = 1
+    # print(matches_2d3d['kp_q'].shape[0], 'matches')
     ret = pycolmap.absolute_pose_estimation(
         matches_2d3d['kp_q'], matches_2d3d['p3d'],
         camera.asdict, pnp_error_multiplier * keypoint_noise,
         return_covariance=return_covariance)
+    # t2 = time.time()
+    # print('absolute_pose_estimation', t2 - t1)
 
     if ret['success']:
         if return_covariance: