import_colmap.py

"""Convert COLMAP text-model cameras/images into a Simam3D capture manifest."""

from __future__ import annotations

import argparse
import hashlib
import json
from pathlib import Path

import numpy as np


def _records(path: Path):
    for raw in path.read_text(encoding="utf-8").splitlines():
        line = raw.strip()
        if line and not line.startswith("#"):
            yield line.split()


def _image_records(path: Path):
    """Yield COLMAP image headers while skipping each observation line."""
    lines = path.read_text(encoding="utf-8").splitlines()
    index = 0
    while index < len(lines):
        line = lines[index].strip()
        if not line or line.startswith("#"):
            index += 1
            continue
        fields = line.split()
        if len(fields) < 10:
            raise ValueError("invalid COLMAP image record")
        yield fields
        index += 2


def _camera_matrix(fields: list[str]) -> np.ndarray:
    model = fields[1]
    width, height = int(fields[2]), int(fields[3])
    params = [float(value) for value in fields[4:]]
    if model in {"SIMPLE_PINHOLE", "SIMPLE_RADIAL", "RADIAL"}:
        focal, cx, cy = params[:3]
        fx = fy = focal
    elif model in {"PINHOLE", "OPENCV", "OPENCV_FISHEYE", "FULL_OPENCV"}:
        fx, fy, cx, cy = params[:4]
    else:
        raise ValueError(f"unsupported COLMAP camera model: {model}")
    return np.array([[fx, 0.0, cx], [0.0, fy, cy], [0.0, 0.0, 1.0]], dtype=float), width, height


def _quaternion_rotation(qw: float, qx: float, qy: float, qz: float) -> np.ndarray:
    norm = np.linalg.norm([qw, qx, qy, qz])
    if norm <= 1e-12:
        raise ValueError("COLMAP quaternion is zero")
    qw, qx, qy, qz = np.asarray([qw, qx, qy, qz], dtype=float) / norm
    return np.array([
        [1 - 2 * (qy * qy + qz * qz), 2 * (qx * qy - qz * qw), 2 * (qx * qz + qy * qw)],
        [2 * (qx * qy + qz * qw), 1 - 2 * (qx * qx + qz * qz), 2 * (qy * qz - qx * qw)],
        [2 * (qx * qz - qy * qw), 2 * (qy * qz + qx * qw), 1 - 2 * (qx * qx + qy * qy)],
    ])


def import_colmap(cameras_path: str | Path, images_path: str | Path, image_root: str | Path, held_out: set[str] | None = None) -> dict[str, object]:
    cameras = {}
    for fields in _records(Path(cameras_path)):
        matrix, width, height = _camera_matrix(fields)
        cameras[int(fields[0])] = (matrix, width, height)
    image_root = Path(image_root)
    held_out = held_out or set()
    views = []
    for fields in _image_records(Path(images_path)):
        image_id, qw, qx, qy, qz, tx, ty, tz, camera_id = fields[:9]
        name = " ".join(fields[9:])
        if int(camera_id) not in cameras:
            raise ValueError(f"image {image_id} references missing camera {camera_id}")
        matrix, width, height = cameras[int(camera_id)]
        extrinsics = np.eye(4, dtype=float)
        extrinsics[:3, :3] = _quaternion_rotation(float(qw), float(qx), float(qy), float(qz))
        extrinsics[:3, 3] = [float(tx), float(ty), float(tz)]
        relative_image = Path(name)
        absolute_image = image_root / relative_image
        view = {
            "id": str(image_id),
            "image": relative_image.as_posix(),
            "split": "held_out" if name in held_out or str(image_id) in held_out else "train",
            "width": width,
            "height": height,
            "intrinsics": matrix.tolist(),
            "extrinsics": extrinsics.tolist(),
        }
        if absolute_image.is_file():
            view["sha256"] = hashlib.sha256(absolute_image.read_bytes()).hexdigest()
        views.append(view)
    return {
        "project": "Simam3D",
        "dataset": "COLMAP text model",
        "coordinate_convention": "right-handed; extrinsics are world-to-camera",
        "views": views,
    }


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("cameras", type=Path)
    parser.add_argument("images", type=Path)
    parser.add_argument("image_root", type=Path)
    parser.add_argument("output", type=Path)
    parser.add_argument("--held-out", nargs="*", default=[])
    args = parser.parse_args(argv)
    if not args.held_out:
        parser.error("provide at least one image name or IMAGE_ID with --held-out")
    manifest = import_colmap(args.cameras, args.images, args.image_root, set(args.held_out))
    args.output.write_text(json.dumps(manifest, indent=2), encoding="utf-8")
    print(f"wrote {args.output} with {len(manifest['views'])} views")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())