Add fisheye support to data processing (#369)

add fisheye support to data processing

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.
 

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)

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__":