WIP : support for hloc toolbox (superpoint, superglue, netvlad,...)

nerfstudio/process_data/hloc_utils.py

@@ -0,0 +1,54 @@
+from pathlib import Path
+from typing_extensions import Literal
+
+from nerfstudio.process_data.process_data_utils import CameraModel
+from hloc import extract_features, match_features, reconstruction, visualization, pairs_from_retrieval, pairs_from_retrieval, pairs_from_exhaustive
+
+from hloc.triangulation import (import_features, import_matches)
+from hloc.utils.io import get_keypoints, get_matches
+
+import pycolmap
+
+def run_hloc(
+    image_dir: Path,
+    colmap_dir: Path,
+    camera_model: CameraModel,
+    gpu: bool = True,
+    verbose: bool = False,
+    matching_method: Literal["vocab_tree", "exhaustive", "sequential"] = "vocab_tree",
+    feature_type: Literal["sift", "superpoint_aachen", "superpoint_max", "superpoint_inloc", "r2d2", "d2net-ss", "sosnet", "disk"] = "superpoint_aachen",
+    matcher_type: Literal["superglue", "superglue-fast", "NN-superpoint", "NN-ratio", "NN-mutual", "adalam"] = "superglue",
+    num_matched: int = 25,
+) -> None:
+    """Runs hloc on the images.
+
+    Args:
+        image_dir: Path to the directory containing the images.
+        colmap_dir: Path to the output directory.
+        camera_model: Camera model to use.
+        gpu: If True, use GPU.
+        verbose: If True, logs the output of the command.
+    """
+    outputs = colmap_dir
+    sfm_pairs = outputs / 'pairs-netvlad.txt'
+    sfm_dir = outputs / 'sparse' / '0'
+    features = outputs / 'features.h5'
+    matches = outputs / 'matches.h5'
+
+    retrieval_conf = extract_features.confs['netvlad']
+    feature_conf = extract_features.confs[feature_type]
+    matcher_conf = match_features.confs[matcher_type]
+
+    references = image_dir.iterdir()
+    extract_features.main(feature_conf, image_dir, image_list=references, feature_path=features)
+    if matching_method == "exhaustive":
+        pairs_from_exhaustive.main(sfm_pairs, image_list=references)
+    else:
+        retrieval_path = extract_features.main(retrieval_conf, image_dir, outputs)
+        pairs_from_retrieval.main(retrieval_path, sfm_pairs, num_matched=num_matched)
+    match_features.main(matcher_conf, sfm_pairs, features=features, matches=matches)
+    match_path = match_features.main(matcher_conf, sfm_pairs, features=features, matches=matches)
+
+    image_options=pycolmap.ImageReaderOptions(camera_model=camera_model.value)
+    model = reconstruction.main(sfm_dir, image_dir, sfm_pairs, features, matches, camera_mode=pycolmap.CameraMode.SINGLE, image_options=image_options, verbose=verbose)
+

scripts/process_data.py

@@ -15,6 +15,7 @@
 
 from nerfstudio.process_data import (
     colmap_utils,
+    hloc_utils,
     insta360_utils,
     metashape_utils,
     polycam_utils,
@@ -26,6 +27,26 @@
 
 CONSOLE = Console(width=120)
 
+def findToolFeatureMatcherCombination(sfm_tool, feature_type, matcher_type):
+    if sfm_tool == "any":
+        if (feature_type == "any" or feature_type == "sift") and (matcher_type == "any" or matcher_type == "NN"):
+            sfm_tool = "colmap"
+        else:
+            sfm_tool = "hloc"
+
+    if sfm_tool == "colmap":
+        return ("colmap", "sift", "NN")
+    elif sfm_tool == "hloc":
+        if feature_type == "any" or feature_type == "superpoint":
+            feature_type = "superpoint_aachen"
+
+        if matcher_type == "any":
+            matcher_type = "superglue"
+        elif matcher_type == "NN":
+            matcher_type = "NN-mutual"
+
+        return (sfm_tool, feature_type, matcher_type)
+    return (None, None, None)
 
 @dataclass
 class ProcessImages:
@@ -46,6 +67,12 @@ class ProcessImages:
     matching_method: Literal["exhaustive", "sequential", "vocab_tree"] = "vocab_tree"
     """Feature matching method to use. Vocab tree is recommended for a balance of speed and
         accuracy. Exhaustive is slower but more accurate. Sequential is faster but should only be used for videos."""
+    sfm_tool: Literal["any", "colmap", "hloc"] = "any"
+    """Structure from motion tool to use. Colmap will use sift features, hloc can use many modern methods such as superpoint features and superglue matcher"""
+    feature_type: Literal["any", "sift", "superpoint", "superpoint_aachen", "superpoint_max", "superpoint_inloc", "r2d2", "d2net-ss", "sosnet", "disk"] = "any"
+    """Type of feature to use."""
+    matcher_type: Literal["any", "NN", "superglue", "superglue-fast", "NN-superpoint", "NN-ratio", "NN-mutual", "adalam"] = "any"
+    """Matching algorithm."""
     num_downscales: int = 3
     """Number of times to downscale the images. Downscales by 2 each time. For example a value of 3
         will downscale the images by 2x, 4x, and 8x."""
@@ -80,16 +107,32 @@ def main(self) -> None:
         colmap_dir = self.output_dir / "colmap"
         if not self.skip_colmap:
             colmap_dir.mkdir(parents=True, exist_ok=True)
-
-            colmap_utils.run_colmap(
-                image_dir=image_dir,
-                colmap_dir=colmap_dir,
-                camera_model=CAMERA_MODELS[self.camera_type],
-                gpu=self.gpu,
-                verbose=self.verbose,
-                matching_method=self.matching_method,
-                colmap_cmd=self.colmap_cmd,
-            )
+
+            (sfm_tool, feature_type, matcher_type) = findToolFeatureMatcherCombination(self.sfm_tool, self.feature_type, self.matcher_type)
+
+            if sfm_tool == "colmap":
+                colmap_utils.run_colmap(
+                    image_dir=image_dir,
+                    colmap_dir=colmap_dir,
+                    camera_model=CAMERA_MODELS[self.camera_type],
+                    gpu=self.gpu,
+                    verbose=self.verbose,
+                    matching_method=self.matching_method,
+                    colmap_cmd=self.colmap_cmd,
+                )
+            elif sfm_tool == "hloc":
+                hloc_utils.run_hloc(
+                    image_dir=image_dir,
+                    colmap_dir=colmap_dir,
+                    camera_model=CAMERA_MODELS[self.camera_type],
+                    gpu=self.gpu,
+                    verbose=self.verbose,
+                    matching_method=self.matching_method,
+                    feature_type=feature_type,
+                    matcher_type=matcher_type,
+                )
+            else:
+                CONSOLE.log("[bold yellow]Warning: Invalid combination of sfm_tool, feature_type, and matcher_type ")
 
         # Save transforms.json
         if (colmap_dir / "sparse" / "0" / "cameras.bin").exists():
@@ -134,6 +177,12 @@ class ProcessVideo:
     matching_method: Literal["exhaustive", "sequential", "vocab_tree"] = "vocab_tree"
     """Feature matching method to use. Vocab tree is recommended for a balance of speed and
         accuracy. Exhaustive is slower but more accurate. Sequential is faster but should only be used for videos."""
+    sfm_tool: Literal["any", "colmap", "hloc"] = "any"
+    """Structure from motion tool to use. Colmap will use sift features, hloc can use many modern methods such as superpoint features and superglue matcher"""
+    feature_type: Literal["any", "sift", "superpoint", "superpoint_aachen", "superpoint_max", "superpoint_inloc", "r2d2", "d2net-ss", "sosnet", "disk"] = "any"
+    """Type of feature to use."""
+    matcher_type: Literal["any", "NN", "superglue", "superglue-fast", "NN-superpoint", "NN-ratio", "NN-mutual", "adalam"] = "any"
+    """Matching algorithm."""
     num_downscales: int = 3
     """Number of times to downscale the images. Downscales by 2 each time. For example a value of 3
         will downscale the images by 2x, 4x, and 8x."""
@@ -167,16 +216,32 @@ def main(self) -> None:
         colmap_dir = self.output_dir / "colmap"
         if not self.skip_colmap:
             colmap_dir.mkdir(parents=True, exist_ok=True)
-
-            colmap_utils.run_colmap(
-                image_dir=image_dir,
-                colmap_dir=colmap_dir,
-                camera_model=CAMERA_MODELS[self.camera_type],
-                gpu=self.gpu,
-                verbose=self.verbose,
-                matching_method=self.matching_method,
-                colmap_cmd=self.colmap_cmd,
-            )
+
+            (sfm_tool, feature_type, matcher_type) = findToolFeatureMatcherCombination(self.sfm_tool, self.feature_type, self.matcher_type)
+
+            if sfm_tool == "colmap":
+                colmap_utils.run_colmap(
+                    image_dir=image_dir,
+                    colmap_dir=colmap_dir,
+                    camera_model=CAMERA_MODELS[self.camera_type],
+                    gpu=self.gpu,
+                    verbose=self.verbose,
+                    matching_method=self.matching_method,
+                    colmap_cmd=self.colmap_cmd,
+                )
+            elif sfm_tool == "hloc":
+                hloc_utils.run_hloc(
+                    image_dir=image_dir,
+                    colmap_dir=colmap_dir,
+                    camera_model=CAMERA_MODELS[self.camera_type],
+                    gpu=self.gpu,
+                    verbose=self.verbose,
+                    matching_method=self.matching_method,
+                    feature_type=feature_type,
+                    matcher_type=matcher_type,
+                )
+            else:
+                CONSOLE.log("[bold yellow]Warning: Invalid combination of sfm_tool, feature_type, and matcher_type ")
 
         # Save transforms.json
         if (colmap_dir / "sparse" / "0" / "cameras.bin").exists():