Frontend 1/5 speedup: CPU reuse and upstream MegaLoc

vidmap/frontend/cache/metadata.py

@@ -51,6 +51,8 @@ def canonical_value(value: Any) -> Any:
     if isinstance(value, bytes):
         return value.decode("utf-8")
     if isinstance(value, np.ndarray):
+        if value.ndim and value.dtype.kind in "biufU":
+            return value.tolist()
         return [canonical_value(item) for item in value.tolist()]
     if isinstance(value, np.generic):
         return value.item()

vidmap/frontend/initial_reconstruction.py

@@ -6,18 +6,15 @@
 def build_initial_reconstruction(
     scene_parser, reference_image_names: list[str] | None = None
 ) -> pycolmap.Reconstruction:
-    """Build the initial COLMAP reconstruction from a scene reconstruction."""
-    if scene_parser.reconstruction_dir is not None:
-        refrec = pycolmap.Reconstruction(scene_parser.reconstruction_dir)
-    else:
-        refrec = scene_parser.rec
+    """Copy selected images and their cameras from the prepared scene."""
+    source = scene_parser.rec
+    names = None if reference_image_names is None else set(reference_image_names)
+    images = [image for image in source.images.values() if names is None or image.name in names]
     rec = pycolmap.Reconstruction()
-    add_cameras = {camera.camera_id: camera for camera in refrec.cameras.values()}
-    for camera in add_cameras.values():
-        rec.add_camera_with_trivial_rig(deepcopy(camera))
-    for imid, image in refrec.images.items():
-        if reference_image_names is not None and image.name not in reference_image_names:
-            continue
-        image_ = pycolmap.Image(image_id=imid, name=image.name, camera_id=image.camera_id)
-        rec.add_image_with_trivial_frame(image_)
+    for camera_id in dict.fromkeys(image.camera_id for image in images):
+        rec.add_camera_with_trivial_rig(deepcopy(source.cameras[camera_id]))
+    for image in images:
+        rec.add_image_with_trivial_frame(
+            pycolmap.Image(image_id=image.image_id, name=image.name, camera_id=image.camera_id)
+        )
     return rec

vidmap/frontend/keyframes/matching.py

@@ -13,14 +13,18 @@ class ImagePairDataset(torch.utils.data.Dataset):
     def __init__(self, image_dataset, pair_indices):
         self.image_dataset = image_dataset
         self.pair_indices = pair_indices
+        self.previous = None
 
     def __len__(self):
         return len(self.pair_indices)
 
     def __getitem__(self, index):
         first, second = self.pair_indices[index]
-        image_a = self.image_dataset[first]
+        image_a = (
+            self.previous[1] if self.previous is not None and self.previous[0] == first else self.image_dataset[first]
+        )
         image_b = self.image_dataset[second]
+        self.previous = (second, image_b)
         return image_a["image"], image_b["image"], image_a["name"], image_b["name"]
 
 

vidmap/frontend/models/megaloc.py

@@ -14,7 +14,7 @@
 from vidmap.frontend.cache import file_fingerprint
 from vidmap.model_sources import model_package_root
 
-MEGALOC_SOURCE_REVISION = "1af071c68fc3ab6c6018c5c868391763516e50f7"
+MEGALOC_SOURCE_REVISION = "5fe0dd697c4a70ba3e23607f6716ab3c606b16db"
 MEGALOC_MODEL_REVISION = "7cb9f7970d366fdf059963d04d372e503e8e9df9"
 MEGALOC_MODEL_SHA256 = "d4f9f2bcb60018f91eb6a8e061ed054fd55654e10c2569cf13841ea986ffb4f8"
 MEGALOC_SOURCE = model_package_root("MegaLoc", "third_party/MegaLoc")
@@ -44,8 +44,9 @@ def __init__(self):
             raise RuntimeError(
                 f"MegaLoc weights at {pinned_weights} have sha256 {actual}, expected {MEGALOC_MODEL_SHA256}"
             )
-        self.net = _create_megaloc_model()
-        self.net.load_state_dict(load_file(pinned_weights))
+        with torch.device("meta"):
+            self.net = _create_megaloc_model()
+        self.net.load_state_dict(load_file(pinned_weights), strict=True, assign=True)
         self.net.eval()
         mean = [0.485, 0.456, 0.406]
         std = [0.229, 0.224, 0.225]
@@ -68,4 +69,5 @@ def megaloc_cache_identity():
         "source_sha256": file_fingerprint(MEGALOC_SOURCE_FILE),
         "model_revision": MEGALOC_MODEL_REVISION,
         "model_sha256": MEGALOC_MODEL_SHA256,
+        "torch": torch.__version__,
     }

vidmap/frontend/tracking/kernels.py

@@ -76,24 +76,20 @@ def select_lc_matches_from_dense(
     valid = certainty_sampled > lc_match_thresh
     mkp_ids1[~valid] = -1
 
+    matches0 = np.full(len(kpts0_np), -1, dtype=np.int32)
+    matching_scores0 = np.zeros(len(kpts0_np), dtype=np.float32)
     valid_mask = mkp_ids1 >= 0
     if valid_mask.any():
         ref_indices = mkp_ids1[valid_mask]
         query_indices = np.where(valid_mask)[0]
         certainties = certainty_sampled[valid_mask]
 
-        matches0 = np.full(len(kpts0_np), -1, dtype=np.int32)
-        matching_scores0 = np.zeros(len(kpts0_np), dtype=np.float32)
-        for ref_idx in np.unique(ref_indices):
-            mask = ref_indices == ref_idx
-            query_candidates = query_indices[mask]
-            certainty_candidates = certainties[mask]
-            best_idx = query_candidates[np.argmax(certainty_candidates)]
-            matches0[best_idx] = ref_idx
-            matching_scores0[best_idx] = certainty_sampled[best_idx]
-    else:
-        matches0 = np.full(len(kpts0_np), -1, dtype=np.int32)
-        matching_scores0 = np.zeros(len(kpts0_np), dtype=np.float32)
+        # Preserve first-source tie breaking within each target keypoint.
+        order = np.lexsort((query_indices, -certainties, ref_indices))
+        first = np.r_[True, ref_indices[order][1:] != ref_indices[order][:-1]]
+        best = query_indices[order[first]]
+        matches0[best] = ref_indices[order[first]]
+        matching_scores0[best] = certainty_sampled[best]
 
     return {
         "matches0": matches0,

vidmap/frontend/tracking/propagation.py

@@ -42,19 +42,19 @@ class DenseMatchField:
     target_size: tuple[int, int]
 
 
-@dataclass
+@dataclass(frozen=True)
 class DenseHopBuffers:
-    matches: np.ndarray
-    certainty: np.ndarray
-    covariance: np.ndarray
+    matches: dict[int, np.ndarray]
+    certainty: dict[int, np.ndarray]
+    covariance: dict[int, np.ndarray]
 
     @classmethod
-    def create(cls, max_hop, current_size):
-        width, height = current_size
+    def from_fields(cls, records, fields):
+        """Borrow CPU tensor storage only for this target's scheduled hops."""
         return cls(
-            matches=np.full((max_hop, height, width, 2), -1, dtype=np.float32),
-            certainty=np.full((max_hop, height, width), -1, dtype=np.float32),
-            covariance=np.full((max_hop, height, width, 2, 2), -1, dtype=np.float32),
+            matches={r.slot: f.matches.numpy().astype(np.float32, copy=False) for r, f in zip(records, fields)},
+            certainty={r.slot: f.certainty.numpy().astype(np.float32, copy=False) for r, f in zip(records, fields)},
+            covariance={r.slot: f.covariance.numpy().astype(np.float32, copy=False) for r, f in zip(records, fields)},
         )
 
 
@@ -285,7 +285,6 @@ def _cleanup(self):
             "highres_dataset",
             "lowres_dataset",
             "state",
-            "dense_hops",
         ):
             if attribute in vars(self):
                 delattr(self, attribute)
@@ -397,15 +396,10 @@ def _run_streaming_loop(
                     current_size=direct.source_size,
                     scheduled_hops=self.conf.multiflow_hops,
                 )
-                self.dense_hops = DenseHopBuffers.create(max(self.conf.multiflow_hops), direct.source_size)
             state = self.state
-            dense_hops = self.dense_hops
             if any(field.source_size != state.current_size for field in fields):
                 raise RuntimeError("Dense multiflow fields changed processing-grid size within one run")
-            for record, field in zip(records, fields):
-                dense_hops.matches[record.slot] = field.matches
-                dense_hops.certainty[record.slot] = field.certainty
-                dense_hops.covariance[record.slot] = field.covariance
+            dense_hops = DenseHopBuffers.from_fields(records, fields)
 
             certainty = direct.certainty.clone()
             matches = direct.matches

vidmap/frontend/tracking/state.py

@@ -287,8 +287,9 @@ def _select_source_keypoints(self, *, kf_id, certainty, covariance, salient_keyp
         if self.prev_keypoints is not None:
             check_existing_kps = scale_keypoints(self.prev_keypoints[previous_mask], self.scale_ratio)
             if len(check_existing_kps) > 0 and len(salient_keypoints) > 0:
-                diff = salient_keypoints[:, None, :] - check_existing_kps[None, :, :]
-                sq_dists = (diff**2).sum(axis=-1)
+                dx = salient_keypoints[:, None, 0] - check_existing_kps[None, :, 0]
+                dy = salient_keypoints[:, None, 1] - check_existing_kps[None, :, 1]
+                sq_dists = dx**2 + dy**2
                 keep_salient_mask = (sq_dists > self.conf.nms_radius**2).all(axis=1)
 
         sample = nn_sample_2d