Add ODM data processor

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

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)

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

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")],    
 ]