colmap_io.py

"""COLMAP scene loading + view selection.

Verbatim helpers from the ReSplat reference inference script
(https://github.com/cvg/resplat, scripts/infer_colmap.py, MIT), which shares
the DepthSplat camera conventions used by CodecSplat.
"""

import collections
import os
import struct

import numpy as np
import torch

CameraModel = collections.namedtuple(
    "CameraModel", ["model_id", "model_name", "num_params"]
)
ColmapCamera = collections.namedtuple(
    "ColmapCamera", ["id", "model", "width", "height", "params"]
)
BaseImage = collections.namedtuple(
    "ColmapImage",
    ["id", "qvec", "tvec", "camera_id", "name", "xys", "point3D_ids"],
)

CAMERA_MODELS = {
    CameraModel(model_id=0, model_name="SIMPLE_PINHOLE", num_params=3),
    CameraModel(model_id=1, model_name="PINHOLE", num_params=4),
    CameraModel(model_id=2, model_name="SIMPLE_RADIAL", num_params=4),
    CameraModel(model_id=3, model_name="RADIAL", num_params=5),
    CameraModel(model_id=4, model_name="OPENCV", num_params=8),
    CameraModel(model_id=5, model_name="OPENCV_FISHEYE", num_params=8),
    CameraModel(model_id=6, model_name="FULL_OPENCV", num_params=12),
    CameraModel(model_id=7, model_name="FOV", num_params=5),
    CameraModel(model_id=8, model_name="SIMPLE_RADIAL_FISHEYE", num_params=4),
    CameraModel(model_id=9, model_name="RADIAL_FISHEYE", num_params=5),
    CameraModel(model_id=10, model_name="THIN_PRISM_FISHEYE", num_params=12),
}
CAMERA_MODEL_IDS = dict(
    [(camera_model.model_id, camera_model) for camera_model in CAMERA_MODELS]
)


class ColmapImage(BaseImage):
    def qvec2rotmat(self):
        return qvec2rotmat(self.qvec)


def qvec2rotmat(qvec):
    return np.array(
        [
            [
                1 - 2 * qvec[2] ** 2 - 2 * qvec[3] ** 2,
                2 * qvec[1] * qvec[2] - 2 * qvec[0] * qvec[3],
                2 * qvec[3] * qvec[1] + 2 * qvec[0] * qvec[2],
            ],
            [
                2 * qvec[1] * qvec[2] + 2 * qvec[0] * qvec[3],
                1 - 2 * qvec[1] ** 2 - 2 * qvec[3] ** 2,
                2 * qvec[2] * qvec[3] - 2 * qvec[0] * qvec[1],
            ],
            [
                2 * qvec[3] * qvec[1] - 2 * qvec[0] * qvec[2],
                2 * qvec[2] * qvec[3] + 2 * qvec[0] * qvec[1],
                1 - 2 * qvec[1] ** 2 - 2 * qvec[2] ** 2,
            ],
        ]
    )


def read_next_bytes(fid, num_bytes, format_char_sequence, endian_character="<"):
    data = fid.read(num_bytes)
    return struct.unpack(endian_character + format_char_sequence, data)


def read_extrinsics_binary(path_to_model_file):
    images = {}
    with open(path_to_model_file, "rb") as fid:
        num_reg_images = read_next_bytes(fid, 8, "Q")[0]
        for _ in range(num_reg_images):
            binary_image_properties = read_next_bytes(
                fid, num_bytes=64, format_char_sequence="idddddddi"
            )
            image_id = binary_image_properties[0]
            qvec = np.array(binary_image_properties[1:5])
            tvec = np.array(binary_image_properties[5:8])
            camera_id = binary_image_properties[8]
            image_name = ""
            current_char = read_next_bytes(fid, 1, "c")[0]
            while current_char != b"\x00":
                image_name += current_char.decode("utf-8")
                current_char = read_next_bytes(fid, 1, "c")[0]
            num_points2D = read_next_bytes(fid, num_bytes=8, format_char_sequence="Q")[
                0
            ]
            x_y_id_s = read_next_bytes(
                fid,
                num_bytes=24 * num_points2D,
                format_char_sequence="ddq" * num_points2D,
            )
            xys = np.column_stack(
                [
                    tuple(map(float, x_y_id_s[0::3])),
                    tuple(map(float, x_y_id_s[1::3])),
                ]
            )
            point3D_ids = np.array(tuple(map(int, x_y_id_s[2::3])))
            images[image_id] = ColmapImage(
                id=image_id,
                qvec=qvec,
                tvec=tvec,
                camera_id=camera_id,
                name=image_name,
                xys=xys,
                point3D_ids=point3D_ids,
            )
    return images


def read_extrinsics_text(path):
    images = {}
    with open(path, "r") as fid:
        while True:
            line = fid.readline()
            if not line:
                break
            line = line.strip()
            if len(line) > 0 and line[0] != "#":
                elems = line.split()
                image_id = int(elems[0])
                qvec = np.array(tuple(map(float, elems[1:5])))
                tvec = np.array(tuple(map(float, elems[5:8])))
                camera_id = int(elems[8])
                image_name = elems[9]
                elems = fid.readline().split()
                xys = np.column_stack(
                    [
                        tuple(map(float, elems[0::3])),
                        tuple(map(float, elems[1::3])),
                    ]
                )
                point3D_ids = np.array(tuple(map(int, elems[2::3])))
                images[image_id] = ColmapImage(
                    id=image_id,
                    qvec=qvec,
                    tvec=tvec,
                    camera_id=camera_id,
                    name=image_name,
                    xys=xys,
                    point3D_ids=point3D_ids,
                )
    return images


def read_intrinsics_binary(path_to_model_file):
    cameras = {}
    with open(path_to_model_file, "rb") as fid:
        num_cameras = read_next_bytes(fid, 8, "Q")[0]
        for _ in range(num_cameras):
            camera_properties = read_next_bytes(
                fid, num_bytes=24, format_char_sequence="iiQQ"
            )
            camera_id = camera_properties[0]
            model_id = camera_properties[1]
            model_name = CAMERA_MODEL_IDS[camera_properties[1]].model_name
            width = camera_properties[2]
            height = camera_properties[3]
            num_params = CAMERA_MODEL_IDS[model_id].num_params
            params = read_next_bytes(
                fid, num_bytes=8 * num_params, format_char_sequence="d" * num_params
            )
            cameras[camera_id] = ColmapCamera(
                id=camera_id,
                model=model_name,
                width=width,
                height=height,
                params=np.array(params),
            )
        assert len(cameras) == num_cameras
    return cameras


def read_intrinsics_text(path):
    cameras = {}
    with open(path, "r") as fid:
        while True:
            line = fid.readline()
            if not line:
                break
            line = line.strip()
            if len(line) > 0 and line[0] != "#":
                elems = line.split()
                camera_id = int(elems[0])
                model = elems[1]
                width = int(elems[2])
                height = int(elems[3])
                params = np.array(tuple(map(float, elems[4:])))
                cameras[camera_id] = ColmapCamera(
                    id=camera_id,
                    model=model,
                    width=width,
                    height=height,
                    params=params,
                )
    return cameras


# =============================================================================
# Scene Loading
# =============================================================================


def load_colmap_scene(scene_path, sparse_dir="sparse/0", images_dir="images"):
    """Load COLMAP reconstruction and convert to ReSplat format.

    Returns dict with:
        image_paths: list[str] - full paths to images
        image_names: list[str] - image filenames
        c2w: np.ndarray [N, 4, 4] - camera-to-world matrices (OpenCV convention)
        intrinsics: np.ndarray [N, 3, 3] - normalized intrinsic matrices
        image_sizes: list[tuple[int, int]] - (width, height) per image
    """
    sparse_path = os.path.join(scene_path, sparse_dir)

    # Try binary format first, fall back to text
    cameras_bin = os.path.join(sparse_path, "cameras.bin")
    cameras_txt = os.path.join(sparse_path, "cameras.txt")
    images_bin = os.path.join(sparse_path, "images.bin")
    images_txt = os.path.join(sparse_path, "images.txt")

    if os.path.exists(cameras_bin) and os.path.exists(images_bin):
        cam_intrinsics = read_intrinsics_binary(cameras_bin)
        cam_extrinsics = read_extrinsics_binary(images_bin)
        print(f"Loaded COLMAP binary format from {sparse_path}")
    elif os.path.exists(cameras_txt) and os.path.exists(images_txt):
        cam_intrinsics = read_intrinsics_text(cameras_txt)
        cam_extrinsics = read_extrinsics_text(images_txt)
        print(f"Loaded COLMAP text format from {sparse_path}")
    else:
        raise FileNotFoundError(
            f"No COLMAP reconstruction found in {sparse_path}. "
            "Expected cameras.bin/txt and images.bin/txt"
        )

    images_root = os.path.join(scene_path, images_dir)

    # Sort images by name for deterministic ordering
    sorted_images = sorted(cam_extrinsics.values(), key=lambda x: x.name)

    image_paths = []
    image_names = []
    c2w_list = []
    intrinsics_list = []
    image_sizes = []
    point3D_ids_list = []

    for img in sorted_images:
        # Check image exists
        img_path = os.path.join(images_root, img.name)
        if not os.path.exists(img_path):
            print(f"Warning: Image not found, skipping: {img_path}")
            continue

        # Get camera intrinsics
        cam = cam_intrinsics[img.camera_id]

        # Only support PINHOLE and SIMPLE_PINHOLE (undistorted images)
        if cam.model == "PINHOLE":
            fx, fy, cx, cy = cam.params[:4]
        elif cam.model == "SIMPLE_PINHOLE":
            f, cx, cy = cam.params[:3]
            fx = fy = f
        else:
            print(
                f"Warning: Unsupported camera model '{cam.model}' for image "
                f"{img.name}, skipping. Only PINHOLE and SIMPLE_PINHOLE are "
                f"supported. Please undistort images first using COLMAP."
            )
            continue

        # Build normalized intrinsic matrix (fx/width, fy/height, cx/width, cy/height)
        K = np.eye(3, dtype=np.float32)
        K[0, 0] = fx / cam.width
        K[1, 1] = fy / cam.height
        K[0, 2] = cx / cam.width
        K[1, 2] = cy / cam.height

        # Build W2C matrix from COLMAP qvec + tvec
        # COLMAP: qvec is world-to-camera rotation, tvec is translation in camera frame
        R = qvec2rotmat(img.qvec)
        w2c = np.eye(4, dtype=np.float32)
        w2c[:3, :3] = R
        w2c[:3, 3] = img.tvec

        # Invert to get C2W (camera-to-world)
        # COLMAP is already in OpenCV convention (Y-down, Z-forward)
        c2w = np.linalg.inv(w2c).astype(np.float32)

        # Extract valid point3D_ids (filter out -1 which means unmatched)
        valid_pts = img.point3D_ids[img.point3D_ids >= 0]

        image_paths.append(img_path)
        image_names.append(img.name)
        c2w_list.append(c2w)
        intrinsics_list.append(K)
        image_sizes.append((cam.width, cam.height))
        point3D_ids_list.append(set(valid_pts.tolist()))

    if len(image_paths) == 0:
        raise RuntimeError(f"No valid images found in {images_root}")

    return {
        "image_paths": image_paths,
        "image_names": image_names,
        "c2w": np.stack(c2w_list, axis=0),  # [N, 4, 4]
        "intrinsics": np.stack(intrinsics_list, axis=0),  # [N, 3, 3]
        "image_sizes": image_sizes,
        "point3D_ids": point3D_ids_list,  # list of sets, one per image
    }


# =============================================================================
# View Selection
# =============================================================================


def farthest_point_sample(xyz, npoint):
    """Farthest point sampling on camera positions.

    Adapted from src/dataset/view_sampler/view_sampler_bounded_v2.py

    Args:
        xyz: [B, N, 3] point cloud data
        npoint: number of samples
    Returns:
        centroids: [B, npoint] sampled indices
    """
    device = xyz.device
    B, N, C = xyz.shape

    centroids = torch.zeros(B, npoint, dtype=torch.long).to(device)
    distance = torch.ones(B, N).to(device) * 1e10

    batch_indices = torch.arange(B, dtype=torch.long).to(device)

    barycenter = torch.sum(xyz, 1)
    barycenter = barycenter / xyz.shape[1]
    barycenter = barycenter.view(B, 1, 3)

    dist = torch.sum((xyz - barycenter) ** 2, -1)
    farthest = torch.max(dist, 1)[1]

    for i in range(npoint):
        centroids[:, i] = farthest
        centroid = xyz[batch_indices, farthest, :].view(B, 1, 3)
        dist = torch.sum((xyz - centroid) ** 2, -1)
        mask = dist < distance
        distance[mask] = dist[mask]
        farthest = torch.max(distance, -1)[1]

    return centroids


def select_context_views(c2w, num_context, strategy="fps"):
    """Select context view indices from all available views.

    Args:
        c2w: [N, 4, 4] camera-to-world matrices
        num_context: number of context views to select
        strategy: "fps" | "uniform"
    Returns:
        context_indices: sorted numpy array of indices
    """
    N = len(c2w)
    if num_context >= N:
        return np.arange(N)

    if strategy == "fps":
        positions = torch.tensor(c2w[:, :3, 3], dtype=torch.float32).unsqueeze(0)
        indices = farthest_point_sample(positions, num_context)[0].numpy()
        return np.sort(indices)
    elif strategy == "uniform":
        return np.linspace(0, N - 1, num_context, dtype=int)
    else:
        raise ValueError(f"Unknown context selection strategy: {strategy}")


def select_target_views(num_images, context_indices, target_selection="remaining",
                        num_target=None):
    """Select target views for rendering.

    Args:
        num_images: total number of images
        context_indices: indices used as context
        target_selection: "remaining" | "all"
        num_target: optional max number of target views
    Returns:
        target_indices: numpy array of indices
    """
    if target_selection == "remaining":
        all_indices = np.arange(num_images)
        target_indices = np.setdiff1d(all_indices, context_indices)
    elif target_selection == "all":
        target_indices = np.arange(num_images)
    else:
        raise ValueError(f"Unknown target selection: {target_selection}")

    if num_target is not None and len(target_indices) > num_target:
        step = len(target_indices) / num_target
        selected = [int(i * step) for i in range(num_target)]
        target_indices = target_indices[selected]

    return target_indices