From d9b3347964a204058feb0cfba9eca281b1db986d Mon Sep 17 00:00:00 2001 From: akristoffersen Date: Fri, 2 Dec 2022 02:31:33 +0000 Subject: [PATCH] update. --- nerfstudio/cameras/cameras.py | 18 +- nerfstudio/data/pixel_samplers.py | 94 +++++++++ nerfstudio/fields/nerfacto_field.py | 2 +- nerfstudio/process_data/colmap_utils.py | 2 + nerfstudio/process_data/opensfm_utils.py | 194 ++++++++++++++++++ nerfstudio/process_data/process_data_utils.py | 17 +- scripts/process_data.py | 156 +++++++++----- scripts/viewer/view_dataset.py | 6 +- 8 files changed, 423 insertions(+), 66 deletions(-) create mode 100644 nerfstudio/process_data/opensfm_utils.py diff --git a/nerfstudio/cameras/cameras.py b/nerfstudio/cameras/cameras.py index 5057509c9..04893f5ec 100644 --- a/nerfstudio/cameras/cameras.py +++ b/nerfstudio/cameras/cameras.py @@ -36,6 +36,7 @@ class CameraType(Enum): PERSPECTIVE = auto() FISHEYE = auto() + EQUIRECTANGULAR = auto() CAMERA_MODEL_TO_TYPE = { @@ -45,6 +46,7 @@ CAMERA_MODEL_TO_TYPE = { "RADIAL": CameraType.PERSPECTIVE, "OPENCV": CameraType.PERSPECTIVE, "OPENCV_FISHEYE": CameraType.FISHEYE, + "EQUIRECTANGULAR": CameraType.EQUIRECTANGULAR, } @@ -248,7 +250,7 @@ class Cameras: image_coords = torch.stack(image_coords, dim=-1) + pixel_offset # stored as (y, x) coordinates return image_coords - def generate_rays( + def generate_rays( # pylint: disable=too-many-statements self, camera_indices: Union[TensorType["num_rays":...], int], coords: Optional[TensorType["num_rays":..., 2]] = None, @@ -277,8 +279,8 @@ class Cameras: coords = self.get_image_coords().to(self.device) assert coords is not None - y = coords[..., 0] # (..., 1) - x = coords[..., 1] # (..., 1) + y = coords[..., 0] + x = coords[..., 1] fx, fy = self.fx[camera_indices], self.fy[camera_indices] cx, cy = self.cx[camera_indices], self.cy[camera_indices] @@ -312,6 +314,16 @@ class Cameras: [coord_stack[..., 0] * sin_theta / theta, coord_stack[..., 1] * sin_theta / theta, -torch.cos(theta)], dim=-1, ) + elif self.camera_type[0] == CameraType.EQUIRECTANGULAR.value: + # u goes from -1 to 1 + # v goes from -1/2 to 1/2 (height is half of width) + theta = -torch.pi * coord_stack[..., 0] + phi = torch.pi * (0.5 - coord_stack[..., 1]) + # use spherical in local camera coordinates (+y up, x=0 and z<0 is theta=0) + directions_stack = torch.stack( + [-torch.sin(theta) * torch.sin(phi), torch.cos(phi), -torch.cos(theta) * torch.sin(phi)], + dim=-1, + ) else: raise ValueError(f"Camera type {CameraType(self.camera_type[0])} not supported.") diff --git a/nerfstudio/data/pixel_samplers.py b/nerfstudio/data/pixel_samplers.py index 6da0dff33..cec02176a 100644 --- a/nerfstudio/data/pixel_samplers.py +++ b/nerfstudio/data/pixel_samplers.py @@ -92,3 +92,97 @@ class PixelSampler: # pylint: disable=too-few-public-methods image_batch, self.num_rays_per_batch, keep_full_image=self.keep_full_image ) return pixel_batch + + +# TODO(kevinddchen): Implement equirect pixel sampler + +# def collate_image_dataset_batch_equirectangular(batch: Dict, num_rays_per_batch: int, keep_full_image: bool = False): +# """ +# Operates on a batch of images and samples pixels to use for generating rays. +# Returns a collated batch which is input to the Graph. +# It will sample only within the valid 'mask' if it's specified. + +# Args: +# batch: batch of images to sample from +# num_rays_per_batch: number of rays to sample per batch +# keep_full_image: whether or not to include a reference to the full image in returned batch +# """ +# device = batch["image"].device +# num_images, image_height, image_width, _ = batch["image"].shape + +# # only sample within the mask, if the mask is in the batch +# # if "mask" in batch: +# # nonzero_indices = torch.nonzero(batch["mask"][..., 0], as_tuple=False) +# # chosen_indices = random.sample(range(len(nonzero_indices)), k=num_rays_per_batch) +# # indices = nonzero_indices[chosen_indices] +# # else: +# # indices = torch.floor( +# # torch.rand((num_rays_per_batch, 3), device=device) +# # * torch.tensor([num_images, image_height, image_width], device=device) +# # ).long() + +# number_of_images = torch.rand(num_rays_per_batch, device=device) +# phi = torch.acos(1 - 2 * torch.rand(num_rays_per_batch, device=device)) / torch.pi +# theta = torch.rand(num_rays_per_batch, device=device) +# indices = torch.floor( +# torch.stack((number_of_images, phi, theta), dim=-1) +# * torch.tensor([num_images, image_height, image_width], device=device) +# ).long() + +# c, y, x = (i.flatten() for i in torch.split(indices, 1, dim=-1)) +# image = batch["image"][c, y, x] +# # mask, semantics = None, None +# # if "mask" in batch: +# # mask = batch["mask"][c, y, x] +# # if "semantics" in batch: +# # semantics = batch["semantics"][c, y, x] +# # assert image.shape == (num_rays_per_batch, 3), image.shape + +# # Needed to correct the random indices to their actual camera idx locations. +# local_indices = indices.clone() +# indices[:, 0] = batch["image_idx"][c] +# collated_batch = { +# "local_indices": local_indices, # local to the batch returned +# "indices": indices, # with the abs camera indices +# "image": image, +# } +# # if mask is not None: +# # collated_batch["mask"] = mask +# # if semantics is not None: +# # collated_batch["semantics"] = semantics + +# # if keep_full_image: +# # collated_batch["full_image"] = batch["image"] + +# return collated_batch + +# class EquirectangularSampler: # pylint: disable=too-few-public-methods +# """Samples 'pixel_batch's from 'image_batch's. + +# Args: +# num_rays_per_batch: number of rays to sample per batch +# keep_full_image: whether or not to include a reference to the full image in returned batch +# """ + +# def __init__(self, num_rays_per_batch: int, keep_full_image: bool = False) -> None: +# self.num_rays_per_batch = num_rays_per_batch +# self.keep_full_image = keep_full_image + +# def set_num_rays_per_batch(self, num_rays_per_batch: int): +# """Set the number of rays to sample per batch. + +# Args: +# num_rays_per_batch: number of rays to sample per batch +# """ +# self.num_rays_per_batch = num_rays_per_batch + +# def sample(self, image_batch: Dict): +# """Sample an image batch and return a pixel batch. + +# Args: +# image_batch: batch of images to sample from +# """ +# pixel_batch = collate_image_dataset_batch_equirectangular( +# image_batch, self.num_rays_per_batch, keep_full_image=self.keep_full_image +# ) +# return pixel_batch diff --git a/nerfstudio/fields/nerfacto_field.py b/nerfstudio/fields/nerfacto_field.py index e3273ec9d..ebb74bff4 100644 --- a/nerfstudio/fields/nerfacto_field.py +++ b/nerfstudio/fields/nerfacto_field.py @@ -115,7 +115,7 @@ class TCNNNerfactoField(Field): self.use_pred_normals = use_pred_normals num_levels = 16 - max_res = 1024 + max_res = 2048 base_res = 16 log2_hashmap_size = 19 features_per_level = 2 diff --git a/nerfstudio/process_data/colmap_utils.py b/nerfstudio/process_data/colmap_utils.py index c36cc7d5f..dc7b93ddc 100644 --- a/nerfstudio/process_data/colmap_utils.py +++ b/nerfstudio/process_data/colmap_utils.py @@ -46,6 +46,8 @@ from typing import Dict, Optional, Tuple import appdirs import numpy as np import requests +import yaml +from PIL import Image as pil_Image from rich.console import Console from rich.progress import track from typing_extensions import Literal diff --git a/nerfstudio/process_data/opensfm_utils.py b/nerfstudio/process_data/opensfm_utils.py new file mode 100644 index 000000000..d46ff58f0 --- /dev/null +++ b/nerfstudio/process_data/opensfm_utils.py @@ -0,0 +1,194 @@ +# Copyright (c) 2022, ETH Zurich and UNC Chapel Hill. +# All rights reserved. +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# * Redistributions of source code must retain the above copyright +# notice, this list of conditions and the following disclaimer. +# +# * Redistributions in binary form must reproduce the above copyright +# notice, this list of conditions and the following disclaimer in the +# documentation and/or other materials provided with the distribution. +# +# * Neither the name of ETH Zurich and UNC Chapel Hill nor the names of +# its contributors may be used to endorse or promote products derived +# from this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE +# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN +# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) +# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE +# POSSIBILITY OF SUCH DAMAGE. +# +# Author: Johannes L. Schoenberger (jsch-at-demuc-dot-de) + +import json +from pathlib import Path + +import numpy as np +import yaml +from PIL import Image +from rich.console import Console + +from nerfstudio.process_data.colmap_utils import qvec2rotmat +from nerfstudio.process_data.process_data_utils import CameraModel +from nerfstudio.utils.rich_utils import status +from nerfstudio.utils.scripts import run_command + +CONSOLE = Console(width=120) + + +def run_opensfm( + image_dir: Path, opensfm_dir: Path, camera_model: CameraModel, opensfm_install: Path, verbose=False +) -> None: + exe = opensfm_install / "bin" / "opensfm" + # TODO cleanup config + config = { + "processes": 12, + "matching_order_neighbors": 20, + "feature_process_size": 1024, + "triangulation_threshold": 0.006, + "triangulation_type": "ROBUST", + "min_track_length": 2, + "retriangulation_ratio": 1.5, + "bundle_new_points_ratio": 1.5, + } + with open(opensfm_dir / "config.yaml", "w") as config_file: + yaml.dump(config, config_file) + example_img = image_dir.glob("*").__next__() + pil_im = Image.open(example_img) + if camera_model == CameraModel.OPENCV: + proj_type = "brown" + elif camera_model == CameraModel.OPENCV_FISHEYE: + proj_type = "fisheye_opencv" + elif camera_model == CameraModel.EQUIRECTANGULAR: + proj_type = "equirectangular" + cam_overrides = { + "all": { + "projection_type": proj_type, + "width": pil_im.width, + "height": pil_im.height, + "focal_x": 0.85, + "focal_y": 0.85, + } + } + with open(opensfm_dir / "camera_models_overrides.json", "w") as cam_file: + json.dump(cam_overrides, cam_file) + + run_command(f"cp -r {image_dir} {opensfm_dir/'images'}") + metadata_cmd = [f"{exe} extract_metadata", f"{opensfm_dir}"] + metadata_cmd = " ".join(metadata_cmd) + with status(msg="[bold yellow]Extracting OpenSfM metadata...", spinner="moon", verbose=verbose): + run_command(metadata_cmd, verbose=verbose) + CONSOLE.log("[bold green]:tada: Done OpenSfM metadata.") + + feature_extractor_cmd = [f"{exe} detect_features", f"{opensfm_dir}"] + feature_extractor_cmd = " ".join(feature_extractor_cmd) + with status(msg="[bold yellow]Running OpenSfM feature extractor...", spinner="runner", verbose=verbose): + run_command(feature_extractor_cmd, verbose=verbose) + CONSOLE.log("[bold green]:tada: Done extracting image features.") + + matcher_cmd = [f"{exe} match_features", f"{opensfm_dir}"] + matcher_cmd = " ".join(matcher_cmd) + with status(msg="[bold yellow]Running OpenSfM feature matcher...", spinner="runner", verbose=verbose): + run_command(matcher_cmd, verbose=verbose) + CONSOLE.log("[bold green]:tada: Done matching image features.") + + tracks_cmd = [f"{exe} create_tracks", f"{opensfm_dir}"] + tracks_cmd = " ".join(tracks_cmd) + with status(msg="[bold yellow]Merging OpenSfM features into tracked points...", spinner="dqpb", verbose=verbose): + run_command(tracks_cmd, verbose=verbose) + CONSOLE.log("[bold green]:tada: Done merging features.") + + recon_cmd = [f"{exe} reconstruct", f"{opensfm_dir}"] + recon_cmd = " ".join(recon_cmd) + with status( + msg="[bold yellow]Running OpenSfM bundle adjustment... (this can take 10-15 minutes)", + spinner="dqpb", + verbose=verbose, + ): + run_command(recon_cmd, verbose=verbose) + CONSOLE.log("[bold green]:tada: Done running bundle adjustment.") + + +def opensfm_to_json(reconstruction_path: Path, output_dir: Path, camera_model: CameraModel) -> int: + sfm_out = (json.load(open(reconstruction_path, "r")))[0] + cams = sfm_out["cameras"] + shots = sfm_out["shots"] + frames = [] + for shot_name, shot_info in shots.items(): + # enumerates through images in the capture + aa_vec = np.array(shot_info["rotation"]) # 3D axis angle repr + angle = np.linalg.norm(aa_vec) + if angle > 1e-8: + # normalize the axis-angle repr if angle is large enough + aa_vec /= angle + qx = aa_vec[0] * np.sin(angle / 2) + qy = aa_vec[1] * np.sin(angle / 2) + qz = aa_vec[2] * np.sin(angle / 2) + qw = np.cos(angle / 2) + else: + qx, qy, qz, qw = 0.0, 0.0, 0.0, 1.0 + rotation = qvec2rotmat(np.array([qw, qx, qy, qz])) + translation = np.array(shot_info["translation"]).reshape(3, 1) + w2c = np.concatenate([rotation, translation], 1) + w2c = np.concatenate([w2c, np.array([[0, 0, 0, 1]])], 0) + c2w = np.linalg.inv(w2c) + # Convert camera +z forwards (OpenSfM) convention to -z forwards (NeRF) convention + # Equivalent to right-multiplication by diag([1, -1, -1, 1]) + c2w[0:3, 1:3] *= -1 + # Rotation around global z-axis by -90 degrees + c2w = c2w[np.array([1, 0, 2, 3]), :] + c2w[2, :] *= -1 + + name = Path(f"./images/{shot_name}") + + frame = { + "file_path": str(name), + "transform_matrix": c2w.tolist(), + } + frames.append(frame) + # For now just assume it's all the same camera + cam = cams[list(cams.keys())[0]] + out = { + "fl_x": float(cam["focal_x"]) if "focal_x" in cam else float(cam["height"]), + "fl_y": float(cam["focal_y"]) if "focal_y" in cam else float(cam["height"]), + "cx": float(cam["c_x"]) if "c_x" in cam else 0.5 * cam["width"], + "cy": float(cam["c_y"]) if "c_y" in cam else 0.5 * cam["height"], + "w": int(cam["width"]), + "h": int(cam["height"]), + "camera_model": camera_model.value, + } + if camera_model == CameraModel.OPENCV: + out.update( + { + "k1": float(cam["k1"]), + "k2": float(cam["k2"]), + "p1": float(cam["p1"]), + "p2": float(cam["p2"]), + } + ) + if camera_model == CameraModel.OPENCV_FISHEYE: + # TODO: find the opensfm camera model that uses these params + out.update( + { + "k1": float(cam["k1"]), + "k2": float(cam["k2"]), + "k3": float(cam["k3"]), + "k4": float(cam["k4"]), + } + ) + + out["frames"] = frames + + with open(output_dir / "transforms.json", "w", encoding="utf-8") as f: + json.dump(out, f, indent=4) + + return len(frames) diff --git a/nerfstudio/process_data/process_data_utils.py b/nerfstudio/process_data/process_data_utils.py index 106d87f77..cd814b6de 100644 --- a/nerfstudio/process_data/process_data_utils.py +++ b/nerfstudio/process_data/process_data_utils.py @@ -33,11 +33,13 @@ class CameraModel(Enum): OPENCV = "OPENCV" OPENCV_FISHEYE = "OPENCV_FISHEYE" + EQUIRECTANGULAR = "EQUIRECTANGULAR" CAMERA_MODELS = { "perspective": CameraModel.OPENCV, "fisheye": CameraModel.OPENCV_FISHEYE, + "equirectangular": CameraModel.EQUIRECTANGULAR, } @@ -108,7 +110,11 @@ def convert_video_to_images( def copy_images_list( - image_paths: List[Path], image_dir: Path, crop_border_pixels: Optional[int] = None, verbose: bool = False + image_paths: List[Path], + image_dir: Path, + crop_border_pixels: Optional[int] = None, + verbose: bool = False, + rename: bool = True, ) -> List[Path]: """Copy all images in a list of Paths. Useful for filtering from a directory. Args: @@ -131,7 +137,10 @@ def copy_images_list( for idx, image_path in enumerate(image_paths): if verbose: CONSOLE.log(f"Copying image {idx + 1} of {len(image_paths)}...") - copied_image_path = image_dir / f"frame_{idx + 1:05d}{image_path.suffix}" + if rename: + copied_image_path = image_dir / f"frame_{idx + 1:05d}{image_path.suffix}" + else: + copied_image_path = image_dir / image_path.name shutil.copy(image_path, copied_image_path) copied_image_paths.append(copied_image_path) @@ -152,7 +161,7 @@ def copy_images_list( return copied_image_paths -def copy_images(data: Path, image_dir: Path, verbose) -> int: +def copy_images(data: Path, image_dir: Path, verbose, rename: bool = True) -> int: """Copy images from a directory to a new directory. Args: @@ -166,7 +175,7 @@ def copy_images(data: Path, image_dir: Path, verbose) -> int: allowed_exts = [".jpg", ".jpeg", ".png", ".tif", ".tiff"] image_paths = sorted([p for p in data.glob("[!.]*") if p.suffix.lower() in allowed_exts]) - num_frames = len(copy_images_list(image_paths, image_dir, verbose)) + num_frames = len(copy_images_list(image_paths, image_dir, verbose, rename=rename)) return num_frames diff --git a/scripts/process_data.py b/scripts/process_data.py index ed0b0d54c..aa2ca18fa 100755 --- a/scripts/process_data.py +++ b/scripts/process_data.py @@ -6,7 +6,7 @@ import sys import zipfile from dataclasses import dataclass from pathlib import Path -from typing import Union +from typing import Optional, Union import numpy as np import tyro @@ -17,6 +17,7 @@ from nerfstudio.process_data import ( colmap_utils, insta360_utils, metashape_utils, + opensfm_utils, polycam_utils, process_data_utils, record3d_utils, @@ -41,8 +42,12 @@ class ProcessImages: """Path the data, either a video file or a directory of images.""" output_dir: Path """Path to the output directory.""" - camera_type: Literal["perspective", "fisheye"] = "perspective" + camera_type: Literal["perspective", "fisheye", "equirectangular"] = "perspective" """Camera model to use.""" + sfm_method: Literal["colmap", "OpenSfM"] = "colmap" + """Which sfm solver to use""" + opensfm_dir: Optional[Path] = None + """If sfm_method is OpenSfM, specify the executable to use here""" 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.""" @@ -70,16 +75,19 @@ class ProcessImages: summary_log = [] # Copy images to output directory - num_frames = process_data_utils.copy_images(self.data, image_dir=image_dir, verbose=self.verbose) + + rename = self.sfm_method != "OpenSfM" + num_frames = process_data_utils.copy_images(self.data, image_dir=image_dir, verbose=self.verbose, rename=rename) summary_log.append(f"Starting with {num_frames} images") # Downscale images summary_log.append(process_data_utils.downscale_images(image_dir, self.num_downscales, verbose=self.verbose)) - # Run COLMAP - colmap_dir = self.output_dir / "colmap" - if not self.skip_colmap: - colmap_dir.mkdir(parents=True, exist_ok=True) + if self.sfm_method == "colmap": + # Run COLMAP + 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, @@ -91,19 +99,37 @@ class ProcessImages: colmap_cmd=self.colmap_cmd, ) - # Save transforms.json - if (colmap_dir / "sparse" / "0" / "cameras.bin").exists(): - with CONSOLE.status("[bold yellow]Saving results to transforms.json", spinner="balloon"): - num_matched_frames = colmap_utils.colmap_to_json( - cameras_path=colmap_dir / "sparse" / "0" / "cameras.bin", - images_path=colmap_dir / "sparse" / "0" / "images.bin", - output_dir=self.output_dir, - camera_model=CAMERA_MODELS[self.camera_type], + # Save transforms.json + if (colmap_dir / "sparse" / "0" / "cameras.bin").exists(): + with CONSOLE.status("[bold yellow]Saving results to transforms.json", spinner="balloon"): + num_matched_frames = colmap_utils.colmap_to_json( + cameras_path=colmap_dir / "sparse" / "0" / "cameras.bin", + images_path=colmap_dir / "sparse" / "0" / "images.bin", + output_dir=self.output_dir, + camera_model=CAMERA_MODELS[self.camera_type], + ) + summary_log.append(f"Colmap matched {num_matched_frames} images") + else: + CONSOLE.log( + "[bold yellow]Warning: could not find existing COLMAP results. Not generating transforms.json" + ) + + elif self.sfm_method == "OpenSfM": + assert self.opensfm_dir is not None, "Please provide the path to the OpenSfM directory <...>/bin/opensfm" + # Run the OpenSfM solver + opensfm_dir = self.output_dir / "opensfm" + opensfm_dir.mkdir(parents=True, exist_ok=True) + opensfm_utils.run_opensfm(image_dir, opensfm_dir, CAMERA_MODELS[self.camera_type], self.opensfm_dir) + # Save transforms.json + if (opensfm_dir / "reconstruction.json").exists(): + with CONSOLE.status("[bold yellow]Saving results to transforms.json", spinner="balloon"): + opensfm_utils.opensfm_to_json( + opensfm_dir / "reconstruction.json", self.output_dir, CAMERA_MODELS[self.camera_type] + ) + else: + CONSOLE.log( + "[bold yellow]Warning: could not find existing OpenSfM results. Not generating transforms.json" ) - summary_log.append(f"Colmap matched {num_matched_frames} images") - 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") CONSOLE.rule("[bold green]:tada: :tada: :tada: All DONE :tada: :tada: :tada:") @@ -127,10 +153,14 @@ class ProcessVideo: """Path the data, either a video file or a directory of images.""" output_dir: Path """Path to the output directory.""" - num_frames_target: int = 300 + num_frames_target: int = 150 """Target number of frames to use for the dataset, results may not be exact.""" - camera_type: Literal["perspective", "fisheye"] = "perspective" + camera_type: Literal["perspective", "fisheye", "equirectangular"] = "perspective" """Camera model to use.""" + sfm_method: Literal["colmap", "OpenSfM"] = "OpenSfM" + """Which sfm solver to use""" + opensfm_dir: Optional[Path] = None + """If sfm_method is OpenSfM, specify the executable to use here""" 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.""" @@ -163,10 +193,11 @@ class ProcessVideo: # Downscale images summary_log.append(process_data_utils.downscale_images(image_dir, self.num_downscales, verbose=self.verbose)) - # Run Colmap - colmap_dir = self.output_dir / "colmap" - if not self.skip_colmap: - colmap_dir.mkdir(parents=True, exist_ok=True) + if self.sfm_method == "colmap": + # Run Colmap + 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, @@ -178,19 +209,37 @@ class ProcessVideo: colmap_cmd=self.colmap_cmd, ) - # Save transforms.json - if (colmap_dir / "sparse" / "0" / "cameras.bin").exists(): - with CONSOLE.status("[bold yellow]Saving results to transforms.json", spinner="balloon"): - num_matched_frames = colmap_utils.colmap_to_json( - cameras_path=colmap_dir / "sparse" / "0" / "cameras.bin", - images_path=colmap_dir / "sparse" / "0" / "images.bin", - output_dir=self.output_dir, - camera_model=CAMERA_MODELS[self.camera_type], + # Save transforms.json + if (colmap_dir / "sparse" / "0" / "cameras.bin").exists(): + with CONSOLE.status("[bold yellow]Saving results to transforms.json", spinner="balloon"): + num_matched_frames = colmap_utils.colmap_to_json( + cameras_path=colmap_dir / "sparse" / "0" / "cameras.bin", + images_path=colmap_dir / "sparse" / "0" / "images.bin", + output_dir=self.output_dir, + camera_model=CAMERA_MODELS[self.camera_type], + ) + summary_log.append(f"Colmap matched {num_matched_frames} images") + else: + CONSOLE.log( + "[bold yellow]Warning: could not find existing COLMAP results. Not generating transforms.json" + ) + + elif self.sfm_method == "OpenSfM": + assert self.opensfm_dir is not None, "Please provide the path to the OpenSfM directory <...>/bin/opensfm" + # Run the OpenSfM solver + opensfm_dir = self.output_dir / "opensfm" + opensfm_dir.mkdir(parents=True, exist_ok=True) + opensfm_utils.run_opensfm(image_dir, opensfm_dir, CAMERA_MODELS[self.camera_type], self.opensfm_dir) + # Save transforms.json + if (opensfm_dir / "reconstruction.json").exists(): + with CONSOLE.status("[bold yellow]Saving results to transforms.json", spinner="balloon"): + opensfm_utils.opensfm_to_json( + opensfm_dir / "reconstruction.json", self.output_dir, CAMERA_MODELS[self.camera_type] + ) + else: + CONSOLE.log( + "[bold yellow]Warning: could not find existing OpenSfM results. Not generating transforms.json" ) - summary_log.append(f"Colmap matched {num_matched_frames} images") - summary_log.append(colmap_utils.get_matching_summary(num_extracted_frames, num_matched_frames)) - else: - CONSOLE.log("[bold yellow]Warning: could not find existing COLMAP results. Not generating transforms.json") CONSOLE.rule("[bold green]:tada: :tada: :tada: All DONE :tada: :tada: :tada:") @@ -259,31 +308,28 @@ class ProcessInsta360: summary_log.append(process_data_utils.downscale_images(image_dir, self.num_downscales, verbose=self.verbose)) # Run Colmap - colmap_dir = self.output_dir / "colmap" + opensfm_dir = self.output_dir / "opensfm" if not self.skip_colmap: - colmap_dir.mkdir(parents=True, exist_ok=True) + opensfm_dir.mkdir(parents=True, exist_ok=True) - colmap_utils.run_colmap( + opensfm_utils.run_opensfm( image_dir=image_dir, - colmap_dir=colmap_dir, - camera_model=CAMERA_MODELS["fisheye"], - gpu=self.gpu, + opensfm_dir=opensfm_dir, + camera_model=CAMERA_MODELS["equirectangular"], + opensfm_install=Path(""), verbose=self.verbose, - matching_method=self.matching_method, - colmap_cmd=self.colmap_cmd, ) # Save transforms.json - if (colmap_dir / "sparse" / "0" / "cameras.bin").exists(): - with CONSOLE.status("[bold yellow]Saving results to transforms.json", spinner="balloon"): - num_matched_frames = colmap_utils.colmap_to_json( - cameras_path=colmap_dir / "sparse" / "0" / "cameras.bin", - images_path=colmap_dir / "sparse" / "0" / "images.bin", - output_dir=self.output_dir, - camera_model=CAMERA_MODELS["fisheye"], - ) - summary_log.append(f"Colmap matched {num_matched_frames} images") - summary_log.append(colmap_utils.get_matching_summary(num_extracted_frames, num_matched_frames)) + if (opensfm_dir / "sparse" / "0" / "cameras.bin").exists(): + # with CONSOLE.status("[bold yellow]Saving results to transforms.json", spinner="balloon"): + # num_matched_frames = opensfm_utils.opensfm_to_json( + # output_dir=self.output_dir, + # camera_model=CAMERA_MODELS["equirectangular"], + # ) + # summary_log.append(f"Colmap matched {num_matched_frames} images") + # summary_log.append(colmap_utils.get_matching_summary(num_extracted_frames, num_matched_frames)) + pass else: CONSOLE.log("[bold yellow]Warning: could not find existing COLMAP results. Not generating transforms.json") diff --git a/scripts/viewer/view_dataset.py b/scripts/viewer/view_dataset.py index 68b54e42d..1f5aa2201 100644 --- a/scripts/viewer/view_dataset.py +++ b/scripts/viewer/view_dataset.py @@ -13,7 +13,7 @@ from pathlib import Path import torch import tyro -from nerfstudio.configs import method_configs as cfg +from nerfstudio.configs.base_config import ViewerConfig from nerfstudio.data.datamanagers import AnnotatedDataParserUnion from nerfstudio.data.datasets.base_dataset import InputDataset from nerfstudio.viewer.server import viewer_utils @@ -26,7 +26,7 @@ torch.backends.cudnn.benchmark = True # type: ignore def main( dataparser: AnnotatedDataParserUnion, - viewer: cfg.ViewerConfig, + viewer: ViewerConfig, log_base_dir: Path = Path("/tmp/nerfstudio_viewer_logs"), ) -> None: """Main function.""" @@ -34,7 +34,7 @@ def main( viewer, log_filename=log_base_dir / viewer.relative_log_filename, ) - dataset = InputDataset(dataparser.setup().get_dataset_inputs(split="train")) + dataset = InputDataset(dataparser.setup().get_dataparser_outputs(split="train")) viewer_state.init_scene(dataset=dataset, start_train=False) CONSOLE.log("Please refresh and load page at: %s", viewer_state.viewer_url) time.sleep(30) # allowing time to refresh page