Function read_colmap_model

flowmap/export/colmap.py:114–171

(
    path: Path,
    device: torch.device = torch.device("cpu"),
    reorder: bool = True,
)

def read_colmap_model(
    path: Path,
    device: torch.device = torch.device("cpu"),
    reorder: bool = True,
) -> tuple[
    Float[Tensor, "frame 4 4"],  # extrinsics
    Float[Tensor, "frame 3 3"],  # intrinsics
    list[str],  # image names
]:
    model = read_model(path)
    if model is None:
        raise FileNotFoundError()
    cameras, images, _ = model

    all_extrinsics = []
    all_intrinsics = []
    all_image_names = []

    for image in images.values():
        camera: Camera = cameras[image.camera_id]

        # Read the camera intrinsics.
        intrinsics = torch.eye(3, dtype=torch.float32, device=device)
        if camera.model == "SIMPLE_PINHOLE":
            fx, cx, cy = camera.params
            fy = fx
        elif camera.model == "PINHOLE":
            fx, fy, cx, cy = camera.params
        intrinsics[0, 0] = fx
        intrinsics[1, 1] = fy
        intrinsics[0, 2] = cx
        intrinsics[1, 2] = cy
        intrinsics[0] /= camera.width
        intrinsics[1] /= camera.height
        all_intrinsics.append(intrinsics)

        # Read the camera extrinsics.
        qw, qx, qy, qz = image.qvec
        w2c = torch.eye(4, dtype=torch.float32, device=device)
        rotation = R.from_quat([qx, qy, qz, qw]).as_matrix()
        w2c[:3, :3] = torch.tensor(rotation, dtype=torch.float32, device=device)
        w2c[:3, 3] = torch.tensor(image.tvec, dtype=torch.float32, device=device)
        extrinsics = w2c.inverse()
        all_extrinsics.append(extrinsics)

        # Read the image name.
        all_image_names.append(image.name)

    # Since COLMAP shuffles the images, we generally want to re-order them according
    # to their file names so that they form a video again.
    if reorder:
        ordered = sorted([(name, index) for index, name in enumerate(all_image_names)])
        indices = torch.tensor([index for _, index in ordered])
        all_extrinsics = [all_extrinsics[index] for index in indices]
        all_intrinsics = [all_intrinsics[index] for index in indices]
        all_image_names = [all_image_names[index] for index in indices]

    return torch.stack(all_extrinsics), torch.stack(all_intrinsics), all_image_names