From 09561e42727c7c7e0e9ce065692a4f344ba70ed5 Mon Sep 17 00:00:00 2001 From: Piero Toffanin Date: Thu, 12 Oct 2023 16:54:51 -0400 Subject: [PATCH] Add ODM data processor --- README.md | 1 + docs/quickstart/custom_dataset.md | 29 +++++ nerfstudio/process_data/odm_utils.py | 153 +++++++++++++++++++++++++++ nerfstudio/scripts/process_data.py | 87 +++++++++++++++ 4 files changed, 270 insertions(+) create mode 100644 nerfstudio/process_data/odm_utils.py diff --git a/README.md b/README.md index a853192ad..3c7e827b1 100644 --- a/README.md +++ b/README.md @@ -250,6 +250,7 @@ Using an existing dataset is great, but likely you want to use your own data! We | 📱 [Record3D](https://docs.nerf.studio/quickstart/custom_dataset.html#record3d-capture) | IOS with LiDAR | [Record3D app](https://record3d.app/) | 🐇 | | 🖥 [Metashape](https://docs.nerf.studio/quickstart/custom_dataset.html#metashape) | Any | [Metashape](https://www.agisoft.com/) | 🐇 | | 🖥 [RealityCapture](https://docs.nerf.studio/quickstart/custom_dataset.html#realitycapture) | Any | [RealityCapture](https://www.capturingreality.com/realitycapture) | 🐇 | +| 🖥 [ODM](https://docs.nerf.studio/quickstart/custom_dataset.html#ODM) | Any | [ODM](https://github.com/OpenDroneMap/ODM) | 🐇 | | 🛠 [Custom](https://docs.nerf.studio/quickstart/data_conventions.html) | Any | Camera Poses | 🐇 | ## 5. Advanced Options diff --git a/docs/quickstart/custom_dataset.md b/docs/quickstart/custom_dataset.md index 723a97c91..14a89e4a2 100644 --- a/docs/quickstart/custom_dataset.md +++ b/docs/quickstart/custom_dataset.md @@ -21,6 +21,7 @@ We Currently support the following custom data types: | 📱 [Record3D](record3d) | IOS with LiDAR | [Record3D app](https://record3d.app/) | 🐇 | | 🖥 [Metashape](metashape) | Any | [Metashape](https://www.agisoft.com/) | 🐇 | | 🖥 [RealityCapture](realitycapture) | Any | [RealityCapture](https://www.capturingreality.com/realitycapture) | 🐇 | +| 🖥 [ODM](odm) | Any | [ODM](https://github.com/OpenDroneMap/ODM) | 🐇 | (images_and_video)= @@ -319,6 +320,34 @@ ns-process-data realitycapture --data {data directory} --csv {csv file} --output ns-train nerfacto --data {output directory} ``` +(odm)= + +## ODM + +All images/videos must be captured with the same camera. + +1. Process a dataset using [ODM](https://github.com/OpenDroneMap/ODM#quickstart) + +```bash +$ ls /path/to/dataset +images +odm_report +odm_orthophoto +... +``` + +2. Convert to nerfstudio format. + +```bash +ns-process-data odm --data /path/to/dataset --output-dir {output directory} +``` + +4. Train! + +```bash +ns-train nerfacto --data {output directory} +``` + (360_data)= ## 360 Data (Equirectangular) diff --git a/nerfstudio/process_data/odm_utils.py b/nerfstudio/process_data/odm_utils.py new file mode 100644 index 000000000..6a6547888 --- /dev/null +++ b/nerfstudio/process_data/odm_utils.py @@ -0,0 +1,153 @@ +# Copyright 2022 the Regents of the University of California, Nerfstudio Team and contributors. 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. + +"""Helper utils for processing ODM data into the nerfstudio format.""" + +import json +from pathlib import Path +from typing import Dict, List +import os +import sys +import math + +import numpy as np + +from nerfstudio.process_data.process_data_utils import CAMERA_MODELS +from nerfstudio.utils.rich_utils import CONSOLE + +def rodrigues_vec_to_rotation_mat(rodrigues_vec): + theta = np.linalg.norm(rodrigues_vec) + if theta < sys.float_info.epsilon: + rotation_mat = np.eye(3, dtype=float) + else: + r = rodrigues_vec / theta + I = np.eye(3, dtype=float) + r_rT = np.array([ + [r[0]*r[0], r[0]*r[1], r[0]*r[2]], + [r[1]*r[0], r[1]*r[1], r[1]*r[2]], + [r[2]*r[0], r[2]*r[1], r[2]*r[2]] + ]) + r_cross = np.array([ + [0, -r[2], r[1]], + [r[2], 0, -r[0]], + [-r[1], r[0], 0] + ]) + rotation_mat = math.cos(theta) * I + (1 - math.cos(theta)) * r_rT + math.sin(theta) * r_cross + return rotation_mat + +def cameras2nerfds( + image_filename_map: Dict[str, Path], + cameras_file: Path, + shots_file: Path, + output_dir: Path, + verbose: bool = False, +) -> List[str]: + """Convert ODM cameras into a nerfstudio dataset. + + Args: + image_filename_map: Mapping of original image filenames to their saved locations. + shots_file: Path to ODM's shots.geojson + output_dir: Path to the output directory. + verbose: Whether to print verbose output. + + Returns: + Summary of the conversion. + """ + + with open(cameras_file, "r", encoding="utf-8") as f: + cameras = json.loads(f.read()) + with open(shots_file, "r", encoding="utf-8") as f: + shots = json.loads(f.read()) + + camera_ids = list(cameras.keys()) + if len(camera_ids) > 1: + raise ValueError("Only one camera is supported") + camera_id = camera_ids[0] + camera = cameras[camera_id] + data = {} + if camera['projection_type'] in ['brown', 'perspective']: + data["camera_model"] = CAMERA_MODELS["perspective"].value + elif camera['projection_type'] in ['fisheye', 'fisheye_opencv']: + data["camera_model"] = CAMERA_MODELS["fisheye"].value + elif camera['projection_type'] in ['spherical', 'equirectangular']: + data["camera_model"] = CAMERA_MODELS["equirectangular"].value + else: + raise ValueError("Unsupported ODM camera model: " + data["camera_model"]) + + sensor_dict = {} + s = { + "w": int(camera["width"]), + "h": int(camera["height"]) + } + + s["fl_x"] = camera.get('focal_x', camera.get('focal')) * max(s["w"], s["h"]) + s["fl_y"] = camera.get('focal_y', camera.get('focal')) * max(s["w"], s["h"]) + + s["cx"] = camera["c_x"] + (s["w"] - 1.0) / 2.0 + s["cy"] = camera["c_y"] + (s["h"] - 1.0) / 2.0 + + for p in ["k1","k2","p1","p2","k3"]: + if p in camera: + s[p] = camera[p] + + sensor_dict[camera_id] = s + + shots = shots['features'] + shots_dict = {} + for shot in shots: + props = shot['properties'] + filename = props['filename'] + rotation = rodrigues_vec_to_rotation_mat(np.array(props['rotation']) * -1) + translation = np.array(props['translation']) + + m = np.eye(4) + m[:3, :3] = rotation + m[:3, 3] = translation + + name, ext = os.path.splitext(filename) + shots_dict[name] = m + + frames = [] + num_skipped = 0 + + + for fname in shots_dict: + transform = shots_dict[fname] + if fname not in image_filename_map: + num_skipped += 1 + continue + + frame = {} + frame["file_path"] = image_filename_map[fname].as_posix() + frame.update(sensor_dict[camera_id]) + + transform = transform[[2, 0, 1, 3], :] + transform[:, 1:3] *= -1 + frame["transform_matrix"] = transform.tolist() + frames.append(frame) + + data["frames"] = frames + + with open(output_dir / "transforms.json", "w", encoding="utf-8") as f: + json.dump(data, f, indent=4) + + summary = [] + if num_skipped == 1: + summary.append(f"{num_skipped} image skipped because it was missing its camera pose.") + if num_skipped > 1: + summary.append(f"{num_skipped} images were skipped because they were missing camera poses.") + + summary.append(f"Final dataset is {len(data['frames'])} frames.") + + return summary diff --git a/nerfstudio/scripts/process_data.py b/nerfstudio/scripts/process_data.py index 8b4151f75..aecaf2c16 100644 --- a/nerfstudio/scripts/process_data.py +++ b/nerfstudio/scripts/process_data.py @@ -32,6 +32,7 @@ from nerfstudio.process_data import ( process_data_utils, realitycapture_utils, record3d_utils, + odm_utils, ) from nerfstudio.process_data.colmap_converter_to_nerfstudio_dataset import BaseConverterToNerfstudioDataset from nerfstudio.process_data.images_to_nerfstudio_dataset import ImagesToNerfstudioDataset @@ -387,6 +388,91 @@ class ProcessRealityCapture(BaseConverterToNerfstudioDataset, _NoDefaultProcessR CONSOLE.rule() +@dataclass +class ProcessODM(BaseConverterToNerfstudioDataset): + """Process ODM data into a nerfstudio dataset. + + This script does the following: + + 1. Scales images to a specified size. + 2. Converts ODM poses into the nerfstudio format. + """ + + 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.""" + max_dataset_size: int = 600 + """Max number of images to train on. If the dataset has more, images will be sampled approximately evenly. If -1, + use all images.""" + + def main(self) -> None: + """Process images into a nerfstudio dataset.""" + + orig_images_dir = self.data / "images" + cameras_file = self.data / "cameras.json" + shots_file = self.data / "odm_report" / "shots.geojson" + + if not shots_file.exists: + raise ValueError(f"shots file {shots_file} doesn't exist") + if not shots_file.exists: + raise ValueError(f"cameras file {cameras_file} doesn't exist") + + if not orig_images_dir.exists: + raise ValueError(f"Images dir {orig_images_dir} doesn't exist") + + if self.eval_data is not None: + raise ValueError("Cannot use eval_data since cameras were already aligned with ODM.") + + self.output_dir.mkdir(parents=True, exist_ok=True) + image_dir = self.output_dir / "images" + image_dir.mkdir(parents=True, exist_ok=True) + + summary_log = [] + + # Copy images to output directory + image_filenames, num_orig_images = process_data_utils.get_image_filenames(orig_images_dir, self.max_dataset_size) + copied_image_paths = process_data_utils.copy_images_list( + image_filenames, + image_dir=image_dir, + verbose=self.verbose, + num_downscales=self.num_downscales, + ) + num_frames = len(copied_image_paths) + + copied_image_paths = [Path("images/" + copied_image_path.name) for copied_image_path in copied_image_paths] + original_names = [image_path.stem for image_path in image_filenames] + image_filename_map = dict(zip(original_names, copied_image_paths)) + + if self.max_dataset_size > 0 and num_frames != num_orig_images: + summary_log.append(f"Started with {num_frames} images out of {num_orig_images} total") + summary_log.append( + "To change the size of the dataset add the argument [yellow]--max_dataset_size[/yellow] to " + f"larger than the current value ({self.max_dataset_size}), or -1 to use all images." + ) + else: + summary_log.append(f"Started with {num_frames} images") + + # Save json + if num_frames == 0: + CONSOLE.print("[bold red]No images found, exiting") + sys.exit(1) + summary_log.extend( + odm_utils.cameras2nerfds( + image_filename_map=image_filename_map, + cameras_file=cameras_file, + shots_file=shots_file, + output_dir=self.output_dir, + verbose=self.verbose, + ) + ) + + CONSOLE.rule("[bold green]:tada: :tada: :tada: All DONE :tada: :tada: :tada:") + + for summary in summary_log: + CONSOLE.print(summary, justify="center") + CONSOLE.rule() + + Commands = Union[ Annotated[ImagesToNerfstudioDataset, tyro.conf.subcommand(name="images")], Annotated[VideoToNerfstudioDataset, tyro.conf.subcommand(name="video")], @@ -394,6 +480,7 @@ Commands = Union[ Annotated[ProcessMetashape, tyro.conf.subcommand(name="metashape")], Annotated[ProcessRealityCapture, tyro.conf.subcommand(name="realitycapture")], Annotated[ProcessRecord3D, tyro.conf.subcommand(name="record3d")], + Annotated[ProcessODM, tyro.conf.subcommand(name="odm")], ]