# Copyright [2020] [Michał Tyszkiewicz, Pascal Fua, Eduard Trulls] # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import os, argparse, h5py, warnings import numpy as np from tqdm import tqdm from PIL import Image, ExifTags from database import COLMAPDatabase, image_ids_to_pair_id def get_focal(image_path, err_on_default=False): image = Image.open(image_path) max_size = max(image.size) exif = image.getexif() focal = None if exif is not None: focal_35mm = None # https://github.com/colmap/colmap/blob/d3a29e203ab69e91eda938d6e56e1c7339d62a99/src/util/bitmap.cc#L299 for tag, value in exif.items(): focal_35mm = None if ExifTags.TAGS.get(tag, None) == 'FocalLengthIn35mmFilm': focal_35mm = float(value) break if focal_35mm is not None: focal = focal_35mm / 35. * max_size if focal is None: if err_on_default: raise RuntimeError("Failed to find focal length") # failed to find it in exif, use prior FOCAL_PRIOR = 1.2 focal = FOCAL_PRIOR * max_size return focal def create_camera(db, image_path, camera_model): image = Image.open(image_path) width, height = image.size focal = get_focal(image_path) if camera_model == 'simple-pinhole': model = 0 # simple pinhole param_arr = np.array([focal, width / 2, height / 2]) if camera_model == 'pinhole': model = 1 # pinhole param_arr = np.array([focal, focal, width / 2, height / 2]) elif camera_model == 'simple-radial': model = 2 # simple radial param_arr = np.array([focal, width / 2, height / 2, 0.1]) elif camera_model == 'opencv': model = 4 # opencv param_arr = np.array([focal, focal, width / 2, height / 2, 0., 0., 0., 0.]) return db.add_camera(model, width, height, param_arr) def add_keypoints(db, h5_path, image_path, img_ext, camera_model, single_camera=True): import h5py import numpy as np import os import glob from tqdm import tqdm from PIL import Image keypoint_f = h5py.File(os.path.join(h5_path, 'keypoints.h5'), 'r') camera_id = None fname_to_id = {} # Map filenames to their full paths all_images = {} for ext in ['.jpg', '.jpeg', '.png', '.JPG', '.JPEG', '.PNG']: for img_file in glob.glob(os.path.join(image_path, f'*{ext}')): base_name = os.path.splitext(os.path.basename(img_file))[0] all_images[base_name] = img_file # ★★★ Camera Model Mapping ★★★ camera_model_ids = { 'SIMPLE_PINHOLE': 0, 'PINHOLE': 1, 'SIMPLE_RADIAL': 2, 'RADIAL': 3, 'OPENCV': 4, } if camera_model not in camera_model_ids: raise ValueError(f"Unknown camera model: {camera_model}") model_id = camera_model_ids[camera_model] for filename in tqdm(list(keypoint_f.keys())): keypoints = keypoint_f[filename][()] if filename not in all_images: raise IOError(f'Image not found for key: {filename}') path = all_images[filename] fname_with_ext = os.path.basename(path) # Initialize camera once (if single_camera=True) if camera_id is None: img = Image.open(path) width, height = img.size # Use the max dimension as a heuristic for initial focal length focal_length = float(max(width, height)) cx = float(width / 2.0) cy = float(height / 2.0) # ★★★ Set parameters based on the chosen camera model ★★★ if camera_model == 'PINHOLE': # PINHOLE: [fx, fy, cx, cy] params = np.array([focal_length, focal_length, cx, cy], dtype=np.float64) print(f"Camera created: PINHOLE, f={focal_length}, cx={cx}, cy={cy}") elif camera_model == 'SIMPLE_PINHOLE': # SIMPLE_PINHOLE: [f, cx, cy] params = np.array([focal_length, cx, cy], dtype=np.float64) print(f"Camera created: SIMPLE_PINHOLE, f={focal_length}, cx={cx}, cy={cy}") elif camera_model == 'SIMPLE_RADIAL': # SIMPLE_RADIAL: [f, cx, cy, k] k = 0.0 # Initial radial distortion coefficient params = np.array([focal_length, cx, cy, k], dtype=np.float64) print(f"Camera created: SIMPLE_RADIAL, f={focal_length}, cx={cx}, cy={cy}, k={k}") else: raise ValueError(f"Unsupported camera model: {camera_model}") camera_id = db.add_camera( model=model_id, # Set based on argument width=width, height=height, params=params ) image_id = db.add_image(fname_with_ext, camera_id) fname_to_id[filename] = image_id db.add_keypoints(image_id, keypoints) return fname_to_id def add_matches(db, h5_path, fname_to_id): """ Directly add matches from matches.h5 (compatible with COLMAP format). """ import h5py import numpy as np import os from tqdm import tqdm match_file = h5py.File(os.path.join(h5_path, 'matches.h5'), 'r') added_pairs = set() n_keys = len(match_file.keys()) with tqdm(total=n_keys, desc="Importing matches to database") as pbar: for pair_key in match_file.keys(): # Parse the pair key to extract individual image filenames parts = pair_key.split('_') mid = len(parts) // 2 key_1 = '_'.join(parts[:mid]) key_2 = '_'.join(parts[mid:]) if key_1 not in fname_to_id or key_2 not in fname_to_id: print(f"Warning: Filename lookup failed for {key_1} or {key_2}") pbar.update(1) continue id_1 = fname_to_id[key_1] id_2 = fname_to_id[key_2] pair_id = image_ids_to_pair_id(id_1, id_2) # Skip if this pair has already been processed if pair_id in added_pairs: pbar.update(1) continue matches = match_file[pair_key][()] # Ensure the match array is not empty if len(matches) == 0: pbar.update(1) continue # COLMAP requirement: uint32 array with shape (N, 2) matches = matches.astype(np.uint32) # Add raw matches to the database db.add_matches(id_1, id_2, matches) # Add two-view geometry (required by COLMAP for reconstruction) # This step typically marks matches as "verified" db.add_two_view_geometry(id_1, id_2, matches) added_pairs.add(pair_id) pbar.update(1) match_file.close() def import_into_colmap(img_dir, feature_dir ='.featureout', database_path = 'colmap.db', img_ext='.jpg'): db = COLMAPDatabase.connect(database_path) db.create_tables() single_camera = False fname_to_id = add_keypoints(db, feature_dir, img_dir, img_ext, 'simple-radial', single_camera) add_matches( db, feature_dir, fname_to_id, ) db.commit() return if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('h5_path', help=('Path to the directory with ' 'keypoints.h5 and matches.h5')) parser.add_argument('image_path', help='Path to source images') parser.add_argument( '--image-extension', default='.jpg', type=str, help='Extension of files in image_path' ) parser.add_argument('--database-path', default='database.db', help='Location where the COLMAP .db file will be created' ) parser.add_argument( '--single-camera', action='store_true', help=('Consider all photos to be made with a single camera (COLMAP ' 'will reduce the number of degrees of freedom'), ) parser.add_argument( '--camera-model', choices=['simple-pinhole', 'pinhole', 'simple-radial', 'opencv'], default='simple-radial', help=('Camera model to use in COLMAP. ' 'See https://github.com/colmap/colmap/blob/master/src/base/camera_models.h' ' for explanations') ) args = parser.parse_args() if args.camera_model == 'opencv' and not args.single_camera: raise RuntimeError("Cannot use --camera-model=opencv camera without " "--single-camera (the COLMAP optimisation will " "likely fail to converge)") if os.path.exists(args.database_path): raise RuntimeError("database path already exists - will not modify it.") db = COLMAPDatabase.connect(args.database_path) db.create_tables() fname_to_id = add_keypoints(db, args.h5_path, args.image_path, args.image_extension, args.camera_model, args.single_camera) add_matches( db, args.h5_path, fname_to_id, ) db.commit()