140 lines
6.3 KiB
Python
140 lines
6.3 KiB
Python
from typing import *
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
from scipy.spatial.transform import Rotation
|
|
|
|
|
|
__all__ = ['read_extrinsics_from_colmap', 'read_intrinsics_from_colmap', 'write_extrinsics_as_colmap', 'write_intrinsics_as_colmap']
|
|
|
|
|
|
def write_extrinsics_as_colmap(file: Union[str, Path], extrinsics: np.ndarray, image_names: Union[str, List[str]] = 'image_{i:04d}.png', camera_ids: List[int] = None):
|
|
"""
|
|
Write extrinsics to colmap `images.txt` file.
|
|
Args:
|
|
file: Path to `images.txt` file.
|
|
extrinsics: (N, 4, 4) array of extrinsics.
|
|
image_names: str or List of str, image names. Length is N.
|
|
If str, it should be a format string with `i` as the index. (i starts from 1, in correspondence with IMAGE_ID in colmap)
|
|
camera_ids: List of int, camera ids. Length is N.
|
|
If None, it will be set to [1, 2, ..., N].
|
|
"""
|
|
assert extrinsics.shape[1:] == (4, 4) and extrinsics.ndim == 3 or extrinsics.shape == (4, 4)
|
|
if extrinsics.ndim == 2:
|
|
extrinsics = extrinsics[np.newaxis, ...]
|
|
quats = Rotation.from_matrix(extrinsics[:, :3, :3]).as_quat()
|
|
trans = extrinsics[:, :3, 3]
|
|
if camera_ids is None:
|
|
camera_ids = list(range(1, len(extrinsics) + 1))
|
|
if isinstance(image_names, str):
|
|
image_names = [image_names.format(i=i) for i in range(1, len(extrinsics) + 1)]
|
|
assert len(extrinsics) == len(image_names) == len(camera_ids), \
|
|
f'Number of extrinsics ({len(extrinsics)}), image_names ({len(image_names)}), and camera_ids ({len(camera_ids)}) must be the same'
|
|
with open(file, 'w') as fp:
|
|
print("# IMAGE_ID, QW, QX, QY, QZ, TX, TY, TZ, CAMERA_ID, NAME", file=fp)
|
|
for i, (quat, t, name, camera_id) in enumerate(zip(quats.tolist(), trans.tolist(), image_names, camera_ids)):
|
|
# Colmap has wxyz order while scipy.spatial.transform.Rotation has xyzw order. Haha, wcnm.
|
|
qx, qy, qz, qw = quat
|
|
tx, ty, tz = t
|
|
print(f'{i + 1} {qw:f} {qx:f} {qy:f} {qz:f} {tx:f} {ty:f} {tz:f} {camera_id:d} {name}', file=fp)
|
|
print()
|
|
|
|
|
|
def write_intrinsics_as_colmap(file: Union[str, Path], intrinsics: np.ndarray, width: int, height: int, normalized: bool = False):
|
|
"""
|
|
Write intrinsics to colmap `cameras.txt` file. Currently only support PINHOLE model (no distortion)
|
|
Args:
|
|
file: Path to `cameras.txt` file.
|
|
intrinsics: (N, 3, 3) array of intrinsics.
|
|
width: Image width.
|
|
height: Image height.
|
|
normalized: Whether the intrinsics are normalized. If True, the intrinsics will unnormalized for writing.
|
|
"""
|
|
assert intrinsics.shape[1:] == (3, 3) and intrinsics.ndim == 3 or intrinsics.shape == (3, 3)
|
|
if intrinsics.ndim == 2:
|
|
intrinsics = intrinsics[np.newaxis, ...]
|
|
if normalized:
|
|
intrinsics = intrinsics * np.array([width, height, 1])[:, None]
|
|
with open(file, 'w') as fp:
|
|
print("# CAMERA_ID, MODEL, WIDTH, HEIGHT, PARAMS[]", file=fp)
|
|
for i, intr in enumerate(intrinsics):
|
|
fx, fy, cx, cy = intr[0, 0], intr[1, 1], intr[0, 2], intr[1, 2]
|
|
print(f'{i + 1} PINHOLE {width:d} {height:d} {fx:f} {fy:f} {cx:f} {cy:f}', file=fp)
|
|
|
|
|
|
def read_extrinsics_from_colmap(file: Union[str, Path]) -> Union[np.ndarray, List[int], List[str]]:
|
|
"""
|
|
Read extrinsics from colmap `images.txt` file.
|
|
Args:
|
|
file: Path to `images.txt` file.
|
|
Returns:
|
|
extrinsics: (N, 4, 4) array of extrinsics.
|
|
camera_ids: List of int, camera ids. Length is N. Note that camera ids in colmap typically starts from 1.
|
|
image_names: List of str, image names. Length is N.
|
|
"""
|
|
with open(file) as fp:
|
|
lines = fp.readlines()
|
|
image_names, quats, trans, camera_ids = [], [], [], []
|
|
i_line = 0
|
|
for line in lines:
|
|
line = line.strip()
|
|
if line.startswith('#'):
|
|
continue
|
|
i_line += 1
|
|
if i_line % 2 == 0:
|
|
continue
|
|
image_id, qw, qx, qy, qz, tx, ty, tz, camera_id, name = line.split()
|
|
quats.append([float(qx), float(qy), float(qz), float(qw)])
|
|
trans.append([float(tx), float(ty), float(tz)])
|
|
camera_ids.append(int(camera_id))
|
|
image_names.append(name)
|
|
|
|
quats = np.array(quats, dtype=np.float32)
|
|
trans = np.array(trans, dtype=np.float32)
|
|
rotation = Rotation.from_quat(quats).as_matrix()
|
|
extrinsics = np.concatenate([
|
|
np.concatenate([rotation, trans[..., None]], axis=-1),
|
|
np.array([0, 0, 0, 1], dtype=np.float32)[None, None, :].repeat(len(quats), axis=0)
|
|
], axis=-2)
|
|
|
|
return extrinsics, camera_ids, image_names
|
|
|
|
|
|
def read_intrinsics_from_colmap(file: Union[str, Path], normalize: bool = False) -> Tuple[List[int], np.ndarray, np.ndarray]:
|
|
"""
|
|
Read intrinsics from colmap `cameras.txt` file.
|
|
Args:
|
|
file: Path to `cameras.txt` file.
|
|
normalize: Whether to normalize the intrinsics. If True, the intrinsics will be normalized. (mapping coordinates to [0, 1] range)
|
|
Returns:
|
|
camera_ids: List of int, camera ids. Length is N. Note that camera ids in colmap typically starts from 1.
|
|
intrinsics: (N, 3, 3) array of intrinsics.
|
|
distortions: (N, 5) array of distortions.
|
|
"""
|
|
with open(file) as fp:
|
|
lines = fp.readlines()
|
|
intrinsics, distortions, camera_ids = [], [], []
|
|
for line in lines:
|
|
line = line.strip()
|
|
if not line or line.startswith('#'):
|
|
continue
|
|
camera_id, model, width, height, *params = line.split()
|
|
camera_id, width, height = int(camera_id), int(width), int(height)
|
|
if model == 'PINHOLE':
|
|
fx, fy, cx, cy = map(float, params[:4])
|
|
k1 = k2 = k3 = p1 = p2 = 0.0
|
|
elif model == 'OPENCV':
|
|
fx, fy, cx, cy, k1, k2, p1, p2, k3 = *map(float, params[:8]), 0.0
|
|
elif model == 'SIMPLE_RADIAL':
|
|
f, cx, cy, k = map(float, params[:4])
|
|
fx = fy = f
|
|
k1, k2, p1, p2, k3 = k, 0.0, 0.0, 0.0, 0.0
|
|
camera_ids.append(camera_id)
|
|
if normalize:
|
|
fx, fy, cx, cy = fx / width, fy / height, cx / width, cy / height
|
|
intrinsics.append([[fx, 0, cx], [0, fy, cy], [0, 0, 1]])
|
|
distortions.append([k1, k2, p1, p2, k3])
|
|
intrinsics = np.array(intrinsics, dtype=np.float32)
|
|
distortions = np.array(distortions, dtype=np.float32)
|
|
return camera_ids, intrinsics, distortions
|