import argparse import sqlite3 from collections import defaultdict from pathlib import Path import numpy as np from tqdm import tqdm from . import logger from .utils.read_write_model import ( CAMERA_MODEL_NAMES, Camera, Image, Point3D, write_model, ) def recover_database_images_and_ids(database_path): images = {} cameras = {} db = sqlite3.connect(str(database_path)) ret = db.execute("SELECT name, image_id, camera_id FROM images;") for name, image_id, camera_id in ret: images[name] = image_id cameras[name] = camera_id db.close() logger.info( f"Found {len(images)} images and {len(cameras)} cameras in database." ) return images, cameras def quaternion_to_rotation_matrix(qvec): qvec = qvec / np.linalg.norm(qvec) w, x, y, z = qvec R = np.array( [ [ 1 - 2 * y * y - 2 * z * z, 2 * x * y - 2 * z * w, 2 * x * z + 2 * y * w, ], [ 2 * x * y + 2 * z * w, 1 - 2 * x * x - 2 * z * z, 2 * y * z - 2 * x * w, ], [ 2 * x * z - 2 * y * w, 2 * y * z + 2 * x * w, 1 - 2 * x * x - 2 * y * y, ], ] ) return R def camera_center_to_translation(c, qvec): R = quaternion_to_rotation_matrix(qvec) return (-1) * np.matmul(R, c) def read_nvm_model( nvm_path, intrinsics_path, image_ids, camera_ids, skip_points=False ): with open(intrinsics_path, "r") as f: raw_intrinsics = f.readlines() logger.info(f"Reading {len(raw_intrinsics)} cameras...") cameras = {} for intrinsics in raw_intrinsics: intrinsics = intrinsics.strip("\n").split(" ") name, camera_model, width, height = intrinsics[:4] params = [float(p) for p in intrinsics[4:]] camera_model = CAMERA_MODEL_NAMES[camera_model] assert len(params) == camera_model.num_params camera_id = camera_ids[name] camera = Camera( id=camera_id, model=camera_model.model_name, width=int(width), height=int(height), params=params, ) cameras[camera_id] = camera nvm_f = open(nvm_path, "r") line = nvm_f.readline() while line == "\n" or line.startswith("NVM_V3"): line = nvm_f.readline() num_images = int(line) assert num_images == len(cameras) logger.info(f"Reading {num_images} images...") image_idx_to_db_image_id = [] image_data = [] i = 0 while i < num_images: line = nvm_f.readline() if line == "\n": continue data = line.strip("\n").split(" ") image_data.append(data) image_idx_to_db_image_id.append(image_ids[data[0]]) i += 1 line = nvm_f.readline() while line == "\n": line = nvm_f.readline() num_points = int(line) if skip_points: logger.info(f"Skipping {num_points} points.") num_points = 0 else: logger.info(f"Reading {num_points} points...") points3D = {} image_idx_to_keypoints = defaultdict(list) i = 0 pbar = tqdm(total=num_points, unit="pts") while i < num_points: line = nvm_f.readline() if line == "\n": continue data = line.strip("\n").split(" ") x, y, z, r, g, b, num_observations = data[:7] obs_image_ids, point2D_idxs = [], [] for j in range(int(num_observations)): s = 7 + 4 * j img_index, kp_index, kx, ky = data[s : s + 4] image_idx_to_keypoints[int(img_index)].append( (int(kp_index), float(kx), float(ky), i) ) db_image_id = image_idx_to_db_image_id[int(img_index)] obs_image_ids.append(db_image_id) point2D_idxs.append(kp_index) point = Point3D( id=i, xyz=np.array([x, y, z], float), rgb=np.array([r, g, b], int), error=1.0, # fake image_ids=np.array(obs_image_ids, int), point2D_idxs=np.array(point2D_idxs, int), ) points3D[i] = point i += 1 pbar.update(1) pbar.close() logger.info("Parsing image data...") images = {} for i, data in enumerate(image_data): # Skip the focal length. Skip the distortion and terminal 0. name, _, qw, qx, qy, qz, cx, cy, cz, _, _ = data qvec = np.array([qw, qx, qy, qz], float) c = np.array([cx, cy, cz], float) t = camera_center_to_translation(c, qvec) if i in image_idx_to_keypoints: # NVM only stores triangulated 2D keypoints: add dummy ones keypoints = image_idx_to_keypoints[i] point2D_idxs = np.array([d[0] for d in keypoints]) tri_xys = np.array([[x, y] for _, x, y, _ in keypoints]) tri_ids = np.array([i for _, _, _, i in keypoints]) num_2Dpoints = max(point2D_idxs) + 1 xys = np.zeros((num_2Dpoints, 2), float) point3D_ids = np.full(num_2Dpoints, -1, int) xys[point2D_idxs] = tri_xys point3D_ids[point2D_idxs] = tri_ids else: xys = np.zeros((0, 2), float) point3D_ids = np.full(0, -1, int) image_id = image_ids[name] image = Image( id=image_id, qvec=qvec, tvec=t, camera_id=camera_ids[name], name=name, xys=xys, point3D_ids=point3D_ids, ) images[image_id] = image return cameras, images, points3D def main(nvm, intrinsics, database, output, skip_points=False): assert nvm.exists(), nvm assert intrinsics.exists(), intrinsics assert database.exists(), database image_ids, camera_ids = recover_database_images_and_ids(database) logger.info("Reading the NVM model...") model = read_nvm_model( nvm, intrinsics, image_ids, camera_ids, skip_points=skip_points ) logger.info("Writing the COLMAP model...") output.mkdir(exist_ok=True, parents=True) write_model(*model, path=str(output), ext=".bin") logger.info("Done.") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--nvm", required=True, type=Path) parser.add_argument("--intrinsics", required=True, type=Path) parser.add_argument("--database", required=True, type=Path) parser.add_argument("--output", required=True, type=Path) parser.add_argument("--skip_points", action="store_true") args = parser.parse_args() main(**args.__dict__)