Skip to content

Commit

Permalink
Ruff format
Browse files Browse the repository at this point in the history
  • Loading branch information
@luishresende
luishresende committed Oct 28, 2024
1 parent a9f1cf3 commit 000e344
Show file tree
Hide file tree
Showing 2 changed files with 37 additions and 94 deletions.
35 changes: 9 additions & 26 deletions nerfstudio/process_data/colmap_converter_to_nerfstudio_dataset.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,10 +19,8 @@
from pathlib import Path
from typing import Dict, List, Literal, Optional, Tuple

from nerfstudio.process_data import (colmap_utils, hloc_utils,
process_data_utils)
from nerfstudio.process_data.base_converter_to_nerfstudio_dataset import \
BaseConverterToNerfstudioDataset
from nerfstudio.process_data import colmap_utils, hloc_utils, process_data_utils
from nerfstudio.process_data.base_converter_to_nerfstudio_dataset import BaseConverterToNerfstudioDataset
from nerfstudio.process_data.process_data_utils import CAMERA_MODELS
from nerfstudio.utils import install_checks
from nerfstudio.utils.rich_utils import CONSOLE
Expand All @@ -32,9 +30,7 @@
class ColmapConverterToNerfstudioDataset(BaseConverterToNerfstudioDataset):
"""Base class to process images or video into a nerfstudio dataset using colmap"""

camera_type: Literal[
"perspective", "fisheye", "equirectangular", "pinhole", "simple_pinhole"
] = "perspective"
camera_type: Literal["perspective", "fisheye", "equirectangular", "pinhole", "simple_pinhole"] = "perspective"
"""Camera model to use."""
matching_method: Literal["exhaustive", "sequential", "vocab_tree"] = "vocab_tree"
"""Feature matching method to use. Vocab tree is recommended for a balance of speed
Expand Down Expand Up @@ -159,9 +155,7 @@ def _save_transforms(

best_model_index = matched_frame_counts.index(max(matched_frame_counts))
best_model_path = models_path / models[best_model_index]
with CONSOLE.status(
"[bold yellow]Saving results to transforms.json", spinner="balloon"
):
with CONSOLE.status("[bold yellow]Saving results to transforms.json", spinner="balloon"):
num_matched_frames = colmap_utils.colmap_to_json(
recon_dir=best_model_path,
output_dir=self.output_dir,
Expand All @@ -171,13 +165,10 @@ def _save_transforms(
use_single_camera_mode=self.use_single_camera_mode,
)
summary_log.append(f"Colmap matched {num_matched_frames} images")
summary_log.append(
colmap_utils.get_matching_summary(num_frames, num_matched_frames)
)
summary_log.append(colmap_utils.get_matching_summary(num_frames, num_matched_frames))
else:
CONSOLE.log(
"[bold yellow]Warning: Could not find existing COLMAP results. "
"Not generating transforms.json"
"[bold yellow]Warning: Could not find existing COLMAP results. " "Not generating transforms.json"
)
return summary_log

Expand Down Expand Up @@ -233,9 +224,7 @@ def _run_colmap(self, mask_path: Optional[Path] = None):

# check that sfm_tool is hloc if using use_single_camera_mode
if not self.use_single_camera_mode:
assert (
sfm_tool == "hloc"
), "not_use_single_camera_mode only works with sfm_tool hloc"
assert sfm_tool == "hloc", "not_use_single_camera_mode only works with sfm_tool hloc"

# set the image_dir if didn't copy
if self.skip_image_processing:
Expand All @@ -257,10 +246,7 @@ def _run_colmap(self, mask_path: Optional[Path] = None):
)
elif sfm_tool == "hloc":
if mask_path is not None:
raise RuntimeError(
"Cannot use a mask with hloc. Please remove the cropping options "
"and try again."
)
raise RuntimeError("Cannot use a mask with hloc. Please remove the cropping options " "and try again.")

assert feature_type is not None
assert matcher_type is not None
Expand All @@ -277,10 +263,7 @@ def _run_colmap(self, mask_path: Optional[Path] = None):
use_single_camera_mode=self.use_single_camera_mode,
)
else:
raise RuntimeError(
"Invalid combination of sfm_tool, feature_type, and matcher_type, "
"exiting"
)
raise RuntimeError("Invalid combination of sfm_tool, feature_type, and matcher_type, " "exiting")

def __post_init__(self) -> None:
super().__post_init__()
Expand Down
96 changes: 28 additions & 68 deletions nerfstudio/process_data/colmap_utils.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,19 +25,22 @@
import numpy as np
import requests
import torch
from packaging.version import Version
from rich.progress import track

# TODO(1480) use pycolmap instead of colmap_parsing_utils
# import pycolmap
from nerfstudio.data.utils.colmap_parsing_utils import (qvec2rotmat,
read_cameras_binary,
read_images_binary,
read_points3D_binary,
read_points3D_text)
from nerfstudio.data.utils.colmap_parsing_utils import (
qvec2rotmat,
read_cameras_binary,
read_images_binary,
read_points3D_binary,
read_points3D_text,
)
from nerfstudio.process_data.process_data_utils import CameraModel
from nerfstudio.utils import colormaps
from nerfstudio.utils.rich_utils import CONSOLE, status
from nerfstudio.utils.scripts import run_command
from packaging.version import Version
from rich.progress import track


def get_colmap_version(colmap_cmd: str, default_version: str = "3.8") -> Version:
Expand All @@ -57,9 +60,7 @@ def get_colmap_version(colmap_cmd: str, default_version: str = "3.8") -> Version
version = line.split(" ")[1]
version = Version(version)
return version
CONSOLE.print(
f"[bold red]Could not find COLMAP version. Using default {default_version}"
)
CONSOLE.print(f"[bold red]Could not find COLMAP version. Using default {default_version}")
return Version(default_version)


Expand All @@ -72,9 +73,7 @@ def get_vocab_tree() -> Path:
vocab_tree_filename = Path(appdirs.user_data_dir("nerfstudio")) / "vocab_tree.fbow"

if not vocab_tree_filename.exists():
r = requests.get(
"https://demuc.de/colmap/vocab_tree_flickr100K_words32K.bin", stream=True
)
r = requests.get("https://demuc.de/colmap/vocab_tree_flickr100K_words32K.bin", stream=True)
vocab_tree_filename.parent.mkdir(parents=True, exist_ok=True)
with open(vocab_tree_filename, "wb") as f:
total_length = r.headers.get("content-length")
Expand Down Expand Up @@ -130,9 +129,7 @@ def run_colmap(
f"--SiftExtraction.use_gpu {int(gpu)}",
]
if camera_mask_path is not None:
feature_extractor_cmd.append(
f"--ImageReader.camera_mask_path {camera_mask_path}"
)
feature_extractor_cmd.append(f"--ImageReader.camera_mask_path {camera_mask_path}")
feature_extractor_cmd = " ".join(feature_extractor_cmd)
with status(
msg="[bold yellow]Running COLMAP feature extractor...",
Expand All @@ -151,9 +148,7 @@ def run_colmap(
]
if matching_method == "vocab_tree":
vocab_tree_filename = get_vocab_tree()
feature_matcher_cmd.append(
f'--VocabTreeMatching.vocab_tree_path "{vocab_tree_filename}"'
)
feature_matcher_cmd.append(f'--VocabTreeMatching.vocab_tree_path "{vocab_tree_filename}"')
feature_matcher_cmd = " ".join(feature_matcher_cmd)
with status(
msg="[bold yellow]Running COLMAP feature matcher...",
Expand Down Expand Up @@ -186,9 +181,7 @@ def run_colmap(
CONSOLE.log("[bold green]:tada: Done COLMAP bundle adjustment.")

if refine_intrinsics:
with status(
msg="[bold yellow]Refine intrinsics...", spinner="dqpb", verbose=verbose
):
with status(msg="[bold yellow]Refine intrinsics...", spinner="dqpb", verbose=verbose):
bundle_adjuster_cmd = [
f"{colmap_cmd} bundle_adjuster",
f"--input_path {sparse_dir}/0",
Expand Down Expand Up @@ -438,10 +431,7 @@ def colmap_to_json(
cam_id_to_camera = read_cameras_binary(recon_dir / "cameras.bin")
im_id_to_image = read_images_binary(recon_dir / "images.bin")
if set(cam_id_to_camera.keys()) != {1}:
CONSOLE.print(
f"[bold yellow]Warning: More than one camera is found in {recon_dir}"
)
print(cam_id_to_camera)
CONSOLE.print(f"[bold yellow]Warning: More than one camera is found in {recon_dir}")
use_single_camera_mode = False # update bool: one camera per frame
out = {} # out = {"camera_model": parse_colmap_camera_params(cam_id_to_camera[1])["camera_model"]}
else: # one camera for all frames
Expand Down Expand Up @@ -479,19 +469,13 @@ def colmap_to_json(
"colmap_im_id": im_id,
}
if camera_mask_path is not None:
frame["mask_path"] = camera_mask_path.relative_to(
camera_mask_path.parent.parent
).as_posix()
frame["mask_path"] = camera_mask_path.relative_to(camera_mask_path.parent.parent).as_posix()
if image_id_to_depth_path is not None:
depth_path = image_id_to_depth_path[im_id]
frame["depth_file_path"] = str(
depth_path.relative_to(depth_path.parent.parent)
)
frame["depth_file_path"] = str(depth_path.relative_to(depth_path.parent.parent))

if not use_single_camera_mode: # add the camera parameters for this frame
frame.update(
parse_colmap_camera_params(cam_id_to_camera[im_data.camera_id])
)
frame.update(parse_colmap_camera_params(cam_id_to_camera[im_data.camera_id]))

frames.append(frame)

Expand All @@ -505,16 +489,12 @@ def colmap_to_json(
out["applied_transform"] = applied_transform.tolist()

# create ply from colmap
assert ply_filename.endswith(
".ply"
), f"ply_filename: {ply_filename} does not end with '.ply'"
assert ply_filename.endswith(".ply"), f"ply_filename: {ply_filename} does not end with '.ply'"
create_ply_from_colmap(
ply_filename,
recon_dir,
output_dir,
torch.from_numpy(applied_transform).float()
if applied_transform is not None
else None,
torch.from_numpy(applied_transform).float() if applied_transform is not None else None,
)
out["ply_file_path"] = ply_filename

Expand Down Expand Up @@ -601,13 +581,7 @@ def create_sfm_depth(
z = (rotation @ xyz_world.T)[-1] + im_data.tvec[-1]
errors = np.array([ptid_to_info[pid].error for pid in pids])
n_visible = np.array([len(ptid_to_info[pid].image_ids) for pid in pids])
uv = np.array(
[
im_data.xys[i]
for i in range(len(im_data.xys))
if im_data.point3D_ids[i] != -1
]
)
uv = np.array([im_data.xys[i] for i in range(len(im_data.xys)) if im_data.point3D_ids[i] != -1])
# TODO(1480) END delete when abandoning colmap_parsing_utils

# TODO(1480) BEGIN use pycolmap API
Expand Down Expand Up @@ -662,16 +636,11 @@ def create_sfm_depth(
image_id_to_depth_path[im_id] = depth_path

if include_depth_debug:
assert (
input_images_dir is not None
), "Need explicit input_images_dir for debug images"
assert input_images_dir is not None, "Need explicit input_images_dir for debug images"
assert input_images_dir.exists(), input_images_dir

depth_flat = depth.flatten()[:, None]
overlay = (
255.0
* colormaps.apply_depth_colormap(torch.from_numpy(depth_flat)).numpy()
)
overlay = 255.0 * colormaps.apply_depth_colormap(torch.from_numpy(depth_flat)).numpy()
overlay = overlay.reshape([H, W, 3])
input_image_path = input_images_dir / im_data.name
input_image = cv2.imread(str(input_image_path)) # type: ignore
Expand Down Expand Up @@ -700,9 +669,7 @@ def get_matching_summary(num_initial_frames: int, num_matched_frames: int) -> st
return "[bold green]COLMAP found poses for all images, CONGRATS!"
if match_ratio < 0.4:
result = f"[bold red]COLMAP only found poses for {num_matched_frames / num_initial_frames * 100:.2f}%"
result += (
" of the images. This is low.\nThis can be caused by a variety of reasons,"
)
result += " of the images. This is low.\nThis can be caused by a variety of reasons,"
result += " such poor scene coverage, blurry images, or large exposure changes."
return result
if match_ratio < 0.8:
Expand Down Expand Up @@ -735,20 +702,13 @@ def create_ply_from_colmap(
raise ValueError(f"Could not find points3D.txt or points3D.bin in {recon_dir}")

# Load point Positions
points3D = torch.from_numpy(
np.array([p.xyz for p in colmap_points.values()], dtype=np.float32)
)
points3D = torch.from_numpy(np.array([p.xyz for p in colmap_points.values()], dtype=np.float32))
if applied_transform is not None:
assert applied_transform.shape == (3, 4)
points3D = (
torch.einsum("ij,bj->bi", applied_transform[:3, :3], points3D)
+ applied_transform[:3, 3]
)
points3D = torch.einsum("ij,bj->bi", applied_transform[:3, :3], points3D) + applied_transform[:3, 3]

# Load point colours
points3D_rgb = torch.from_numpy(
np.array([p.rgb for p in colmap_points.values()], dtype=np.uint8)
)
points3D_rgb = torch.from_numpy(np.array([p.rgb for p in colmap_points.values()], dtype=np.uint8))

# write ply
with open(output_dir / filename, "w") as f:
Expand Down

0 comments on commit 000e344

Please sign in to comment.