From 740cdebbb523294a9d9ff0750e8e521eb36f76ec Mon Sep 17 00:00:00 2001 From: RandomPrototypes <84458882+RandomPrototypes@users.noreply.github.com> Date: Sun, 4 Dec 2022 15:45:48 +0900 Subject: [PATCH] WIP : support for hloc toolbox (superpoint, superglue, netvlad,...) (#1026) * basic support for hloc (includes netvlad, super point, superglue in colmap), just a first test, not ready yet * add exhaustive match option for hloc * add command line option to choose the type of feature, matcher,... * make hloc an optional dependency, update doc, and small bug fixes * fix missing imports * fix warnings, pylint,... * fix overwrite previous commits * fix import order from isort * move find_tool_feature_matcher_combination to process_data_utils, add type hints and Args/Returns doc * error message when no valid combination of sfm tool, feature type, and matcher Co-authored-by: RandomPrototypes --- docs/reference/cli/ns_process_data.md | 3 +- nerfstudio/process_data/hloc_utils.py | 108 +++++++++++++++++ nerfstudio/process_data/process_data_utils.py | 54 +++++++++ scripts/process_data.py | 109 +++++++++++++++--- 4 files changed, 257 insertions(+), 17 deletions(-) create mode 100644 nerfstudio/process_data/hloc_utils.py diff --git a/docs/reference/cli/ns_process_data.md b/docs/reference/cli/ns_process_data.md index c2d09728e..bb5bd24ac 100644 --- a/docs/reference/cli/ns_process_data.md +++ b/docs/reference/cli/ns_process_data.md @@ -2,7 +2,8 @@ :::{admonition} Note :class: warning -Make sure to have [COLMAP](https://colmap.github.io) and [FFmpeg](https://ffmpeg.org/download.html) installed. +Make sure to have [COLMAP](https://colmap.github.io) and [FFmpeg](https://ffmpeg.org/download.html) installed. +You may also want to install [hloc](https://github.com/cvg/Hierarchical-Localization) (optional) for more feature detector and matcher options. ::: ```{eval-rst} diff --git a/nerfstudio/process_data/hloc_utils.py b/nerfstudio/process_data/hloc_utils.py new file mode 100644 index 000000000..17bd0120c --- /dev/null +++ b/nerfstudio/process_data/hloc_utils.py @@ -0,0 +1,108 @@ +""" +Code that uses the hierarchical localization toolbox (hloc) +to extract and match image features, estimate camera poses, +and do sparse reconstruction. +Requires hloc module from : https://github.com/cvg/Hierarchical-Localization +""" + +# Copyright 2022 The Nerfstudio Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import sys +from pathlib import Path + +from rich.console import Console +from typing_extensions import Literal + +from nerfstudio.process_data.process_data_utils import CameraModel + +try: + import pycolmap + from hloc import ( + extract_features, + match_features, + pairs_from_exhaustive, + pairs_from_retrieval, + reconstruction, + ) +except ImportError: + _HAS_HLOC = False +else: + _HAS_HLOC = True + +CONSOLE = Console(width=120) + + +def run_hloc( + image_dir: Path, + colmap_dir: Path, + camera_model: CameraModel, + 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 = 50, +) -> 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. + """ + if not _HAS_HLOC: + CONSOLE.print( + f"[bold red]Error: To use this set of parameters ({feature_type}/{matcher_type}/hloc), " + "you must install hloc toolbox!!" + ) + sys.exit(1) + + 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 = [p.relative_to(image_dir).as_posix() for p in 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) + if num_matched >= len(references): + num_matched = len(references) + pairs_from_retrieval.main(retrieval_path, sfm_pairs, num_matched=num_matched) + match_features.main(matcher_conf, sfm_pairs, features=features, matches=matches) + + image_options = pycolmap.ImageReaderOptions(camera_model=camera_model.value) + reconstruction.main( + sfm_dir, + image_dir, + sfm_pairs, + features, + matches, + camera_mode=pycolmap.CameraMode.SINGLE, + image_options=image_options, + verbose=verbose, + ) diff --git a/nerfstudio/process_data/process_data_utils.py b/nerfstudio/process_data/process_data_utils.py index 106d87f77..e291d898a 100644 --- a/nerfstudio/process_data/process_data_utils.py +++ b/nerfstudio/process_data/process_data_utils.py @@ -21,6 +21,7 @@ from pathlib import Path from typing import List, Optional, Tuple from rich.console import Console +from typing_extensions import Literal from nerfstudio.utils.rich_utils import status from nerfstudio.utils.scripts import run_command @@ -209,3 +210,56 @@ def downscale_images(image_dir: Path, num_downscales: int, verbose: bool = False downscale_text = [f"[bold blue]{2**(i+1)}x[/bold blue]" for i in range(num_downscales)] downscale_text = ", ".join(downscale_text[:-1]) + " and " + downscale_text[-1] return f"We downsampled the images by {downscale_text}" + + +def find_tool_feature_matcher_combination( + sfm_tool: Literal["any", "colmap", "hloc"], + feature_type: Literal[ + "any", + "sift", + "superpoint", + "superpoint_aachen", + "superpoint_max", + "superpoint_inloc", + "r2d2", + "d2net-ss", + "sosnet", + "disk", + ], + matcher_type: Literal[ + "any", "NN", "superglue", "superglue-fast", "NN-superpoint", "NN-ratio", "NN-mutual", "adalam" + ], +): + """Find a valid combination of sfm tool, feature type, and matcher type. + Basically, replace the default parameters 'any' by usable value + + Args: + sfm_tool: Sfm tool name (any, colmap, hloc) + feature_type: Type of image features (any, sift, superpoint, ...) + matcher_type: Type of matching algorithm (any, NN, superglue,...) + + Returns: + Tuple of sfm tool, feature type, and matcher type. + Returns (None,None,None) if no valid combination can be found + """ + if sfm_tool == "any": + if (feature_type in ("any", "sift")) and (matcher_type in ("any", "NN")): + sfm_tool = "colmap" + else: + sfm_tool = "hloc" + + if sfm_tool == "colmap": + if (feature_type not in ("any", "sift")) or (matcher_type not in ("any", "NN")): + return (None, None, None) + return ("colmap", "sift", "NN") + if sfm_tool == "hloc": + if feature_type in ("any", "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) diff --git a/scripts/process_data.py b/scripts/process_data.py index 0b8240008..fcea3fb44 100755 --- a/scripts/process_data.py +++ b/scripts/process_data.py @@ -16,6 +16,7 @@ from typing_extensions import Annotated, Literal from nerfstudio.process_data import ( colmap_utils, + hloc_utils, insta360_utils, metashape_utils, polycam_utils, @@ -47,6 +48,26 @@ 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.""" @@ -82,16 +103,34 @@ class ProcessImages: 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) = process_data_utils.find_tool_feature_matcher_combination( + 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], + verbose=self.verbose, + matching_method=self.matching_method, + feature_type=feature_type, + matcher_type=matcher_type, + ) + else: + CONSOLE.log("[bold red]Invalid combination of sfm_tool, feature_type, and matcher_type, exiting") + sys.exit(1) + # Save transforms.json if (colmap_dir / "sparse" / "0" / "cameras.bin").exists(): with CONSOLE.status("[bold yellow]Saving results to transforms.json", spinner="balloon"): @@ -135,6 +174,26 @@ 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.""" @@ -169,16 +228,34 @@ class ProcessVideo: 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) = process_data_utils.find_tool_feature_matcher_combination( + 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], + verbose=self.verbose, + matching_method=self.matching_method, + feature_type=feature_type, + matcher_type=matcher_type, + ) + else: + CONSOLE.log("[bold red]Invalid combination of sfm_tool, feature_type, and matcher_type, exiting") + sys.exit(1) + # Save transforms.json if (colmap_dir / "sparse" / "0" / "cameras.bin").exists(): with CONSOLE.status("[bold yellow]Saving results to transforms.json", spinner="balloon"):