Add Metashape support (#989)

Add metashape support

nerfstudio/process_data/metashape_utils.py

@@ -0,0 +1,113 @@
+# Copyright 2022 The Nerfstudio Team. 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 metashape data into the nerfstudio format."""
+
+import json
+import xml.etree.ElementTree as ET
+from pathlib import Path
+from typing import Dict
+
+import numpy as np
+from rich.console import Console
+
+from nerfstudio.process_data.process_data_utils import CAMERA_MODELS
+
+CONSOLE = Console(width=120)
+
+
+def _find_distortion_param(calib_xml: ET.Element, param_name: str):
+    param = calib_xml.find(param_name)
+    if param is not None:
+        return float(param.text)  # type: ignore
+    return 0.0
+
+
+def metashape_to_json(
+    image_filename_map: Dict[str, Path],
+    xml_filename: Path,
+    output_dir: Path,
+) -> str:
+    """Convert Metashape data into a nerfstudio dataset.
+
+    Args:
+        image_filename_map: Mapping of original image filenames to their saved locations.
+        xml_filename: Path to the metashap cameras xml file.
+        output_dir: Path to the output directory.
+
+    Returns:
+        Summary of the conversion.
+    """
+
+    xml_tree = ET.parse(xml_filename)
+    root = xml_tree.getroot()
+    chunk = root[0]
+    sensors = chunk.find("sensors")
+
+    # TODO Add support for per-frame intrinsics
+    if sensors is None or len(sensors) != 1:
+        raise ValueError("Only one sensor is supported for now")
+
+    sensor = sensors.find("sensor")
+
+    data = {}
+
+    assert sensor is not None, "Sensor not found in Metashape XML"
+    resolution = sensor.find("resolution")
+    assert resolution is not None, "Resolution not found in Metashape xml"
+    data["w"] = int(resolution.get("width"))  # type: ignore
+    data["h"] = int(resolution.get("height"))  # type: ignore
+
+    calib = sensor.find("calibration")
+    assert calib is not None, "Calibration not found in Metashape xml"
+    data["fl_x"] = float(calib.find("f").text)  # type: ignore
+    data["fl_y"] = float(calib.find("f").text)  # type: ignore
+    data["cx"] = float(calib.find("cx").text) + data["w"] / 2.0  # type: ignore
+    data["cy"] = float(calib.find("cy").text) + data["h"] / 2.0  # type: ignore
+
+    data["k1"] = _find_distortion_param(calib, "k1")
+    data["k2"] = _find_distortion_param(calib, "k2")
+    data["k3"] = _find_distortion_param(calib, "k3")
+    data["k4"] = _find_distortion_param(calib, "k4")
+    data["p1"] = _find_distortion_param(calib, "p1")
+    data["p2"] = _find_distortion_param(calib, "p2")
+
+    data["camera_model"] = CAMERA_MODELS["perspective"].value
+
+    frames = []
+    cameras = chunk.find("cameras")
+    assert cameras is not None, "Cameras not found in Metashape xml"
+    for camera in cameras:
+        frame = {}
+        frame["file_path"] = image_filename_map[camera.get("label")].as_posix()  # type: ignore
+        t = [float(x) for x in camera.find("transform").text.split()]  # type: ignore
+        transform = np.array(
+            [
+                [t[8], -t[9], -t[10], t[11]],
+                [t[0], -t[1], -t[2], t[3]],
+                [t[4], -t[5], -t[6], t[7]],
+                [t[12], -t[13], -t[14], t[15]],
+            ]
+        )
+        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 = f"Final dataset is {len(data['frames'])} frames."
+
+    return summary

scripts/process_data.py

@@ -16,6 +16,7 @@ from typing_extensions import Annotated, Literal
 from nerfstudio.process_data import (
     colmap_utils,
     insta360_utils,
+    metashape_utils,
     polycam_utils,
     process_data_utils,
     record3d_utils,
@@ -487,10 +488,102 @@ class ProcessPolycam:
         CONSOLE.rule()
 
 
+@dataclass
+class ProcessMetashape:
+    """Process Metashape data into a nerfstudio dataset.
+
+    This script assumes that cameras have been aligned using Metashape. After alignment, it is necessary to export the
+    camera poses as a `.xml` file. This option can be found under `File > Export > Export Cameras`.
+
+    This script does the following:
+
+    1. Scales images to a specified size.
+    2. Converts Metashape poses into the nerfstudio format.
+    """
+
+    data: Path
+    """Path to a folder of images."""
+    xml: Path
+    """Path to the Metashape xml file."""
+    output_dir: Path
+    """Path to the output directory."""
+    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."""
+    verbose: bool = False
+    """If True, print extra logging."""
+
+    def main(self) -> None:
+        """Process images into a nerfstudio dataset."""
+
+        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 = []
+        for f in self.data.iterdir():
+            image_filenames.append(f)
+        image_filenames = sorted(image_filenames, key=lambda fn: fn.stem)
+        num_images = len(image_filenames)
+        idx = np.arange(num_images)
+        if self.max_dataset_size != -1 and num_images > self.max_dataset_size:
+            idx = np.round(np.linspace(0, num_images - 1, self.max_dataset_size)).astype(int)
+
+        image_filenames = list(np.array(image_filenames)[idx])
+        # Copy images to output directory
+        copied_image_paths = process_data_utils.copy_images_list(
+            image_filenames,
+            image_dir=image_dir,
+            verbose=self.verbose,
+        )
+        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_images:
+            summary_log.append(f"Started with {num_frames} images out of {num_images} total")
+            summary_log.append(
+                "To change the size of the dataset add the argument --max_dataset_size to larger than the "
+                f"current value ({self.max_dataset_size}), or -1 to use all images."
+            )
+        else:
+            summary_log.append(f"Started with {num_frames} images")
+
+        # Downscale images
+        summary_log.append(process_data_utils.downscale_images(image_dir, self.num_downscales, verbose=self.verbose))
+
+        # Save json
+        if num_frames == 0:
+            CONSOLE.print("[bold red]No images found, exiting")
+            sys.exit(1)
+        summary_log.append(
+            metashape_utils.metashape_to_json(
+                image_filename_map=image_filename_map,
+                xml_filename=self.xml,
+                output_dir=self.output_dir,
+            )
+        )
+
+        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[ProcessImages, tyro.conf.subcommand(name="images")],
     Annotated[ProcessVideo, tyro.conf.subcommand(name="video")],
     Annotated[ProcessPolycam, tyro.conf.subcommand(name="polycam")],
+    Annotated[ProcessMetashape, tyro.conf.subcommand(name="metashape")],
     Annotated[ProcessInsta360, tyro.conf.subcommand(name="insta360")],
     Annotated[ProcessRecord3D, tyro.conf.subcommand(name="record3d")],
 ]