h5_to_db.py

#  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()