Spaces:
Running on Zero
Running on Zero
multimodalart HF Staff
Learn2Splat ZeroGPU demo: app, vendored optgs, prebuilt sm_120 wheels
9196862 verified | import numpy as np | |
| import os | |
| import sqlite3 | |
| #------------------------------------------------------------------------------- | |
| # convert SQLite BLOBs to/from numpy arrays | |
| def array_to_blob(arr): | |
| return np.getbuffer(arr) | |
| def blob_to_array(blob, dtype, shape=(-1,)): | |
| return np.frombuffer(blob, dtype).reshape(*shape) | |
| #------------------------------------------------------------------------------- | |
| # convert to/from image pair ids | |
| MAX_IMAGE_ID = 2**31 - 1 | |
| def get_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 get_image_ids_from_pair_id(pair_id): | |
| image_id2 = pair_id % MAX_IMAGE_ID | |
| return (pair_id - image_id2) / MAX_IMAGE_ID, image_id2 | |
| #------------------------------------------------------------------------------- | |
| # create table commands | |
| 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 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 < 2147483647), | |
| FOREIGN KEY(camera_id) REFERENCES cameras(camera_id))""" | |
| CREATE_INLIER_MATCHES_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 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 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 BLOB)""" | |
| CREATE_NAME_INDEX = \ | |
| "CREATE UNIQUE INDEX IF NOT EXISTS index_name ON images(name)" | |
| CREATE_ALL = "; ".join([CREATE_CAMERAS_TABLE, CREATE_DESCRIPTORS_TABLE, | |
| CREATE_IMAGES_TABLE, CREATE_INLIER_MATCHES_TABLE, CREATE_KEYPOINTS_TABLE, | |
| CREATE_MATCHES_TABLE, CREATE_NAME_INDEX]) | |
| #------------------------------------------------------------------------------- | |
| # functional interface for adding objects | |
| def add_camera(db, model, width, height, params, prior_focal_length=False, | |
| camera_id=None): | |
| # TODO: Parameter count checks | |
| params = np.asarray(params, np.float64) | |
| db.execute("INSERT INTO cameras VALUES (?, ?, ?, ?, ?, ?)", | |
| (camera_id, model, width, height, array_to_blob(params), | |
| prior_focal_length)) | |
| def add_descriptors(db, image_id, descriptors): | |
| descriptors = np.ascontiguousarray(descriptors, np.uint8) | |
| db.execute("INSERT INTO descriptors VALUES (?, ?, ?, ?)", | |
| (image_id,) + descriptors.shape + (array_to_blob(descriptors),)) | |
| def add_image(db, name, camera_id, prior_q=np.zeros(4), prior_t=np.zeros(3), | |
| image_id=None): | |
| db.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])) | |
| # config: defaults to fundamental matrix | |
| def add_inlier_matches(db, image_id1, image_id2, matches, config=2, F=None, | |
| E=None, H=None): | |
| assert(len(matches.shape) == 2) | |
| assert(matches.shape[1] == 2) | |
| if image_id1 > image_id2: | |
| matches = matches[:,::-1] | |
| if F is not None: | |
| F = np.asarray(F, np.float64) | |
| if E is not None: | |
| E = np.asarray(E, np.float64) | |
| if H is not None: | |
| H = np.asarray(H, np.float64) | |
| pair_id = get_pair_id(image_id1, image_id2) | |
| matches = np.asarray(matches, np.uint32) | |
| db.execute("INSERT INTO inlier_matches VALUES (?, ?, ?, ?, ?, ?, ?, ?)", | |
| (pair_id,) + matches.shape + (array_to_blob(matches), config, F, E, H)) | |
| def add_keypoints(db, image_id, keypoints): | |
| assert(len(keypoints.shape) == 2) | |
| assert(keypoints.shape[1] in [2, 4, 6]) | |
| keypoints = np.asarray(keypoints, np.float32) | |
| db.execute("INSERT INTO keypoints VALUES (?, ?, ?, ?)", | |
| (image_id,) + keypoints.shape + (array_to_blob(keypoints),)) | |
| # config: defaults to fundamental matrix | |
| def add_matches(db, 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 = get_pair_id(image_id1, image_id2) | |
| matches = np.asarray(matches, np.uint32) | |
| db.execute("INSERT INTO matches VALUES (?, ?, ?, ?)", | |
| (pair_id,) + matches.shape + (array_to_blob(matches),)) | |
| #------------------------------------------------------------------------------- | |
| # simple functional interface | |
| class COLMAPDatabase(sqlite3.Connection): | |
| def connect(database_path): | |
| return sqlite3.connect(database_path, factory=COLMAPDatabase) | |
| def __init__(self, *args, **kwargs): | |
| super(COLMAPDatabase, self).__init__(*args, **kwargs) | |
| self.initialize_tables = lambda: self.executescript(CREATE_ALL) | |
| self.initialize_cameras = \ | |
| lambda: self.executescript(CREATE_CAMERAS_TABLE) | |
| self.initialize_descriptors = \ | |
| lambda: self.executescript(CREATE_DESCRIPTORS_TABLE) | |
| self.initialize_images = \ | |
| lambda: self.executescript(CREATE_IMAGES_TABLE) | |
| self.initialize_inlier_matches = \ | |
| lambda: self.executescript(CREATE_INLIER_MATCHES_TABLE) | |
| self.initialize_keypoints = \ | |
| lambda: self.executescript(CREATE_KEYPOINTS_TABLE) | |
| self.initialize_matches = \ | |
| lambda: self.executescript(CREATE_MATCHES_TABLE) | |
| self.create_name_index = lambda: self.executescript(CREATE_NAME_INDEX) | |
| add_camera = add_camera | |
| add_descriptors = add_descriptors | |
| add_image = add_image | |
| add_inlier_matches = add_inlier_matches | |
| add_keypoints = add_keypoints | |
| add_matches = add_matches | |
| #------------------------------------------------------------------------------- | |
| def main(args): | |
| import os | |
| if os.path.exists(args.database_path): | |
| print("Error: database path already exists -- will not modify it.") | |
| exit() | |
| db = COLMAPDatabase.connect(args.database_path) | |
| # | |
| # for convenience, try creating all the tables upfront | |
| # | |
| db.initialize_tables() | |
| # | |
| # create dummy cameras | |
| # | |
| model1, w1, h1, params1 = 0, 1024, 768, np.array((1024., 512., 384.)) | |
| model2, w2, h2, params2 = 2, 1024, 768, np.array((1024., 512., 384., 0.1)) | |
| db.add_camera(model1, w1, h1, params1) | |
| db.add_camera(model2, w2, h2, params2) | |
| # | |
| # create dummy images | |
| # | |
| db.add_image("image1.png", 0) | |
| db.add_image("image2.png", 0) | |
| db.add_image("image3.png", 2) | |
| db.add_image("image4.png", 2) | |
| # | |
| # create dummy keypoints; note that COLMAP supports 2D keypoints (x, y), | |
| # 4D keypoints (x, y, theta, scale), and 6D affine keypoints | |
| # (x, y, a_11, a_12, a_21, a_22) | |
| # | |
| N = 1000 | |
| kp1 = np.random.rand(N, 2) * (1024., 768.) | |
| kp2 = np.random.rand(N, 2) * (1024., 768.) | |
| kp3 = np.random.rand(N, 2) * (1024., 768.) | |
| kp4 = np.random.rand(N, 2) * (1024., 768.) | |
| db.add_keypoints(1, kp1) | |
| db.add_keypoints(2, kp2) | |
| db.add_keypoints(3, kp3) | |
| db.add_keypoints(4, kp4) | |
| # | |
| # create dummy matches | |
| # | |
| M = 50 | |
| m12 = np.random.randint(N, size=(M, 2)) | |
| m23 = np.random.randint(N, size=(M, 2)) | |
| m34 = np.random.randint(N, size=(M, 2)) | |
| db.add_matches(1, 2, m12) | |
| db.add_matches(2, 3, m23) | |
| db.add_matches(3, 4, m34) | |
| # | |
| # check cameras | |
| # | |
| rows = db.execute("SELECT * FROM cameras") | |
| camera_id, model, width, height, params, prior = next(rows) | |
| params = blob_to_array(params, np.float32) | |
| assert model == model1 and width == w1 and height == h1 | |
| assert np.allclose(params, params1) | |
| camera_id, model, width, height, params, prior = next(rows) | |
| params = blob_to_array(params, np.float32) | |
| assert model == model2 and width == w2 and height == h2 | |
| assert np.allclose(params, params2) | |
| # | |
| # check keypoints | |
| # | |
| kps = dict( | |
| (image_id, blob_to_array(data, np.float32, (-1, 2))) | |
| for image_id, data in db.execute( | |
| "SELECT image_id, data FROM keypoints")) | |
| assert np.allclose(kps[1], kp1) | |
| assert np.allclose(kps[2], kp2) | |
| assert np.allclose(kps[3], kp3) | |
| assert np.allclose(kps[4], kp4) | |
| # | |
| # check matches | |
| # | |
| pair_ids = [get_pair_id(*pair) for pair in [(1, 2), (2, 3), (3, 4)]] | |
| matches = dict( | |
| (get_image_ids_from_pair_id(pair_id), | |
| blob_to_array(data, np.uint32, (-1, 2))) | |
| for pair_id, data in db.execute("SELECT pair_id, data FROM matches")) | |
| assert np.all(matches[(1, 2)] == m12) | |
| assert np.all(matches[(2, 3)] == m23) | |
| assert np.all(matches[(3, 4)] == m34) | |
| # | |
| # clean up | |
| # | |
| db.close() | |
| os.remove(args.database_path) | |
| #------------------------------------------------------------------------------- | |
| if __name__ == "__main__": | |
| import argparse | |
| parser = argparse.ArgumentParser( | |
| formatter_class=argparse.ArgumentDefaultsHelpFormatter) | |
| parser.add_argument("--database_path", type=str, default="database.db") | |
| args = parser.parse_args() | |
| main(args) | |