beverley-gorry/colmap-vslamlab
0
1# Copyright (c), ETH Zurich and UNC Chapel Hill.2# All rights reserved.3#4# Redistribution and use in source and binary forms, with or without5# modification, are permitted provided that the following conditions are met:6#7# * Redistributions of source code must retain the above copyright8# notice, this list of conditions and the following disclaimer.9#10# * Redistributions in binary form must reproduce the above copyright11# notice, this list of conditions and the following disclaimer in the12# documentation and/or other materials provided with the distribution.13#14# * Neither the name of ETH Zurich and UNC Chapel Hill nor the names of15# its contributors may be used to endorse or promote products derived16# from this software without specific prior written permission.17#18# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"19# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE20# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE21# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE22# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR23# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF24# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS25# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN26# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)27# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE28# POSSIBILITY OF SUCH DAMAGE.29 30 31# This script is based on an original implementation by True Price.32 33import sqlite334import sys35 36import numpy as np37 38IS_PYTHON3 = sys.version_info[0] >= 339 40MAX_IMAGE_ID = 2**31 - 141 42CREATE_CAMERAS_TABLE = """CREATE TABLE IF NOT EXISTS cameras (43 camera_id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,44 model INTEGER NOT NULL,45 width INTEGER NOT NULL,46 height INTEGER NOT NULL,47 params BLOB,48 prior_focal_length INTEGER NOT NULL)"""49 50CREATE_DESCRIPTORS_TABLE = """CREATE TABLE IF NOT EXISTS descriptors (51 image_id INTEGER PRIMARY KEY NOT NULL,52 rows INTEGER NOT NULL,53 cols INTEGER NOT NULL,54 data BLOB,55 FOREIGN KEY(image_id) REFERENCES images(image_id) ON DELETE CASCADE)"""56 57CREATE_IMAGES_TABLE = """CREATE TABLE IF NOT EXISTS images (58 image_id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,59 name TEXT NOT NULL UNIQUE,60 camera_id INTEGER NOT NULL,61 CONSTRAINT image_id_check CHECK(image_id >= 0 and image_id < {}),62 FOREIGN KEY(camera_id) REFERENCES cameras(camera_id))63""".format(MAX_IMAGE_ID)64 65CREATE_POSE_PRIORS_TABLE = """CREATE TABLE IF NOT EXISTS pose_priors (66 image_id INTEGER PRIMARY KEY NOT NULL,67 position BLOB,68 coordinate_system INTEGER NOT NULL,69 position_covariance BLOB,70 FOREIGN KEY(image_id) REFERENCES images(image_id) ON DELETE CASCADE)"""71 72CREATE_TWO_VIEW_GEOMETRIES_TABLE = """73CREATE TABLE IF NOT EXISTS two_view_geometries (74 pair_id INTEGER PRIMARY KEY NOT NULL,75 rows INTEGER NOT NULL,76 cols INTEGER NOT NULL,77 data BLOB,78 config INTEGER NOT NULL,79 F BLOB,80 E BLOB,81 H BLOB,82 qvec BLOB,83 tvec BLOB)84"""85 86CREATE_KEYPOINTS_TABLE = """CREATE TABLE IF NOT EXISTS keypoints (87 image_id INTEGER PRIMARY KEY NOT NULL,88 rows INTEGER NOT NULL,89 cols INTEGER NOT NULL,90 data BLOB,91 FOREIGN KEY(image_id) REFERENCES images(image_id) ON DELETE CASCADE)92"""93 94CREATE_MATCHES_TABLE = """CREATE TABLE IF NOT EXISTS matches (95 pair_id INTEGER PRIMARY KEY NOT NULL,96 rows INTEGER NOT NULL,97 cols INTEGER NOT NULL,98 data BLOB)"""99 100CREATE_NAME_INDEX = (101 "CREATE UNIQUE INDEX IF NOT EXISTS index_name ON images(name)"102)103 104CREATE_ALL = "; ".join(105 [106 CREATE_CAMERAS_TABLE,107 CREATE_IMAGES_TABLE,108 CREATE_POSE_PRIORS_TABLE,109 CREATE_KEYPOINTS_TABLE,110 CREATE_DESCRIPTORS_TABLE,111 CREATE_MATCHES_TABLE,112 CREATE_TWO_VIEW_GEOMETRIES_TABLE,113 CREATE_NAME_INDEX,114 ]115)116 117 118def image_ids_to_pair_id(image_id1, image_id2):119 if image_id1 > image_id2:120 image_id1, image_id2 = image_id2, image_id1121 return image_id1 * MAX_IMAGE_ID + image_id2122 123 124def pair_id_to_image_ids(pair_id):125 image_id2 = pair_id % MAX_IMAGE_ID126 image_id1 = (pair_id - image_id2) / MAX_IMAGE_ID127 return image_id1, image_id2128 129 130def array_to_blob(array):131 if IS_PYTHON3:132 return array.tostring()133 else:134 return np.getbuffer(array)135 136 137def blob_to_array(blob, dtype, shape=(-1,)):138 if IS_PYTHON3:139 return np.fromstring(blob, dtype=dtype).reshape(*shape)140 else:141 return np.frombuffer(blob, dtype=dtype).reshape(*shape)142 143 144class COLMAPDatabase(sqlite3.Connection):145 @staticmethod146 def connect(database_path):147 return sqlite3.connect(database_path, factory=COLMAPDatabase)148 149 def __init__(self, *args, **kwargs):150 super(COLMAPDatabase, self).__init__(*args, **kwargs)151 152 self.create_tables = lambda: self.executescript(CREATE_ALL)153 self.create_cameras_table = lambda: self.executescript(154 CREATE_CAMERAS_TABLE155 )156 self.create_descriptors_table = lambda: self.executescript(157 CREATE_DESCRIPTORS_TABLE158 )159 self.create_images_table = lambda: self.executescript(160 CREATE_IMAGES_TABLE161 )162 self.create_pose_priors_table = lambda: self.executescript(163 CREATE_POSE_PRIORS_TABLE164 )165 self.create_two_view_geometries_table = lambda: self.executescript(166 CREATE_TWO_VIEW_GEOMETRIES_TABLE167 )168 self.create_keypoints_table = lambda: self.executescript(169 CREATE_KEYPOINTS_TABLE170 )171 self.create_matches_table = lambda: self.executescript(172 CREATE_MATCHES_TABLE173 )174 self.create_name_index = lambda: self.executescript(CREATE_NAME_INDEX)175 176 def add_camera(177 self,178 model,179 width,180 height,181 params,182 prior_focal_length=False,183 camera_id=None,184 ):185 params = np.asarray(params, np.float64)186 cursor = self.execute(187 "INSERT INTO cameras VALUES (?, ?, ?, ?, ?, ?)",188 (189 camera_id,190 model,191 width,192 height,193 array_to_blob(params),194 prior_focal_length,195 ),196 )197 return cursor.lastrowid198 199 def add_image(200 self,201 name,202 camera_id,203 image_id=None,204 ):205 cursor = self.execute(206 "INSERT INTO images VALUES (?, ?, ?)", (image_id, name, camera_id)207 )208 return cursor.lastrowid209 210 def add_pose_prior(211 self, image_id, position, coordinate_system=-1, position_covariance=None212 ):213 position = np.asarray(position, dtype=np.float64)214 if position_covariance is None:215 position_covariance = np.full((3, 3), np.nan, dtype=np.float64)216 self.execute(217 "INSERT INTO pose_priors VALUES (?, ?, ?, ?)",218 (219 image_id,220 array_to_blob(position),221 coordinate_system,222 array_to_blob(position_covariance),223 ),224 )225 226 def add_keypoints(self, image_id, keypoints):227 assert len(keypoints.shape) == 2228 assert keypoints.shape[1] in [2, 4, 6]229 230 keypoints = np.asarray(keypoints, np.float32)231 self.execute(232 "INSERT INTO keypoints VALUES (?, ?, ?, ?)",233 (image_id,) + keypoints.shape + (array_to_blob(keypoints),),234 )235 236 def add_descriptors(self, image_id, descriptors):237 descriptors = np.ascontiguousarray(descriptors, np.uint8)238 self.execute(239 "INSERT INTO descriptors VALUES (?, ?, ?, ?)",240 (image_id,) + descriptors.shape + (array_to_blob(descriptors),),241 )242 243 def add_matches(self, image_id1, image_id2, matches):244 assert len(matches.shape) == 2245 assert matches.shape[1] == 2246 247 if image_id1 > image_id2:248 matches = matches[:, ::-1]249 250 pair_id = image_ids_to_pair_id(image_id1, image_id2)251 matches = np.asarray(matches, np.uint32)252 self.execute(253 "INSERT INTO matches VALUES (?, ?, ?, ?)",254 (pair_id,) + matches.shape + (array_to_blob(matches),),255 )256 257 def add_two_view_geometry(258 self,259 image_id1,260 image_id2,261 matches,262 F=np.eye(3),263 E=np.eye(3),264 H=np.eye(3),265 qvec=np.array([1.0, 0.0, 0.0, 0.0]),266 tvec=np.zeros(3),267 config=2,268 ):269 assert len(matches.shape) == 2270 assert matches.shape[1] == 2271 272 if image_id1 > image_id2:273 matches = matches[:, ::-1]274 275 pair_id = image_ids_to_pair_id(image_id1, image_id2)276 matches = np.asarray(matches, np.uint32)277 F = np.asarray(F, dtype=np.float64)278 E = np.asarray(E, dtype=np.float64)279 H = np.asarray(H, dtype=np.float64)280 qvec = np.asarray(qvec, dtype=np.float64)281 tvec = np.asarray(tvec, dtype=np.float64)282 self.execute(283 "INSERT INTO two_view_geometries VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",284 (pair_id,)285 + matches.shape286 + (287 array_to_blob(matches),288 config,289 array_to_blob(F),290 array_to_blob(E),291 array_to_blob(H),292 array_to_blob(qvec),293 array_to_blob(tvec),294 ),295 )296 297 298def example_usage():299 import argparse300 import os301 302 parser = argparse.ArgumentParser()303 parser.add_argument("--database_path", default="database.db")304 args = parser.parse_args()305 306 if os.path.exists(args.database_path):307 print("ERROR: database path already exists -- will not modify it.")308 return309 310 # Open the database.311 312 db = COLMAPDatabase.connect(args.database_path)313 314 # For convenience, try creating all the tables upfront.315 316 db.create_tables()317 318 # Create dummy cameras.319 320 model1, width1, height1, params1 = (321 0,322 1024,323 768,324 np.array((1024.0, 512.0, 384.0)),325 )326 model2, width2, height2, params2 = (327 2,328 1024,329 768,330 np.array((1024.0, 512.0, 384.0, 0.1)),331 )332 333 camera_id1 = db.add_camera(model1, width1, height1, params1)334 camera_id2 = db.add_camera(model2, width2, height2, params2)335 336 # Create dummy images.337 338 image_id1 = db.add_image("image1.png", camera_id1)339 image_id2 = db.add_image("image2.png", camera_id1)340 image_id3 = db.add_image("image3.png", camera_id2)341 image_id4 = db.add_image("image4.png", camera_id2)342 343 # Create dummy keypoints.344 #345 # Note that COLMAP supports:346 # - 2D keypoints: (x, y)347 # - 4D keypoints: (x, y, theta, scale)348 # - 6D affine keypoints: (x, y, a_11, a_12, a_21, a_22)349 350 num_keypoints = 1000351 keypoints1 = np.random.rand(num_keypoints, 2) * (width1, height1)352 keypoints2 = np.random.rand(num_keypoints, 2) * (width1, height1)353 keypoints3 = np.random.rand(num_keypoints, 2) * (width2, height2)354 keypoints4 = np.random.rand(num_keypoints, 2) * (width2, height2)355 356 db.add_keypoints(image_id1, keypoints1)357 db.add_keypoints(image_id2, keypoints2)358 db.add_keypoints(image_id3, keypoints3)359 db.add_keypoints(image_id4, keypoints4)360 361 # Create dummy matches.362 363 M = 50364 matches12 = np.random.randint(num_keypoints, size=(M, 2))365 matches23 = np.random.randint(num_keypoints, size=(M, 2))366 matches34 = np.random.randint(num_keypoints, size=(M, 2))367 368 db.add_matches(image_id1, image_id2, matches12)369 db.add_matches(image_id2, image_id3, matches23)370 db.add_matches(image_id3, image_id4, matches34)371 372 # Create dummy pose_priors.373 374 pos1 = np.random.rand(3, 1) * np.random.randint(10)375 pos2 = np.random.rand(3, 1) * np.random.randint(10)376 pos3 = np.random.rand(3, 1) * np.random.randint(10)377 378 cov3 = np.random.rand(3, 3) * np.random.randint(10)379 380 pose_prior1 = [image_id1, pos1, 1, None]381 pose_prior2 = [image_id2, pos2, -1, None]382 pose_prior3 = [image_id3, pos3, 0, cov3]383 384 db.add_pose_prior(*pose_prior1)385 db.add_pose_prior(*pose_prior2)386 db.add_pose_prior(*pose_prior3)387 388 # Convert unset covariance to nan matrix for later check389 pose_prior1[3] = np.full((3, 3), np.nan, dtype=np.float64)390 pose_prior2[3] = np.full((3, 3), np.nan, dtype=np.float64)391 392 # Commit the data to the file.393 394 db.commit()395 396 # Read and check cameras.397 398 rows = db.execute("SELECT * FROM cameras")399 400 camera_id, model, width, height, params, prior = next(rows)401 params = blob_to_array(params, np.float64)402 assert camera_id == camera_id1403 assert model == model1 and width == width1 and height == height1404 assert np.allclose(params, params1)405 406 camera_id, model, width, height, params, prior = next(rows)407 params = blob_to_array(params, np.float64)408 assert camera_id == camera_id2409 assert model == model2 and width == width2 and height == height2410 assert np.allclose(params, params2)411 412 # Read and check keypoints.413 414 keypoints = dict(415 (image_id, blob_to_array(data, np.float32, (-1, 2)))416 for image_id, data in db.execute("SELECT image_id, data FROM keypoints")417 )418 419 assert np.allclose(keypoints[image_id1], keypoints1)420 assert np.allclose(keypoints[image_id2], keypoints2)421 assert np.allclose(keypoints[image_id3], keypoints3)422 assert np.allclose(keypoints[image_id4], keypoints4)423 424 # Read and check matches.425 426 matches = dict(427 (pair_id_to_image_ids(pair_id), blob_to_array(data, np.uint32, (-1, 2)))428 for pair_id, data in db.execute("SELECT pair_id, data FROM matches")429 )430 431 assert np.all(matches[(image_id1, image_id2)] == matches12)432 assert np.all(matches[(image_id2, image_id3)] == matches23)433 assert np.all(matches[(image_id3, image_id4)] == matches34)434 435 # Read and check pose_priors436 437 rows = db.execute("SELECT * FROM pose_priors")438 439 img_id1, pos1, coord_sys1, cov1 = next(rows)440 img_id2, pos2, coord_sys2, cov2 = next(rows)441 img_id3, pos3, coord_sys3, cov3 = next(rows)442 443 assert pose_prior1[0] == img_id1444 assert pose_prior2[0] == img_id2445 assert pose_prior3[0] == img_id3446 447 assert pose_prior1[1].all() == blob_to_array(pos1, np.float64, (3, 1)).all()448 assert pose_prior2[1].all() == blob_to_array(pos2, np.float64, (3, 1)).all()449 assert pose_prior3[1].all() == blob_to_array(pos3, np.float64, (3, 1)).all()450 451 assert pose_prior1[2] == coord_sys1452 assert pose_prior2[2] == coord_sys2453 assert pose_prior3[2] == coord_sys3454 455 assert pose_prior1[3].all() == blob_to_array(cov1, np.float64, (3, 3)).all()456 assert pose_prior2[3].all() == blob_to_array(cov2, np.float64, (3, 3)).all()457 assert pose_prior3[3].all() == blob_to_array(cov3, np.float64, (3, 3)).all()458 459 # Clean up.460 461 db.close()462 463 if os.path.exists(args.database_path):464 os.remove(args.database_path)465 466 467if __name__ == "__main__":468 example_usage()469