diff --git a/nerfstudio/process_data/metashape_utils.py b/nerfstudio/process_data/metashape_utils.py new file mode 100644 index 000000000..b2e820eee --- /dev/null +++ b/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 diff --git a/scripts/process_data.py b/scripts/process_data.py index e726084c3..ed0b0d54c 100755 --- a/scripts/process_data.py +++ b/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")], ]