tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-10 18:17:31 +08:00
Code

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

Browse Source 
* 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 <delmotte_ar@yahoo.fr>
This commit is contained in:
RandomPrototypes 2022-12-04 15:45:48 +09:00 committed by GitHub
parent f6ac1d72ff
commit 740cdebbb5
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 257 additions and 17 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
docs/reference/cli/ns_process_data.md
View File

@ -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}

108
nerfstudio/process_data/hloc_utils.py Normal file
View File

@ -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,
)

54
nerfstudio/process_data/process_data_utils.py
View File

@ -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)

109
scripts/process_data.py
View File

@ -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"):