File size: 6,465 Bytes
ec0c8fa
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
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