preprecess_colmap.py

preprecess_colmap.py

import numpy as np
import os
import cv2 as cv
from glob import glob
from scipy.spatial.transform import Rotation as Rot
from shutil import copytree

scans = ['106']

# scans = []

colmap_out_dir = './exp/colmap'
os.makedirs(colmap_out_dir, exist_ok=True)

def load_K_Rt_from_P(filename, P=None):
    if P is None:
        lines = open(filename).read().splitlines()
        if len(lines) == 4:
            lines = lines[1:]
        lines = [[x[0], x[1], x[2], x[3]] for x in (x.split(" ") for x in lines)]
        P = np.asarray(lines).astype(np.float32).squeeze()

    out = cv.decomposeProjectionMatrix(P)
    K = out[0]
    R = out[1]
    t = out[2]

    K = K/K[2,2]
    intrinsics = np.eye(4)
    intrinsics[:3, :3] = K

    pose = np.eye(4, dtype=np.float32)
    pose[:3, :3] = R.transpose()
    pose[:3, 3] = (t[:3] / t[3])[:, 0]

    return intrinsics, pose

for scan in scans:
    # load cameras
    print(scan)
    os.makedirs(os.path.join(colmap_out_dir, '{}'.format(scan), 'sparse'), exist_ok=True)
    os.makedirs(os.path.join(colmap_out_dir, '{}'.format(scan), 'dense'), exist_ok=True)
    cameras = np.load('./data/{}/cameras_sphere.npz'.format(scan))
    image_lis = glob('./data/{}/image/*.png'.format(scan))
    n_images = len(image_lis)

    # camera
    world_mat = cameras['world_mat_0']
    camera, pose = load_K_Rt_from_P(None, world_mat[:3, :4])
    with open(os.path.join(colmap_out_dir, '{}'.format(scan), 'sparse', 'cameras.txt'), 'w') as f:
        f.write('1 PINHOLE 1920 1080 {} {} {} {}\n'.format(camera[0, 0], camera[1, 1], camera[0, 2], camera[1, 2]))
    with open(os.path.join(colmap_out_dir, '{}'.format(scan), 'sparse', 'points3D.txt'), 'w') as f:
        pass

    # images
    # images_lis = sorted(glob('./data/dtu_scan{}/image/*.png'.format(scan)))
    copytree('./data/{}/image'.format(scan), os.path.join(colmap_out_dir, '{}'.format(scan), 'images'), dirs_exist_ok=True)
    with open(os.path.join(colmap_out_dir, '{}'.format(scan), 'sparse', 'images.txt'), 'w') as f:
        for i in range(n_images):
            print(i)
            # image_path = images_lis[i]
            # image_name = image_path.split('/')[-1]
            image_name = '{:0>3d}.png'.format(i)
            world_mat = cameras['world_mat_{}'.format(i)]
            intrinsics, pose = load_K_Rt_from_P(None, world_mat[:3, :4])
            R = pose[:3, :3].transpose()
            T = -R @ pose[:3, 3:]
            rot = Rot.from_matrix(R)
            rot = rot.as_quat()
            f.write('{} {} {} {} {} {} {} {} {} {}\n\n'.format(i + 1, rot[3], rot[0], rot[1], rot[2], T[0, 0], T[1, 0], T[2, 0], 1, image_name))