diff --git a/scripts/process_data.py b/scripts/process_data.py index 797ee4aa1..5ee8874e8 100755 --- a/scripts/process_data.py +++ b/scripts/process_data.py @@ -15,6 +15,7 @@ import appdirs import numpy as np import requests import tyro +import yaml from rich.console import Console from rich.progress import track from typing_extensions import Annotated @@ -484,6 +485,134 @@ def colmap_to_json(cameras_path: Path, images_path: Path, output_dir: Path, came return len(frames) +def run_opensfm(image_dir: Path, opensfm_dir: Path, camera_model: CameraModel, exe: Path, verbose=False) -> None: + # TODO cleanup config + config = { + "processes": 12, + "matching_bow_neighbors": 6, + "feature_process_size": 1280, + "triangulation_threshold": 0.004, + "triangulation_type": "ROBUST", + "min_track_length": 3, + "retriangulation_ratio": 2, + "bundle_new_points_ratio": 2, + } + with open(opensfm_dir / "config.yaml", "w") as config_file: + yaml.dump(config, config_file) + # TODO change cam type based on camera_model + # TODO get the width and height from images + cam_overrides = {"all": {"projection_type": "brown", "width": 540, "height": 960, "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-5: + # normalilze 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) + rotation = colmap_utils.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 from COLMAP's camera coordinate system to ours + # TODO determine if this is necessary for OpenSfM or not + c2w[0:3, 1:3] *= -1 + # 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"]) * cam["width"], + "fl_y": float(cam["focal_y"]) * cam["height"], + "cx": cam["width"] / 2.0 + float(cam["c_x"]) * cam["width"], + "cy": cam["height"] / 2.0 + float(cam["c_y"]) * cam["height"], + "w": cam["width"], + "h": 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) + + def get_insta360_filenames(data: Path) -> Tuple[Path, Path]: """Returns the filenames of the Insta360 videos from a single video file. @@ -609,10 +738,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 = 100 """Target number of frames to use for the dataset, results may not be exact.""" camera_type: Literal["perspective", "fisheye"] = "perspective" """Camera model to use.""" + sfm_method: Literal["colmap", "OpenSfM"] = "OpenSfM" + """Which sfm solver to use""" + opensfm_exec: 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.""" @@ -643,32 +776,51 @@ class ProcessVideo: # Downscale images summary_log.append(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) - 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, - ) - - # 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_to_json( - cameras_path=colmap_dir / "sparse" / "0" / "cameras.bin", - images_path=colmap_dir / "sparse" / "0" / "images.bin", - output_dir=self.output_dir, + 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, + ) + + # 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_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_exec is not None, "Please provide the path to the OpenSfM executable <...>/bin/opensfm" + # Run the OpenSfM solver + opensfm_dir = self.output_dir / "opensfm" + opensfm_dir.mkdir(parents=True, exist_ok=True) + run_opensfm(image_dir, opensfm_dir, CAMERA_MODELS[self.camera_type], self.opensfm_exec) + # Save transforms.json + if (opensfm_dir / "reconstruction.json").exists(): + with CONSOLE.status("[bold yellow]Saving results to transforms.json", spinner="balloon"): + 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") - 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:") diff --git a/scripts/viewer/view_dataset.py b/scripts/viewer/view_dataset.py index 29317515f..b43922ab1 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.utils.datasets 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