Add a COLMAP export format to polyform (closes #9)

README.md

@@ -50,6 +50,20 @@ python3 -m polyform.convert <path-to-data-folder> --format ingp
 
 Note: you may need to tweak the `scale` parameter in the transforms.json file to get the best results. 
 
+### COLMAP
+
+[COLMAP](https://colmap.github.io/) is a widely-used Structure-from-Motion / Multi-View Stereo pipeline, and its text-based sparse model format (`cameras.txt` / `images.txt` / `points3D.txt`) is also the expected input format for many downstream tools, notably the reference [3D Gaussian Splatting](https://github.com/graphdeco-inria/gaussian-splatting) implementation. You can convert from Polycam's data format to a COLMAP text model by running:
+
+```
+python3 -m polyform.convert <path-to-data-folder> --format colmap
+```
+
+This writes `cameras.txt`, `images.txt` and `points3D.txt` into a `colmap_text` folder at the root of the data folder. Since Polycam's camera poses are already globally optimized (see note above), this lets you skip COLMAP's own (often slow, and occasionally failure-prone on textureless indoor scenes) feature-matching + SfM step entirely, and go straight to tools that consume a COLMAP model.
+
+A couple of things to know:
+- `points3D.txt` is written empty, since Polycam's raw export doesn't include COLMAP-style per-image 2D/3D keypoint correspondences. This is fine for tools that only need camera poses (e.g. 3D Gaussian Splatting can initialize from a random point cloud when `points3D.txt` is empty), but means COLMAP itself won't have a sparse point cloud to display until you triangulate one.
+- By default every image gets its own COLMAP camera entry, since Polycam's per-frame intrinsics vary by a pixel or two frame to frame (the same caveat noted for the instant-ngp convertor above). Pass `--shared_camera` if you need a single shared camera instead (e.g. for tools that assume one camera for the whole capture).
+
 ### Adding additional convertors:
 
 If you would like to add an additional export format you can do so by consulting Polycam's data specification below, and using `polyform/convertors/instant_ngp.py` as an example.

polyform/convert.py

@@ -9,22 +9,27 @@
 from polyform.utils.logging import logger
 from polyform.core.capture_folder import CaptureFolder
 from polyform.convertors.instant_ngp import InstantNGPConvertor
+from polyform.convertors.colmap import COLMAPConvertor
 
 
-def convert(data_folder_path: str, format: str = "ingp"):
+def convert(data_folder_path: str, format: str = "ingp", shared_camera: bool = False):
     """
     Main entry point for the command line convertor
     Args:
         data_folder_path: path to the unzipped Polycam data folder
-        format: Output format time. Supported values are [ingp]
+        format: Output format time. Supported values are [ingp, colmap]
+        shared_camera: (colmap format only) if True, write a single shared COLMAP
+            camera for all images instead of one camera per image
     """
     folder = CaptureFolder(data_folder_path)
     if format.lower() == "ingp" or format.lower() == "instant-ngp":
         convertor = InstantNGPConvertor()
+    elif format.lower() == "colmap":
+        convertor = COLMAPConvertor(shared_camera=shared_camera)
     else:
         logger.error("Format {} is not curently supported. Consider adding a convertor for it".format(format))
         exit(1)
     convertor.convert(folder)
 
 if __name__ == '__main__':
-    fire.Fire(convert)
+    fire.Fire(convert)

polyform/convertors/colmap.py

@@ -0,0 +1,239 @@
+'''
+File: colmap.py
+Polycam Inc.
+
+Converts Polycam raw capture data into the COLMAP text model format
+(cameras.txt / images.txt / points3D.txt), so that Polycam captures can be
+consumed by any tool in the COLMAP ecosystem (e.g. 3D Gaussian Splatting,
+nerfstudio's colmap dataparser, Open3D, MVS pipelines, ...).
+
+See https://colmap.github.io/format.html#text-format for the file spec.
+
+Addresses: https://github.com/PolyCam/polyform/issues/9
+"How to get the raw data into Colmap format to use 3DGaussian?"
+'''
+import os
+import numpy as np
+from polyform.utils.logging import logger
+from polyform.core.capture_folder import *
+from polyform.convertors.convertor_interface import ConvertorInterface
+
+# Polycam's raw camera-to-world transform (rotate=False) is expressed in
+# ARKit / OpenGL camera-local axes: +X right, +Y up, +Z out of the screen
+# (i.e. the camera looks down its local -Z axis).
+#
+# COLMAP (and OpenCV) use the computer-vision convention: +X right, +Y down,
+# +Z into the scene (i.e. the camera looks down its local +Z axis).
+#
+# The two conventions only differ by a 180 degree rotation about the local
+# X axis, so converting one to the other is a matter of flipping the local
+# Y and Z axes.
+_ARKIT_TO_CV = np.diag([1.0, -1.0, -1.0]).astype(np.float32)
+
+
+def _rotation_matrix_to_quaternion(R: np.ndarray) -> np.ndarray:
+    """
+    Converts a 3x3 rotation matrix into a (qw, qx, qy, qz) unit quaternion,
+    using COLMAP's convention (scalar-first, Hamilton quaternions).
+
+    Uses Shepperd's method, which picks whichever diagonal term is largest to
+    avoid dividing by a near-zero number (a plain "qw = sqrt(trace)/2" formula
+    is numerically unstable when qw is close to 0, i.e. for near-180-degree
+    rotations, which do show up in real scan trajectories when the camera
+    loops back on itself).
+    """
+    m = R
+    if not np.allclose(m @ m.T, np.eye(3), atol=1e-3) or not np.isclose(np.linalg.det(m), 1.0, atol=1e-3):
+        raise ValueError(
+            "Input is not a valid rotation matrix (must be orthogonal with determinant 1); "
+            "got matrix with det={:.4f}".format(np.linalg.det(m))
+        )
+
+    trace = m[0, 0] + m[1, 1] + m[2, 2]
+    if trace > 0:
+        s = 0.5 / np.sqrt(trace + 1.0)
+        qw = 0.25 / s
+        qx = (m[2, 1] - m[1, 2]) * s
+        qy = (m[0, 2] - m[2, 0]) * s
+        qz = (m[1, 0] - m[0, 1]) * s
+    elif m[0, 0] > m[1, 1] and m[0, 0] > m[2, 2]:
+        s = 2.0 * np.sqrt(1.0 + m[0, 0] - m[1, 1] - m[2, 2])
+        qw = (m[2, 1] - m[1, 2]) / s
+        qx = 0.25 * s
+        qy = (m[0, 1] + m[1, 0]) / s
+        qz = (m[0, 2] + m[2, 0]) / s
+    elif m[1, 1] > m[2, 2]:
+        s = 2.0 * np.sqrt(1.0 + m[1, 1] - m[0, 0] - m[2, 2])
+        qw = (m[0, 2] - m[2, 0]) / s
+        qx = (m[0, 1] + m[1, 0]) / s
+        qy = 0.25 * s
+        qz = (m[1, 2] + m[2, 1]) / s
+    else:
+        s = 2.0 * np.sqrt(1.0 + m[2, 2] - m[0, 0] - m[1, 1])
+        qw = (m[1, 0] - m[0, 1]) / s
+        qx = (m[0, 2] + m[2, 0]) / s
+        qy = (m[1, 2] + m[2, 1]) / s
+        qz = 0.25 * s
+
+    q = np.asarray([qw, qx, qy, qz], dtype=np.float64)
+    return q / np.linalg.norm(q)
+
+
+def arkit_pose_to_colmap(transform: np.ndarray):
+    """
+    Converts a Polycam/ARKit camera-to-world 4x4 transform (rotate=False
+    convention, see Camera in capture_folder.py) into the COLMAP world-to-camera
+    rotation (as a wxyz quaternion) and translation expected by images.txt.
+
+    Args:
+        transform: 4x4 camera-to-world matrix in ARKit/OpenGL camera-local axes
+
+    Returns:
+        (qvec, tvec): qvec is a length-4 np.ndarray [qw, qx, qy, qz], tvec is a
+        length-3 np.ndarray [tx, ty, tz]
+    """
+    R_c2w_arkit = transform[0:3, 0:3]
+    t_c2w = transform[0:3, 3]
+
+    # Re-express the camera-to-world rotation in OpenCV/COLMAP camera axes
+    R_c2w_cv = R_c2w_arkit @ _ARKIT_TO_CV
+
+    # COLMAP stores WORLD-TO-CAMERA pose, i.e. the inverse of the camera-to-world
+    # transform. For a rotation matrix the inverse is just the transpose.
+    R_w2c = R_c2w_cv.T
+    t_w2c = -R_w2c @ t_c2w
+
+    qvec = _rotation_matrix_to_quaternion(R_w2c)
+    return qvec, t_w2c.astype(np.float64)
+
+
+class COLMAPConvertor(ConvertorInterface):
+    """
+    Converts Polycam data into the COLMAP text model (cameras.txt, images.txt,
+    points3D.txt), so it can be loaded by any COLMAP-compatible tool (e.g. the
+    original 3D Gaussian Splatting reference implementation, nerfstudio's
+    `colmap` dataparser, or COLMAP itself for further processing).
+
+    NOTE on points3D.txt: Polycam's raw export does not include a sparse
+    point cloud with per-image 2D/3D correspondences (the kind COLMAP's own
+    feature-matching + triangulation pipeline produces), so we write an empty
+    points3D.txt. This is sufficient input for tools that only need camera
+    poses/intrinsics (e.g. 3D Gaussian Splatting can initialize from random
+    points when points3D.txt is empty). A natural follow-up (left as future
+    work, see PR description) would be to back-project Polycam's per-frame
+    depth maps into a seed point cloud.
+    """
+
+    def __init__(self, shared_camera: bool = False, corrected_image_padding: int = 5):
+        """
+        Args:
+            shared_camera: if True, all images reference a single COLMAP camera
+                (using the first keyframe's intrinsics). If False (default), each
+                image gets its own camera entry, since Polycam's per-frame
+                intrinsics vary by a pixel or two frame-to-frame (same tradeoff
+                noted in InstantNGPConvertor).
+            corrected_image_padding: cropping applied to the corrected/optimized
+                images, mirroring InstantNGPConvertor's handling of the black
+                border left by undistortion.
+        """
+        self.shared_camera = shared_camera
+        self.corrected_image_padding = corrected_image_padding
+
+    def convert(self, folder: CaptureFolder, output_path: str = ""):
+        """
+        Converts a Polycam CaptureFolder into a COLMAP text-format sparse model
+        by writing cameras.txt, images.txt and points3D.txt.
+
+        Args:
+            folder: the capture folder to convert
+            output_path: directory to write the COLMAP model into. Defaults to
+                a `colmap_text` directory at the root of the CaptureFolder.
+        """
+        keyframes = folder.get_keyframes(rotate=False)
+        if len(keyframes) == 0:
+            logger.error("Capture folder does not have any data! Aborting conversion to COLMAP")
+            return
+
+        if not output_path:
+            output_path = os.path.join(folder.root, "colmap_text")
+        os.makedirs(output_path, exist_ok=True)
+
+        use_corrected = folder.has_optimized_poses()
+        camera_lines = []
+        image_lines = []
+
+        shared_camera_id = 1
+        for idx, keyframe in enumerate(keyframes):
+            image_id = idx + 1
+            camera_id = shared_camera_id if self.shared_camera else image_id
+
+            cam = keyframe.camera
+            if use_corrected:
+                width = cam.width - 2 * self.corrected_image_padding
+                height = cam.height - 2 * self.corrected_image_padding
+                cx = cam.cx - self.corrected_image_padding
+                cy = cam.cy - self.corrected_image_padding
+                image_name = "{}/{}.jpg".format(CaptureArtifact.CORRECTED_IMAGES.value, keyframe.timestamp)
+            else:
+                width = cam.width
+                height = cam.height
+                cx = cam.cx
+                cy = cam.cy
+                image_name = "{}/{}.jpg".format(CaptureArtifact.IMAGES.value, keyframe.timestamp)
+
+            if not self.shared_camera or idx == 0:
+                # PINHOLE model params are: fx, fy, cx, cy
+                camera_lines.append(
+                    "{} PINHOLE {} {} {} {} {} {}".format(
+                        camera_id, width, height, cam.fx, cam.fy, cx, cy
+                    )
+                )
+
+            qvec, tvec = arkit_pose_to_colmap(keyframe.camera.transform)
+            image_lines.append(
+                "{} {} {} {} {} {} {} {} {} {}".format(
+                    image_id,
+                    qvec[0], qvec[1], qvec[2], qvec[3],
+                    tvec[0], tvec[1], tvec[2],
+                    camera_id,
+                    image_name,
+                )
+            )
+            # COLMAP's images.txt alternates: one line of pose data, then one
+            # line of (possibly empty) 2D keypoint / 3D point correspondences.
+            image_lines.append("")
+
+        self._write_cameras_txt(output_path, camera_lines)
+        self._write_images_txt(output_path, image_lines, num_images=len(keyframes))
+        self._write_points3D_txt(output_path)
+
+        logger.info("Successfully wrote COLMAP text model to {}".format(output_path))
+
+    @staticmethod
+    def _write_cameras_txt(output_path: str, camera_lines):
+        path = os.path.join(output_path, "cameras.txt")
+        with open(path, "w") as f:
+            f.write("# Camera list with one line of data per camera:\n")
+            f.write("#   CAMERA_ID, MODEL, WIDTH, HEIGHT, PARAMS[]\n")
+            f.write("# Number of cameras: {}\n".format(len(camera_lines)))
+            for line in camera_lines:
+                f.write(line + "\n")
+
+    @staticmethod
+    def _write_images_txt(output_path: str, image_lines, num_images: int):
+        path = os.path.join(output_path, "images.txt")
+        with open(path, "w") as f:
+            f.write("# Image list with two lines of data per image:\n")
+            f.write("#   IMAGE_ID, QW, QX, QY, QZ, TX, TY, TZ, CAMERA_ID, NAME\n")
+            f.write("#   POINTS2D[] as (X, Y, POINT3D_ID)\n")
+            f.write("# Number of images: {}, mean observations per image: 0\n".format(num_images))
+            for line in image_lines:
+                f.write(line + "\n")
+
+    @staticmethod
+    def _write_points3D_txt(output_path: str):
+        path = os.path.join(output_path, "points3D.txt")
+        with open(path, "w") as f:
+            f.write("# 3D point list with one line of data per point:\n")
+            f.write("#   POINT3D_ID, X, Y, Z, R, G, B, ERROR, TRACK[] as (IMAGE_ID, POINT2D_IDX)\n")
+            f.write("# Number of points: 0\n")

tests/conftest.py

@@ -0,0 +1,93 @@
+import json
+import os
+import numpy as np
+import pytest
+from PIL import Image
+
+from polyform.core.capture_folder import CaptureFolder, CaptureArtifact
+
+
+def _write_camera_json(path, fx, fy, cx, cy, width, height, transform, blur_score=1.0):
+    """
+    Writes a Polycam-style camera json. `transform` is a 3x4 camera-to-world
+    matrix (row-major) using the ARKit raw (t_00..t_23) key layout.
+    """
+    j = {
+        "fx": fx, "fy": fy, "cx": cx, "cy": cy,
+        "width": width, "height": height, "blur_score": blur_score,
+    }
+    for row in range(3):
+        for col in range(4):
+            j["t_{}{}".format(row, col)] = float(transform[row, col])
+    with open(path, "w") as f:
+        json.dump(j, f)
+
+
+def _make_capture_folder(tmp_path, num_frames=3, optimized=False, image_size=(64, 48)):
+    """
+    Builds a minimal, valid Polycam raw-data folder on disk with `num_frames`
+    synthetic keyframes, following the on-disk layout CaptureFolder expects
+    (see CaptureArtifact in polyform/core/capture_folder.py).
+    """
+    root = tmp_path / "capture"
+    for artifact in [CaptureArtifact.IMAGES, CaptureArtifact.CAMERAS, CaptureArtifact.DEPTH_MAPS]:
+        os.makedirs(root / artifact.value, exist_ok=True)
+    if optimized:
+        for artifact in [CaptureArtifact.CORRECTED_IMAGES, CaptureArtifact.CORRECTED_CAMERAS]:
+            os.makedirs(root / artifact.value, exist_ok=True)
+
+    width, height = image_size
+    rng = np.random.default_rng(seed=0)
+
+    for i in range(num_frames):
+        timestamp = 1000 + i
+
+        # Build a simple camera-to-world transform: identity rotation with a
+        # translation that moves along +X per frame, plus one frame with a
+        # non-trivial rotation so the quaternion math is actually exercised.
+        R = np.eye(3, dtype=np.float32)
+        if i == 1:
+            # 90 degree rotation about the camera-local Y axis
+            R = np.asarray([[0, 0, 1], [0, 1, 0], [-1, 0, 0]], dtype=np.float32)
+        t = np.asarray([float(i), 0.0, 0.0], dtype=np.float32)
+        transform = np.concatenate([R, t.reshape(3, 1)], axis=1)
+
+        cam_path = root / CaptureArtifact.CAMERAS.value / "{}.json".format(timestamp)
+        _write_camera_json(
+            cam_path, fx=500 + i, fy=500 + i, cx=width / 2, cy=height / 2,
+            width=width, height=height, transform=transform,
+        )
+
+        img_path = root / CaptureArtifact.IMAGES.value / "{}.jpg".format(timestamp)
+        Image.fromarray((rng.random((height, width, 3)) * 255).astype(np.uint8)).save(img_path)
+
+        depth_path = root / CaptureArtifact.DEPTH_MAPS.value / "{}.png".format(timestamp)
+        Image.fromarray((rng.random((height, width)) * 1000).astype(np.uint16)).save(depth_path)
+
+        if optimized:
+            corrected_cam_path = root / CaptureArtifact.CORRECTED_CAMERAS.value / "{}.json".format(timestamp)
+            _write_camera_json(
+                corrected_cam_path, fx=500 + i, fy=500 + i, cx=width / 2, cy=height / 2,
+                width=width, height=height, transform=transform,
+            )
+            corrected_img_path = root / CaptureArtifact.CORRECTED_IMAGES.value / "{}.jpg".format(timestamp)
+            Image.fromarray((rng.random((height, width, 3)) * 255).astype(np.uint8)).save(corrected_img_path)
+
+    return CaptureFolder(str(root))
+
+
+@pytest.fixture
+def raw_capture_folder(tmp_path):
+    return _make_capture_folder(tmp_path, num_frames=3, optimized=False)
+
+
+@pytest.fixture
+def optimized_capture_folder(tmp_path):
+    return _make_capture_folder(tmp_path, num_frames=3, optimized=True, image_size=(64, 48))
+
+
+@pytest.fixture
+def empty_capture_folder(tmp_path):
+    root = tmp_path / "empty_capture"
+    os.makedirs(root, exist_ok=True)
+    return CaptureFolder(str(root))