diff --git a/colmap_runner/database.py b/colmap_runner/database.py new file mode 100644 index 0000000..42a9730 --- /dev/null +++ b/colmap_runner/database.py @@ -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() diff --git a/colmap_runner/run_colmap_posed.py b/colmap_runner/run_colmap_posed.py new file mode 100644 index 0000000..5d455d7 --- /dev/null +++ b/colmap_runner/run_colmap_posed.py @@ -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)