otthorn/nerf_plus_plus
1
0
Fork 0
Code Issues Pull requests Projects Releases Wiki Activity

add colmap scripts for posed images

Browse source
This commit is contained in:
Kai Zhang 2021-03-20 21:38:48 -04:00 committed by GitHub
parent 03741cf85a
commit cd251d6c40
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 647 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

352
colmap_runner/database.py Normal file
View file

@ -0,0 +1,352 @@
# Copyright (c) 2018, ETH Zurich and UNC Chapel Hill.
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# * Redistributions of source code must retain the above copyright
# notice, this list of conditions and the following disclaimer.
#
# * Redistributions in binary form must reproduce the above copyright
# notice, this list of conditions and the following disclaimer in the
# documentation and/or other materials provided with the distribution.
#
# * Neither the name of ETH Zurich and UNC Chapel Hill nor the names of
# its contributors may be used to endorse or promote products derived
# from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
#
# Author: Johannes L. Schoenberger (jsch-at-demuc-dot-de)
# This script is based on an original implementation by True Price.
import sys
import sqlite3
import numpy as np
IS_PYTHON3 = sys.version_info[0] >= 3
MAX_IMAGE_ID = 2**31 - 1
CREATE_CAMERAS_TABLE = """CREATE TABLE IF NOT EXISTS cameras (
camera_id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
model INTEGER NOT NULL,
width INTEGER NOT NULL,
height INTEGER NOT NULL,
params BLOB,
prior_focal_length INTEGER NOT NULL)"""
CREATE_DESCRIPTORS_TABLE = """CREATE TABLE IF NOT EXISTS descriptors (
image_id INTEGER PRIMARY KEY NOT NULL,
rows INTEGER NOT NULL,
cols INTEGER NOT NULL,
data_500 BLOB,
FOREIGN KEY(image_id) REFERENCES images(image_id) ON DELETE CASCADE)"""
CREATE_IMAGES_TABLE = """CREATE TABLE IF NOT EXISTS images (
image_id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
name TEXT NOT NULL UNIQUE,
camera_id INTEGER NOT NULL,
prior_qw REAL,
prior_qx REAL,
prior_qy REAL,
prior_qz REAL,
prior_tx REAL,
prior_ty REAL,
prior_tz REAL,
CONSTRAINT image_id_check CHECK(image_id >= 0 and image_id < {}),
FOREIGN KEY(camera_id) REFERENCES cameras(camera_id))
""".format(MAX_IMAGE_ID)
CREATE_TWO_VIEW_GEOMETRIES_TABLE = """
CREATE TABLE IF NOT EXISTS two_view_geometries (
pair_id INTEGER PRIMARY KEY NOT NULL,
rows INTEGER NOT NULL,
cols INTEGER NOT NULL,
data_500 BLOB,
config INTEGER NOT NULL,
F BLOB,
E BLOB,
H BLOB)
"""
CREATE_KEYPOINTS_TABLE = """CREATE TABLE IF NOT EXISTS keypoints (
image_id INTEGER PRIMARY KEY NOT NULL,
rows INTEGER NOT NULL,
cols INTEGER NOT NULL,
data_500 BLOB,
FOREIGN KEY(image_id) REFERENCES images(image_id) ON DELETE CASCADE)
"""
CREATE_MATCHES_TABLE = """CREATE TABLE IF NOT EXISTS matches (
pair_id INTEGER PRIMARY KEY NOT NULL,
rows INTEGER NOT NULL,
cols INTEGER NOT NULL,
data_500 BLOB)"""
CREATE_NAME_INDEX = \
"CREATE UNIQUE INDEX IF NOT EXISTS index_name ON images(name)"
CREATE_ALL = "; ".join([
CREATE_CAMERAS_TABLE,
CREATE_IMAGES_TABLE,
CREATE_KEYPOINTS_TABLE,
CREATE_DESCRIPTORS_TABLE,
CREATE_MATCHES_TABLE,
CREATE_TWO_VIEW_GEOMETRIES_TABLE,
CREATE_NAME_INDEX
])
def image_ids_to_pair_id(image_id1, image_id2):
if image_id1 > image_id2:
image_id1, image_id2 = image_id2, image_id1
return image_id1 * MAX_IMAGE_ID + image_id2
def pair_id_to_image_ids(pair_id):
image_id2 = pair_id % MAX_IMAGE_ID
image_id1 = (pair_id - image_id2) / MAX_IMAGE_ID
return image_id1, image_id2
def array_to_blob(array):
if IS_PYTHON3:
return array.tostring()
else:
return np.getbuffer(array)
def blob_to_array(blob, dtype, shape=(-1,)):
if IS_PYTHON3:
return np.fromstring(blob, dtype=dtype).reshape(*shape)
else:
return np.frombuffer(blob, dtype=dtype).reshape(*shape)
class COLMAPDatabase(sqlite3.Connection):
@staticmethod
def connect(database_path):
return sqlite3.connect(database_path, factory=COLMAPDatabase)
def __init__(self, *args, **kwargs):
super(COLMAPDatabase, self).__init__(*args, **kwargs)
self.create_tables = lambda: self.executescript(CREATE_ALL)
self.create_cameras_table = \
lambda: self.executescript(CREATE_CAMERAS_TABLE)
self.create_descriptors_table = \
lambda: self.executescript(CREATE_DESCRIPTORS_TABLE)
self.create_images_table = \
lambda: self.executescript(CREATE_IMAGES_TABLE)
self.create_two_view_geometries_table = \
lambda: self.executescript(CREATE_TWO_VIEW_GEOMETRIES_TABLE)
self.create_keypoints_table = \
lambda: self.executescript(CREATE_KEYPOINTS_TABLE)
self.create_matches_table = \
lambda: self.executescript(CREATE_MATCHES_TABLE)
self.create_name_index = lambda: self.executescript(CREATE_NAME_INDEX)
def add_camera(self, model, width, height, params,
prior_focal_length=False, camera_id=None):
params = np.asarray(params, np.float64)
cursor = self.execute(
"INSERT INTO cameras VALUES (?, ?, ?, ?, ?, ?)",
(camera_id, model, width, height, array_to_blob(params),
prior_focal_length))
return cursor.lastrowid
def add_image(self, name, camera_id,
prior_q=np.zeros(4), prior_t=np.zeros(3), image_id=None):
cursor = self.execute(
"INSERT INTO images VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
(image_id, name, camera_id, prior_q[0], prior_q[1], prior_q[2],
prior_q[3], prior_t[0], prior_t[1], prior_t[2]))
return cursor.lastrowid
def add_keypoints(self, image_id, keypoints):
assert(len(keypoints.shape) == 2)
assert(keypoints.shape[1] in [2, 4, 6])
keypoints = np.asarray(keypoints, np.float32)
self.execute(
"INSERT INTO keypoints VALUES (?, ?, ?, ?)",
(image_id,) + keypoints.shape + (array_to_blob(keypoints),))
def add_descriptors(self, image_id, descriptors):
descriptors = np.ascontiguousarray(descriptors, np.uint8)
self.execute(
"INSERT INTO descriptors VALUES (?, ?, ?, ?)",
(image_id,) + descriptors.shape + (array_to_blob(descriptors),))
def add_matches(self, image_id1, image_id2, matches):
assert(len(matches.shape) == 2)
assert(matches.shape[1] == 2)
if image_id1 > image_id2:
matches = matches[:,::-1]
pair_id = image_ids_to_pair_id(image_id1, image_id2)
matches = np.asarray(matches, np.uint32)
self.execute(
"INSERT INTO matches VALUES (?, ?, ?, ?)",
(pair_id,) + matches.shape + (array_to_blob(matches),))
def add_two_view_geometry(self, image_id1, image_id2, matches,
F=np.eye(3), E=np.eye(3), H=np.eye(3), config=2):
assert(len(matches.shape) == 2)
assert(matches.shape[1] == 2)
if image_id1 > image_id2:
matches = matches[:,::-1]
pair_id = image_ids_to_pair_id(image_id1, image_id2)
matches = np.asarray(matches, np.uint32)
F = np.asarray(F, dtype=np.float64)
E = np.asarray(E, dtype=np.float64)
H = np.asarray(H, dtype=np.float64)
self.execute(
"INSERT INTO two_view_geometries VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
(pair_id,) + matches.shape + (array_to_blob(matches), config,
array_to_blob(F), array_to_blob(E), array_to_blob(H)))
def example_usage():
import os
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--database_path", default="database.db")
args = parser.parse_args()
if os.path.exists(args.database_path):
logging.info("ERROR: database path already exists -- will not modify it.")
return
# Open the database.
db = COLMAPDatabase.connect(args.database_path)
# For convenience, try creating all the tables upfront.
db.create_tables()
# Create dummy cameras.
model1, width1, height1, params1 = \
0, 1024, 768, np.array((1024., 512., 384.))
model2, width2, height2, params2 = \
2, 1024, 768, np.array((1024., 512., 384., 0.1))
camera_id1 = db.add_camera(model1, width1, height1, params1)
camera_id2 = db.add_camera(model2, width2, height2, params2)
# Create dummy images.
image_id1 = db.add_image("image1.png", camera_id1)
image_id2 = db.add_image("image2.png", camera_id1)
image_id3 = db.add_image("image3.png", camera_id2)
image_id4 = db.add_image("image4.png", camera_id2)
# Create dummy keypoints.
#
# Note that COLMAP supports:
# - 2D keypoints: (x, y)
# - 4D keypoints: (x, y, theta, scale)
# - 6D affine keypoints: (x, y, a_11, a_12, a_21, a_22)
num_keypoints = 1000
keypoints1 = np.random.rand(num_keypoints, 2) * (width1, height1)
keypoints2 = np.random.rand(num_keypoints, 2) * (width1, height1)
keypoints3 = np.random.rand(num_keypoints, 2) * (width2, height2)
keypoints4 = np.random.rand(num_keypoints, 2) * (width2, height2)
db.add_keypoints(image_id1, keypoints1)
db.add_keypoints(image_id2, keypoints2)
db.add_keypoints(image_id3, keypoints3)
db.add_keypoints(image_id4, keypoints4)
# Create dummy matches.
M = 50
matches12 = np.random.randint(num_keypoints, size=(M, 2))
matches23 = np.random.randint(num_keypoints, size=(M, 2))
matches34 = np.random.randint(num_keypoints, size=(M, 2))
db.add_matches(image_id1, image_id2, matches12)
db.add_matches(image_id2, image_id3, matches23)
db.add_matches(image_id3, image_id4, matches34)
# Commit the data_500 to the file.
db.commit()
# Read and check cameras.
rows = db.execute("SELECT * FROM cameras")
camera_id, model, width, height, params, prior = next(rows)
params = blob_to_array(params, np.float64)
assert camera_id == camera_id1
assert model == model1 and width == width1 and height == height1
assert np.allclose(params, params1)
camera_id, model, width, height, params, prior = next(rows)
params = blob_to_array(params, np.float64)
assert camera_id == camera_id2
assert model == model2 and width == width2 and height == height2
assert np.allclose(params, params2)
# Read and check keypoints.
keypoints = dict(
(image_id, blob_to_array(data, np.float32, (-1, 2)))
for image_id, data in db.execute(
"SELECT image_id, data_500 FROM keypoints"))
assert np.allclose(keypoints[image_id1], keypoints1)
assert np.allclose(keypoints[image_id2], keypoints2)
assert np.allclose(keypoints[image_id3], keypoints3)
assert np.allclose(keypoints[image_id4], keypoints4)
# Read and check matches.
pair_ids = [image_ids_to_pair_id(*pair) for pair in
((image_id1, image_id2),
(image_id2, image_id3),
(image_id3, image_id4))]
matches = dict(
(pair_id_to_image_ids(pair_id),
blob_to_array(data, np.uint32, (-1, 2)))
for pair_id, data in db.execute("SELECT pair_id, data_500 FROM matches")
)
assert np.all(matches[(image_id1, image_id2)] == matches12)
assert np.all(matches[(image_id2, image_id3)] == matches23)
assert np.all(matches[(image_id3, image_id4)] == matches34)
# Clean up.
db.close()
if os.path.exists(args.database_path):
os.remove(args.database_path)
if __name__ == "__main__":
example_usage()

295
colmap_runner/run_colmap_posed.py Normal file
View file

@ -0,0 +1,295 @@
import os
import json
from database import COLMAPDatabase
from pyquaternion import Quaternion
import numpy as np
import imageio
import subprocess
def bash_run(cmd):
# local install of colmap
env = os.environ.copy()
env['LD_LIBRARY_PATH'] = '/home/zhangka2/code/colmap/build/__install__/lib'
colmap_bin = '/home/zhangka2/code/colmap/build/__install__/bin/colmap'
cmd = colmap_bin + ' ' + cmd
print('\nRunning cmd: ', cmd)
subprocess.check_call(['/bin/bash', '-c', cmd], env=env)
gpu_index = '-1'
def run_sift_matching(img_dir, db_file):
print('Running sift matching...')
# if os.path.exists(db_file): # otherwise colmap will skip sift matching
# os.remove(db_file)
# feature extraction
# if there's no attached display, cannot use feature extractor with GPU
cmd = ' feature_extractor --database_path {} \
--image_path {} \
--ImageReader.camera_model PINHOLE \
--SiftExtraction.max_image_size 5000 \
--SiftExtraction.estimate_affine_shape 0 \
--SiftExtraction.domain_size_pooling 1 \
--SiftExtraction.num_threads 32 \
--SiftExtraction.use_gpu 0 \
--SiftExtraction.gpu_index {}'.format(db_file, img_dir, gpu_index)
bash_run(cmd)
# feature matching
cmd = ' exhaustive_matcher --database_path {} \
--SiftMatching.guided_matching 1 \
--SiftMatching.use_gpu 0 \
--SiftMatching.gpu_index {}'.format(db_file, gpu_index)
bash_run(cmd)
def create_init_files(pinhole_dict_file, db_file, out_dir):
if not os.path.exists(out_dir):
os.mkdir(out_dir)
# create template
with open(pinhole_dict_file) as fp:
pinhole_dict = json.load(fp)
template = {}
cameras_line_template = '{camera_id} PINHOLE {width} {height} {fx} {fy} {cx} {cy}\n'
images_line_template = '{image_id} {qw} {qx} {qy} {qz} {tx} {ty} {tz} {camera_id} {image_name}\n\n'
for img_name in pinhole_dict:
# w, h, fx, fy, cx, cy, qvec, t
params = pinhole_dict[img_name]
w = params[0]
h = params[1]
fx = params[2]
fy = params[3]
cx = params[4]
cy = params[5]
qvec = params[6:10]
tvec = params[10:13]
cam_line = cameras_line_template.format(camera_id="{camera_id}", width=w, height=h, fx=fx, fy=fy, cx=cx, cy=cy)
img_line = images_line_template.format(image_id="{image_id}", qw=qvec[0], qx=qvec[1], qy=qvec[2], qz=qvec[3],
tx=tvec[0], ty=tvec[1], tz=tvec[2], camera_id="{camera_id}", image_name=img_name)
template[img_name] = (cam_line, img_line)
# read database
db = COLMAPDatabase.connect(db_file)
table_images = db.execute("SELECT * FROM images")
img_name2id_dict = {}
for row in table_images:
img_name2id_dict[row[1]] = row[0]
cameras_txt_lines = []
images_txt_lines = []
for img_name, img_id in img_name2id_dict.items():
camera_line = template[img_name][0].format(camera_id=img_id)
cameras_txt_lines.append(camera_line)
image_line = template[img_name][1].format(image_id=img_id, camera_id=img_id)
images_txt_lines.append(image_line)
with open(os.path.join(out_dir, 'cameras.txt'), 'w') as fp:
fp.writelines(cameras_txt_lines)
with open(os.path.join(out_dir, 'images.txt'), 'w') as fp:
fp.writelines(images_txt_lines)
fp.write('\n')
# create an empty points3D.txt
fp = open(os.path.join(out_dir, 'points3D.txt'), 'w')
fp.close()
def run_point_triangulation(img_dir, db_file, out_dir):
print('Running point triangulation...')
# triangulate points
cmd = ' point_triangulator --database_path {} \
--image_path {} \
--input_path {} \
--output_path {} \
--Mapper.tri_ignore_two_view_tracks 1'.format(db_file, img_dir, out_dir, out_dir)
bash_run(cmd)
# this step is optional
def run_global_ba(in_dir, out_dir):
print('Running global BA...')
if not os.path.exists(out_dir):
os.mkdir(out_dir)
cmd = ' bundle_adjuster --input_path {in_dir} --output_path {out_dir}'.format(in_dir=in_dir, out_dir=out_dir)
bash_run(cmd)
def prepare_mvs(img_dir, sfm_dir, mvs_dir):
if not os.path.exists(mvs_dir):
os.mkdir(mvs_dir)
images_symlink = os.path.join(mvs_dir, 'images')
if os.path.exists(images_symlink):
os.unlink(images_symlink)
os.symlink(os.path.relpath(img_dir, mvs_dir),
images_symlink)
sparse_symlink = os.path.join(mvs_dir, 'sparse')
if os.path.exists(sparse_symlink):
os.unlink(sparse_symlink)
os.symlink(os.path.relpath(sfm_dir, mvs_dir),
sparse_symlink)
# prepare stereo directory
stereo_dir = os.path.join(mvs_dir, 'stereo')
for subdir in [stereo_dir,
os.path.join(stereo_dir, 'depth_maps'),
os.path.join(stereo_dir, 'normal_maps'),
os.path.join(stereo_dir, 'consistency_graphs')]:
if not os.path.exists(subdir):
os.mkdir(subdir)
# write patch-match.cfg and fusion.cfg
image_names = sorted(os.listdir(os.path.join(mvs_dir, 'images')))
with open(os.path.join(stereo_dir, 'patch-match.cfg'), 'w') as fp:
for img_name in image_names:
fp.write(img_name + '\n__auto__, 20\n')
# use all images
# fp.write(img_name + '\n__all__\n')
# randomly choose 20 images
# from random import shuffle
# candi_src_images = [x for x in image_names if x != img_name]
# shuffle(candi_src_images)
# max_src_images = 10
# fp.write(img_name + '\n' + ', '.join(candi_src_images[:max_src_images]) + '\n')
with open(os.path.join(stereo_dir, 'fusion.cfg'), 'w') as fp:
for img_name in image_names:
fp.write(img_name + '\n')
def run_photometric_mvs(mvs_dir, window_radius):
print('Running photometric MVS...')
cmd = ' patch_match_stereo --workspace_path {} \
--PatchMatchStereo.window_radius {} \
--PatchMatchStereo.min_triangulation_angle 3.0 \
--PatchMatchStereo.filter 1 \
--PatchMatchStereo.geom_consistency 1 \
--PatchMatchStereo.gpu_index={} \
--PatchMatchStereo.num_samples 15 \
--PatchMatchStereo.num_iterations 12'.format(mvs_dir,
window_radius, gpu_index)
bash_run(cmd)
def run_fuse(mvs_dir, out_ply):
print('Running depth fusion...')
cmd = ' stereo_fusion --workspace_path {} \
--output_path {} \
--input_type geometric'.format(mvs_dir, out_ply)
bash_run(cmd)
def run_possion_mesher(in_ply, out_ply, trim):
print('Running possion mesher...')
cmd = ' poisson_mesher \
--input_path {} \
--output_path {} \
--PoissonMeshing.trim {}'.format(in_ply, out_ply, trim)
bash_run(cmd)
def main(img_dir, pinhole_dict_file, out_dir):
if not os.path.exists(out_dir):
os.mkdir(out_dir)
db_file = os.path.join(out_dir, 'database.db')
run_sift_matching(img_dir, db_file)
sfm_dir = os.path.join(out_dir, 'sfm')
create_init_files(pinhole_dict_file, db_file, sfm_dir)
run_point_triangulation(img_dir, db_file, sfm_dir)
# # optional
# run_global_ba(sfm_dir, sfm_dir)
mvs_dir = os.path.join(out_dir, 'mvs')
prepare_mvs(img_dir, sfm_dir, mvs_dir)
run_photometric_mvs(mvs_dir, window_radius=5)
out_ply = os.path.join(out_dir, 'fused.ply')
run_fuse(mvs_dir, out_ply)
out_mesh_ply = os.path.join(out_dir, 'meshed_trim_3.ply')
run_possion_mesher(out_ply, out_mesh_ply, trim=3)
def convert_cam_dict_to_pinhole_dict(cam_dict_file, pinhole_dict_file, img_dir):
print('Writing pinhole_dict to: ', pinhole_dict_file)
with open(cam_dict_file) as fp:
cam_dict = json.load(fp)
pinhole_dict = {}
for img_name in cam_dict:
data_item = cam_dict[img_name]
if 'img_size' in data_item:
w, h = data_item['img_size']
else:
im = imageio.imread(os.path.join(img_dir, img_name))
h, w = im.shape[:2]
K = np.array(data_item['K']).reshape((4, 4))
W2C = np.array(data_item['W2C']).reshape((4, 4))
# params
fx = K[0, 0]
fy = K[1, 1]
assert(np.isclose(K[0, 1], 0.))
cx = K[0, 2]
cy = K[1, 2]
print(img_name)
R = W2C[:3, :3]
print(R)
u, s_old, vh = np.linalg.svd(R, full_matrices=False)
s = np.round(s_old)
print('s: {} ---> {}'.format(s_old, s))
R = np.dot(u * s, vh)
qvec = Quaternion(matrix=R)
tvec = W2C[:3, 3]
params = [w, h, fx, fy, cx, cy,
qvec[0], qvec[1], qvec[2], qvec[3],
tvec[0], tvec[1], tvec[2]]
pinhole_dict[img_name] = params
with open(pinhole_dict_file, 'w') as fp:
json.dump(pinhole_dict, fp, indent=2, sort_keys=True)
if __name__ == '__main__':
img_dir = ''
cam_dict_file = ''
out_dir = ''
os.makedirs(out_dir, exist_ok=True)
pinhole_dict_file = os.path.join(out_dir, 'pinhole_dict.json')
convert_cam_dict_to_pinhole_dict(cam_dict_file, pinhole_dict_file, img_dir)
main(img_dir, pinhole_dict_file, out_dir)