to_kitti_format.py 4.19 KB
Newer Older
Clément Pinard's avatar
Clément Pinard committed
1
2
3
4
from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter
from path import Path
from imageio import imread, imwrite
import numpy as np
5
from colmap_util import read_model as rm
Clément Pinard's avatar
Clément Pinard committed
6
7
from skimage.transform import rescale
from skimage.measure import block_reduce
8
9
10
import gzip
from pebble import ProcessPool
from tqdm import tqdm
Clément Pinard's avatar
Clément Pinard committed
11

Clément Pinard's avatar
Clément Pinard committed
12
parser = ArgumentParser(description='create a visualization from ground truth created',
Clément Pinard's avatar
Clément Pinard committed
13
14
                        formatter_class=ArgumentDefaultsHelpFormatter)

15
parser.add_argument('--img_dir', metavar='DIR', type=Path)
Clément Pinard's avatar
Clément Pinard committed
16
17
18
19
20
21
parser.add_argument('--depth_dir', metavar='DIR', type=Path)
parser.add_argument('--input_model', metavar='DIR', type=Path)
parser.add_argument('--output_dir', metavar='DIR', default=None, type=Path)
parser.add_argument('--downscale', type=int, default=1)


22
23
24
25
26
27
28
29
def save_intrinsics(cameras, images, output_dir, downscale=1):
    def construct_intrinsics(cam):
        assert('PINHOLE' in cam.model)
        if 'SIMPLE' in cam.model:
            fx, cx, cy = cam.params
            fy = fx
        else:
            fx, fy, cx, cy = cam.params
Clément Pinard's avatar
Clément Pinard committed
30

31
32
33
        return np.array([[fx / downscale, 0, cx / downscale],
                         [0, fy / downscale, cy / downscale],
                         [0, 0, 1]])
Clément Pinard's avatar
Clément Pinard committed
34

35
36
37
38
39
40
41
42
    if len(cameras) == 1:
        cam = cameras[list(cameras.keys())[0]]
        intrinsics = construct_intrinsics(cam)
        np.savetxt(output_dir / 'intrinsics.txt', intrinsics)
    else:
        for _, img in images.items():
            cam = cameras[img.camera_id]
            intrinsics = construct_intrinsics(cam)
Clément Pinard's avatar
Clément Pinard committed
43
            intrinsics_name = output_dir / Path(img.name).stem + "_intrinsics.txt"
44
            np.savetxt(intrinsics_name, intrinsics)
Clément Pinard's avatar
Clément Pinard committed
45
46


47
48
49
50
51
52
53
54
55
56
def process_one_frame(cameras, img, img_root, depth_dir, output_dir, downscale):
    cam = cameras[img.camera_id]
    img_name = Path(img.name)
    depth_path = depth_dir / img_name.basename() + ".gz"
    h, w = cam.width, cam.height
    with gzip.open(depth_path, "rb") as f:
        depth = np.frombuffer(f.read(), np.float32).reshape(h, w)
    downscaled_depth = block_reduce(depth, (downscale, downscale), np.min)
    output_depth_name = output_dir / img_name.basename() + '.npy'
    np.save(output_depth_name, downscaled_depth)
Clément Pinard's avatar
Clément Pinard committed
57

58
59
60
61
    input_img_path = img_root / img_name
    output_img_path = output_dir / img_name.basename()
    image = rescale(imread(input_img_path), 1/downscale, multichannel=True)*255
    imwrite(output_img_path, image.astype(np.uint8))
Clément Pinard's avatar
Clément Pinard committed
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83


def save_positions(images, output_dir):
    starting_pos = None
    positions = []
    for _, img in images.items():
        current_pos = to_transform_matrix(img.qvec, img.tvec)
        if starting_pos is None:
            starting_pos = current_pos
        relative_position = np.linalg.inv(starting_pos) @ current_pos
        positions.append(relative_position[:3])
    positions = np.stack(positions)
    np.savetxt(output_dir/'poses.txt', positions.reshape((len(images), -1)))


def to_transform_matrix(q, t):
    cam_R = rm.qvec2rotmat(q).T
    cam_t = (- cam_R @ t).reshape(3, 1)
    transform = np.vstack((np.hstack([cam_R, cam_t]), [0, 0, 0, 1]))
    return transform


84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
def convert_to_kitti(input_model, img_dir, depth_dir, output_dir, downscale=1, threads=1, **env):
    cameras, images, _ = rm.read_model(input_model, '.txt')
    save_intrinsics(cameras, images, output_dir, downscale)
    save_positions(images, output_dir)
    if threads == 1:
        for _, img in tqdm(images.items()):
            process_one_frame(cameras, img, img_dir, depth_dir, output_dir, downscale)
    else:
        with ProcessPool(max_workers=threads) as pool:

            tasks = pool.map(process_one_frame, [cameras] * len(images),
                             [img for _, img in images.items()],
                             [img_dir] * len(images),
                             [depth_dir] * len(images),
                             [output_dir] * len(images),
                             [downscale] * len(images))
            try:
                for _ in tqdm(tasks.result(), total=len(images)):
                    pass
            except KeyboardInterrupt as e:
                tasks.cancel()
                raise e
Clément Pinard's avatar
Clément Pinard committed
106
107
108


if __name__ == '__main__':
109
110
    args = parser.parse_args()
    convert_to_kitti(vars(args))