diff --git a/nerfactory/cameras/cameras.py b/nerfactory/cameras/cameras.py index 7e5415eb9..d30869114 100644 --- a/nerfactory/cameras/cameras.py +++ b/nerfactory/cameras/cameras.py @@ -36,6 +36,16 @@ class CameraType(Enum): FISHEYE = auto() +CAMERA_MODEL_TO_TYPE = { + "SIMPLE_PINHOLE": CameraType.PERSPECTIVE, + "PINHOLE": CameraType.PERSPECTIVE, + "SIMPLE_RADIAL": CameraType.PERSPECTIVE, + "RADIAL": CameraType.PERSPECTIVE, + "OPENCV": CameraType.PERSPECTIVE, + "OPENCV_FISHEYE": CameraType.FISHEYE, +} + + class Cameras: """Dataset inputs for the image dataset and the ray generator. diff --git a/nerfactory/datamanagers/dataparsers/nerfactory_parser.py b/nerfactory/datamanagers/dataparsers/nerfactory_parser.py index 68538fc05..50493a358 100644 --- a/nerfactory/datamanagers/dataparsers/nerfactory_parser.py +++ b/nerfactory/datamanagers/dataparsers/nerfactory_parser.py @@ -22,7 +22,7 @@ import numpy as np import torch from nerfactory.cameras import utils as camera_utils -from nerfactory.cameras.cameras import Cameras, CameraType +from nerfactory.cameras.cameras import CAMERA_MODEL_TO_TYPE, Cameras, CameraType from nerfactory.configs import base as cfg from nerfactory.datamanagers.dataparsers.base import DataParser from nerfactory.datamanagers.structs import DatasetInputs, SceneBounds @@ -77,10 +77,6 @@ class Nerfactory(DataParser): scale_factor = 1.0 / torch.max(torch.abs(poses[:, :3, 3])) poses[:, :3, 3] *= scale_factor - distortion_params = camera_utils.get_distortion_params( - k1=float(meta["k1"]), k2=float(meta["k2"]), p1=float(meta["p1"]), p2=float(meta["p2"]) - ) - # in x,y,z order # assumes that the scene is centered at the origin aabb_scale = self.config.scene_scale @@ -90,6 +86,20 @@ class Nerfactory(DataParser): ) ) + if "camera_model" in meta: + camera_type = CAMERA_MODEL_TO_TYPE[meta["camera_model"]] + else: + camera_type = CameraType.PERSPECTIVE + + distortion_params = camera_utils.get_distortion_params( + k1=float(meta["k1"]) if "k1" in meta else 0.0, + k2=float(meta["k2"]) if "k2" in meta else 0.0, + k3=float(meta["k3"]) if "k3" in meta else 0.0, + k4=float(meta["k4"]) if "k4" in meta else 0.0, + p1=float(meta["p1"]) if "p1" in meta else 0.0, + p2=float(meta["p2"]) if "p2" in meta else 0.0, + ) + cameras = Cameras( fx=float(meta["fl_x"]), fy=float(meta["fl_y"]), @@ -99,7 +109,7 @@ class Nerfactory(DataParser): height=int(meta["h"]), width=int(meta["w"]), camera_to_worlds=poses[:, :3, :4], - camera_type=CameraType.PERSPECTIVE, + camera_type=camera_type, ) cameras.rescale_output_resolution(scaling_factor=1.0 / self.config.downscale_factor) diff --git a/scripts/process_data.py b/scripts/process_data.py index 8d11ad438..8343b7e7c 100644 --- a/scripts/process_data.py +++ b/scripts/process_data.py @@ -17,7 +17,7 @@ from nerfactory.utils import colmap_utils CONSOLE = Console(width=120) -class CameraType(Enum): +class CameraModel(Enum): """Enum for camera types.""" OPENCV = "OPENCV" @@ -25,8 +25,8 @@ class CameraType(Enum): CAMERA_MODELS = { - "perspective": CameraType.OPENCV, - "fisheye": CameraType.OPENCV_FISHEYE, + "perspective": CameraModel.OPENCV, + "fisheye": CameraModel.OPENCV_FISHEYE, } @@ -143,6 +143,7 @@ def copy_images(data, image_dir, verbose) -> int: """ image_paths = sorted(data.glob("*")) for i, image_path in enumerate(image_paths): + i = i + 1 # 1-indexed if verbose: CONSOLE.log(f"Copying image {i + 1} of {len(image_paths)}...") shutil.copy(image_path, image_dir / f"frame_{i:05d}{image_path.suffix}") @@ -153,6 +154,8 @@ def copy_images(data, image_dir, verbose) -> int: def downscale_images(image_dir: Path, num_downscales: int, verbose: bool = False) -> None: """Downscales the images in the directory. Uses FFMPEG. + Assumes images are named frame_00001.png, frame_00002.png, etc. + Args: image_dir: Path to the directory containing the images. num_downscales: Number of times to downscale the images. Downscales by 2 each time. @@ -177,7 +180,7 @@ def downscale_images(image_dir: Path, num_downscales: int, verbose: bool = False def run_colmap( image_dir: Path, colmap_dir: Path, - camera_model: CameraType, + camera_model: CameraModel, gpu: bool = True, verbose: bool = False, ) -> None: @@ -250,13 +253,14 @@ def run_colmap( CONSOLE.log("[bold green]:tada: Done COLMAP bundle adjustment.") -def colmap_to_json(cameras_path: Path, images_path: Path, output_dir: Path) -> int: +def colmap_to_json(cameras_path: Path, images_path: Path, output_dir: Path, camera_model: CameraModel) -> int: """Converts COLMAP's cameras.bin and images.bin to a JSON file. Args: cameras_path: Path to the cameras.bin file. images_path: Path to the images.bin file. output_dir: Path to the output directory. + camera_model: Camera model used. Returns: The number of registered images. @@ -291,23 +295,41 @@ def colmap_to_json(cameras_path: Path, images_path: Path, output_dir: Path) -> i out = { "fl_x": float(camera_params[0]), "fl_y": float(camera_params[1]), - "k1": float(camera_params[4]), - "k2": float(camera_params[5]), - "p1": float(camera_params[6]), - "p2": float(camera_params[7]), "cx": float(camera_params[2]), "cy": float(camera_params[3]), "w": cameras[1].width, "h": cameras[1].height, - "frames": frames, + "camera_model": camera_model.value, } + if camera_model == CameraModel.OPENCV: + out.update( + { + "k1": float(camera_params[4]), + "k2": float(camera_params[5]), + "p1": float(camera_params[6]), + "p2": float(camera_params[7]), + } + ) + if camera_model == CameraModel.OPENCV_FISHEYE: + out.update( + { + "k1": float(camera_params[4]), + "k2": float(camera_params[5]), + "k3": float(camera_params[6]), + "k4": float(camera_params[7]), + } + ) + + out["frames"] = frames + with open(output_dir / "transforms.json", "w", encoding="utf-8") as f: json.dump(out, f, indent=4) return len(frames) +# pylint: disable=too-many-statements def main( data: Path, output_dir: Path, @@ -334,6 +356,7 @@ def main( camera_type: Camera model to use. num_downscales: 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. + skip_colmap: If True, skips COLMAP and generates transforms.json if possible. gpu: If True, use GPU. verbose: If True, print extra logging. """ @@ -384,24 +407,30 @@ def main( downscale_text = ", ".join(downscale_text[:-1]) + " and " + downscale_text[-1] summary_log.append(f"We downsampled the images by {downscale_text}") + camera_model = CAMERA_MODELS[camera_type] + colmap_dir = output_dir / "colmap" if not skip_colmap: - colmap_dir = output_dir / "colmap" colmap_dir.mkdir(parents=True, exist_ok=True) - camera_model = CAMERA_MODELS[camera_type] run_colmap(image_dir=image_dir, colmap_dir=colmap_dir, camera_model=camera_model, gpu=gpu, verbose=verbose) + 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=output_dir, + camera_model=camera_model, ) 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:") for summary in summary_log: CONSOLE.print(summary, justify="center") + CONSOLE.rule() if __name__ == "__main__":