Add colmap preprocessing script (#338)

* Add colmap preprocessing script

* Fix lint, add dataparser

nerfactory/configs/base.py

@@ -31,6 +31,7 @@ from nerfactory.dataloaders.data_parsers import (
     Friends,
     InstantNGP,
     Mipnerf360,
+    Nerfactory,
     Record3D,
 )
 from nerfactory.models.base import Model
@@ -169,6 +170,17 @@ class DataParserConfig(InstantiateConfig):
     _target: Type = DataParser
 
 
+@dataclass
+class NerfactoryDataParserConfig(DataParserConfig):
+    """Nerfactory dataset config"""
+
+    _target: Type = Nerfactory
+    data_directory: Path = Path("data/ours/posterv2")
+    scale_factor: float = 1.0
+    downscale_factor: int = 1
+    scene_scale: float = 0.33
+
+
 @dataclass
 class BlenderDataParserConfig(DataParserConfig):
     """Blender dataset config"""
@@ -202,7 +214,7 @@ class MipNerf360DataParserConfig(DataParserConfig):
 
 @dataclass
 class InstantNGPDataParserConfig(DataParserConfig):
-    """Mipnerf 360 dataset config"""
+    """Instant-NGP dataset config"""
 
     _target: Type = InstantNGP
     data_directory: Path = Path("data/ours/posterv2")

nerfactory/dataloaders/data_parsers.py

@@ -30,17 +30,17 @@ from scipy.spatial.transform import Rotation
 from nerfactory.cameras import utils as camera_utils
 from nerfactory.cameras.cameras import Cameras, CameraType
 from nerfactory.configs import base as cfg
-from nerfactory.dataloaders.colmap_utils import (
-    read_cameras_binary,
-    read_images_binary,
-    read_pointsTD_binary,
-)
 from nerfactory.dataloaders.structs import (
     DatasetInputs,
     PointCloud,
     SceneBounds,
     Semantics,
 )
+from nerfactory.utils.colmap_utils import (
+    read_cameras_binary,
+    read_images_binary,
+    read_pointsTD_binary,
+)
 from nerfactory.utils.colors import get_color
 from nerfactory.utils.io import (
     get_absolute_path,
@@ -129,6 +129,83 @@ class DataParser:
         return dataset_inputs
 
 
+@dataclass
+class Nerfactory(DataParser):
+    """Nerfactory Dataset
+
+    Args:
+        data_directory: Location of data
+        scale_factor: How much to scale the camera origins by.
+        downscale_factor: How much to downscale images. Defaults to 1.
+        scene_scale: How much to scale the scene. Defaults to 0.33
+    """
+
+    config: cfg.InstantNGPDataParserConfig
+
+    def _generate_dataset_inputs(self, split="train"):
+
+        abs_dir = get_absolute_path(self.config.data_directory)
+
+        meta = load_from_json(abs_dir / "transforms.json")
+        image_filenames = []
+        poses = []
+        num_skipped_image_filenames = 0
+        for frame in meta["frames"]:
+            fname = abs_dir / Path(frame["file_path"])
+            if not fname:
+                num_skipped_image_filenames += 1
+            else:
+                image_filenames.append(fname)
+                poses.append(np.array(frame["transform_matrix"]))
+        if num_skipped_image_filenames >= 0:
+            logging.info("Skipping %s files in dataset split %s.", num_skipped_image_filenames, split)
+        assert (
+            len(image_filenames) != 0
+        ), """
+        No image files found. 
+        You should check the file_paths in the transforms.json file to make sure they are correct.
+        """
+        poses = np.array(poses).astype(np.float32)
+        poses[:3, 3] *= self.config.scene_scale
+
+        camera_to_world = torch.from_numpy(poses[:, :3])  # camera to world transform
+
+        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 = 16
+        scene_bounds = SceneBounds(
+            aabb=torch.tensor(
+                [[-aabb_scale, -aabb_scale, -aabb_scale], [aabb_scale, aabb_scale, aabb_scale]], dtype=torch.float32
+            )
+        )
+
+        cameras = Cameras(
+            fx=float(meta["fl_x"]),
+            fy=float(meta["fl_y"]),
+            cx=float(meta["cx"]),
+            cy=float(meta["cy"]),
+            distortion_params=distortion_params,
+            height=int(meta["h"]),
+            width=int(meta["w"]),
+            camera_to_worlds=camera_to_world,
+            camera_type=CameraType.PERSPECTIVE,
+        )
+
+        cameras.rescale_output_resolution(scaling_factor=1.0 / self.config.downscale_factor)
+
+        dataset_inputs = DatasetInputs(
+            image_filenames=image_filenames,
+            downscale_factor=self.config.downscale_factor,
+            cameras=cameras,
+            scene_bounds=scene_bounds,
+        )
+        return dataset_inputs
+
+
 @dataclass
 class Blender(DataParser):
     """Blender Dataset
@@ -238,20 +315,11 @@ class InstantNGP(DataParser):
         poses = np.array(poses).astype(np.float32)
         poses[:3, 3] *= self.config.scene_scale
 
-        fx = float(meta["fl_x"])
-        fy = float(meta["fl_y"])
-        cx = float(meta["cx"])
-        cy = float(meta["cy"])
-        k1 = float(meta["k1"])
-        k2 = float(meta["k2"])
-        p1 = float(meta["p1"])
-        p2 = float(meta["p2"])
-        image_height = int(meta["h"])
-        image_width = int(meta["w"])
-
         camera_to_world = torch.from_numpy(poses[:, :3])  # camera to world transform
 
-        distortion_params = camera_utils.get_distortion_params(k1=k1, k2=k2, p1=p1, p2=p2)
+        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
@@ -263,13 +331,13 @@ class InstantNGP(DataParser):
         )
 
         cameras = Cameras(
-            fx=fx,
-            fy=fy,
-            cx=cx,
-            cy=cy,
+            fx=float(meta["fl_x"]),
+            fy=float(meta["fl_y"]),
+            cx=float(meta["cx"]),
+            cy=float(meta["cy"]),
             distortion_params=distortion_params,
-            height=image_height,
-            width=image_width,
+            height=int(meta["h"]),
+            width=int(meta["w"]),
             camera_to_worlds=camera_to_world,
             camera_type=CameraType.PERSPECTIVE,
         )

nerfactory/utils/colmap_utils.py

@@ -1,6 +1,6 @@
 # pylint: disable-all
 """
-Here we have code taken from COLMAP for parsing data in the COLMAP format.
+Here we have modified code taken from COLMAP for parsing data in the COLMAP format.
 Original file at https://github.com/colmap/colmap/blob/1a4d0bad2e90aa65ce997c9d1779518eaed998d5/scripts/python/read_write_model.py.
 """
 
@@ -84,24 +84,6 @@ def read_next_bytes(fid, num_bytes, format_char_sequence, endian_character="<"):
     return struct.unpack(endian_character + format_char_sequence, data)
 
 
-def write_next_bytes(fid, data, format_char_sequence, endian_character="<"):
-    """pack and write to a binary file.
-
-    Args:
-        fid:
-        data: data to send, if multiple elements are sent at the same time,
-            they should be encapsuled either in a list or a tuple
-        format_char_sequence: List of {c, e, f, d, h, H, i, I, l, L, q, Q}.
-            should be the same length as the data list or tuple
-        endian_character: Any of {@, =, <, >, !}
-    """
-    if isinstance(data, (list, tuple)):
-        bytes = struct.pack(endian_character + format_char_sequence, *data)
-    else:
-        bytes = struct.pack(endian_character + format_char_sequence, data)
-    fid.write(bytes)
-
-
 def read_cameras_text(path):
     """
     see: src/base/reconstruction.cc
@@ -151,42 +133,6 @@ def read_cameras_binary(path_to_model_file):
     return cameras
 
 
-def write_cameras_text(cameras, path):
-    """
-    see: src/base/reconstruction.cc
-        void Reconstruction::WriteCamerasText(const std::string& path)
-        void Reconstruction::ReadCamerasText(const std::string& path)
-    """
-    HEADER = (
-        "# Camera list with one line of data per camera:\n"
-        + "#   CAMERA_ID, MODEL, WIDTH, HEIGHT, PARAMS[]\n"
-        + f"# Number of cameras: {len(cameras)}\n"
-    )
-    with open(path, "w") as fid:
-        fid.write(HEADER)
-        for _, cam in cameras.items():
-            to_write = [cam.id, cam.model, cam.width, cam.height, *cam.params]
-            line = " ".join([str(elem) for elem in to_write])
-            fid.write(line + "\n")
-
-
-def write_cameras_binary(cameras, path_to_model_file):
-    """
-    see: src/base/reconstruction.cc
-        void Reconstruction::WriteCamerasBinary(const std::string& path)
-        void Reconstruction::ReadCamerasBinary(const std::string& path)
-    """
-    with open(path_to_model_file, "wb") as fid:
-        write_next_bytes(fid, len(cameras), "Q")
-        for _, cam in cameras.items():
-            model_id = CAMERA_MODEL_NAMES[cam.model].model_id
-            camera_properties = [cam.id, model_id, cam.width, cam.height]
-            write_next_bytes(fid, camera_properties, "iiQQ")
-            for p in cam.params:
-                write_next_bytes(fid, float(p), "d")
-    return cameras
-
-
 def read_images_text(path):
     """
     see: src/base/reconstruction.cc
@@ -258,57 +204,6 @@ def read_images_binary(path_to_model_file):
     return images
 
 
-def write_images_text(images, path):
-    """
-    see: src/base/reconstruction.cc
-        void Reconstruction::ReadImagesText(const std::string& path)
-        void Reconstruction::WriteImagesText(const std::string& path)
-    """
-    if len(images) == 0:
-        mean_observations = 0
-    else:
-        mean_observations = sum((len(img.point3D_ids) for _, img in images.items())) / len(images)
-    HEADER = (
-        "# Image list with two lines of data per image:\n"
-        + "#   IMAGE_ID, QW, QX, QY, QZ, TX, TY, TZ, CAMERA_ID, NAME\n"
-        + "#   POINTS2D[] as (X, Y, POINT3D_ID)\n"
-        + f"# Number of images: {len(images)}, mean observations per image: {mean_observations}\n"
-    )
-
-    with open(path, "w") as fid:
-        fid.write(HEADER)
-        for _, img in images.items():
-            image_header = [img.id, *img.qvec, *img.tvec, img.camera_id, img.name]
-            first_line = " ".join(map(str, image_header))
-            fid.write(first_line + "\n")
-
-            points_strings = []
-            for xy, point3D_id in zip(img.xys, img.point3D_ids):
-                points_strings.append(" ".join(map(str, [*xy, point3D_id])))
-            fid.write(" ".join(points_strings) + "\n")
-
-
-def write_images_binary(images, path_to_model_file):
-    """
-    see: src/base/reconstruction.cc
-        void Reconstruction::ReadImagesBinary(const std::string& path)
-        void Reconstruction::WriteImagesBinary(const std::string& path)
-    """
-    with open(path_to_model_file, "wb") as fid:
-        write_next_bytes(fid, len(images), "Q")
-        for _, img in images.items():
-            write_next_bytes(fid, img.id, "i")
-            write_next_bytes(fid, img.qvec.tolist(), "dddd")
-            write_next_bytes(fid, img.tvec.tolist(), "ddd")
-            write_next_bytes(fid, img.camera_id, "i")
-            for char in img.name:
-                write_next_bytes(fid, char.encode("utf-8"), "c")
-            write_next_bytes(fid, b"\x00", "c")
-            write_next_bytes(fid, len(img.point3D_ids), "Q")
-            for xy, p3d_id in zip(img.xys, img.point3D_ids):
-                write_next_bytes(fid, [*xy, p3d_id], "ddq")
-
-
 def read_pointsTD_text(path):
     """
     see: src/base/reconstruction.cc
@@ -361,52 +256,6 @@ def read_pointsTD_binary(path_to_model_file):
     return points3D
 
 
-def write_pointsTD_text(points3D, path):
-    """
-    see: src/base/reconstruction.cc
-        void Reconstruction::ReadPoints3DText(const std::string& path)
-        void Reconstruction::WritePoints3DText(const std::string& path)
-    """
-    if len(points3D) == 0:
-        mean_track_length = 0
-    else:
-        mean_track_length = sum((len(pt.image_ids) for _, pt in points3D.items())) / len(points3D)
-    HEADER = (
-        "# 3D point list with one line of data per point:\n"
-        + "#   POINT3D_ID, X, Y, Z, R, G, B, ERROR, TRACK[] as (IMAGE_ID, POINT2D_IDX)\n"
-        + f"# Number of points: {len(points3D)}, mean track length: {mean_track_length}\n"
-    )
-
-    with open(path, "w") as fid:
-        fid.write(HEADER)
-        for _, pt in points3D.items():
-            point_header = [pt.id, *pt.xyz, *pt.rgb, pt.error]
-            fid.write(" ".join(map(str, point_header)) + " ")
-            track_strings = []
-            for image_id, point2D in zip(pt.image_ids, pt.point2D_idxs):
-                track_strings.append(" ".join(map(str, [image_id, point2D])))
-            fid.write(" ".join(track_strings) + "\n")
-
-
-def write_pointsTD_binary(points3D, path_to_model_file):
-    """
-    see: src/base/reconstruction.cc
-        void Reconstruction::ReadPoints3DBinary(const std::string& path)
-        void Reconstruction::WritePoints3DBinary(const std::string& path)
-    """
-    with open(path_to_model_file, "wb") as fid:
-        write_next_bytes(fid, len(points3D), "Q")
-        for _, pt in points3D.items():
-            write_next_bytes(fid, pt.id, "Q")
-            write_next_bytes(fid, pt.xyz.tolist(), "ddd")
-            write_next_bytes(fid, pt.rgb.tolist(), "BBB")
-            write_next_bytes(fid, pt.error, "d")
-            track_length = pt.image_ids.shape[0]
-            write_next_bytes(fid, track_length, "Q")
-            for image_id, point2D_id in zip(pt.image_ids, pt.point2D_idxs):
-                write_next_bytes(fid, [image_id, point2D_id], "ii")
-
-
 def detect_model_format(path, ext):
     if (
         os.path.isfile(os.path.join(path, "cameras" + ext))
@@ -440,18 +289,6 @@ def read_model(path, ext=""):
     return cameras, images, points3D
 
 
-def write_model(cameras, images, points3D, path, ext=".bin"):
-    if ext == ".txt":
-        write_cameras_text(cameras, os.path.join(path, "cameras" + ext))
-        write_images_text(images, os.path.join(path, "images" + ext))
-        write_pointsTD_text(points3D, os.path.join(path, "points3D") + ext)
-    else:
-        write_cameras_binary(cameras, os.path.join(path, "cameras" + ext))
-        write_images_binary(images, os.path.join(path, "images" + ext))
-        write_pointsTD_binary(points3D, os.path.join(path, "points3D") + ext)
-    return cameras, images, points3D
-
-
 def qvec2rotmat(qvec):
     return np.array(
         [
@@ -492,25 +329,3 @@ def rotmat2qvec(R):
     if qvec[0] < 0:
         qvec *= -1
     return qvec
-
-
-def main():
-    parser = argparse.ArgumentParser(description="Read and write COLMAP binary and text models")
-    parser.add_argument("--input_model", help="path to input model folder")
-    parser.add_argument("--input_format", choices=[".bin", ".txt"], help="input model format", default="")
-    parser.add_argument("--output_model", help="path to output model folder")
-    parser.add_argument("--output_format", choices=[".bin", ".txt"], help="outut model format", default=".txt")
-    args = parser.parse_args()
-
-    cameras, images, points3D = read_model(path=args.input_model, ext=args.input_format)
-
-    print("num_cameras:", len(cameras))
-    print("num_images:", len(images))
-    print("num_points3D:", len(points3D))
-
-    if args.output_model is not None:
-        write_model(cameras, images, points3D, path=args.output_model, ext=args.output_format)
-
-
-if __name__ == "__main__":
-    main()