run_colmap.py

"""
COLMAP Wrapper for Structure from Motion

This module provides a high-level interface to run COLMAP pipeline including
feature extraction, matching, sparse reconstruction, and dense reconstruction.
"""

import argparse
import os
import subprocess
from pathlib import Path
from typing import Optional

import numpy as np
from skimage.io import imread

from colmap.database import COLMAPDatabase
from colmap.read_write_model import CAMERA_MODEL_NAMES


def run_sfm(
    image_dir: str,
    project_dir: str,
    same_camera: bool = False,
    colmap_path: str = "$HOME/code/colmap/build/src/exe/colmap"
) -> None:
    """
    Run complete COLMAP Structure from Motion pipeline.

    Args:
        image_dir: Path to directory containing images
        project_dir: Output directory for COLMAP project
        same_camera: If True, use same camera for all images
        colmap_path: Path to COLMAP executable
    """
    Path(project_dir).mkdir(exist_ok=True, parents=True)
    # create database for all images
    db = COLMAPDatabase.connect(f'{project_dir}/database.db')
    db.create_tables()

    # add images
    img_dir = Path(image_dir)
    img_fns = []
    for pattern in ['*.jpg', '*.png', '*.PNG', '*.JPG']:
        img_fns += [fn for fn in img_dir.glob(pattern)]
    img_fns = sorted(img_fns)
    global_cam_id = None
    for k, img_fn in enumerate(img_fns):
        img = imread(img_fn)
        h, w, _ = img.shape
        focal = np.sqrt(h**2 + w**2)  # guess a focal here
        if same_camera:
            if k == 0:
                global_cam_id = db.add_camera(
                    CAMERA_MODEL_NAMES['SIMPLE_PINHOLE'].model_id,
                    float(w),
                    float(h),
                    np.array([focal, w / 2, h / 2], np.float64),
                    prior_focal_length=True)
            db.add_image(img_fn.name, global_cam_id)
        else:
            cam_id = db.add_camera(
                CAMERA_MODEL_NAMES['SIMPLE_PINHOLE'].model_id,
                float(w),
                float(h),
                np.array([focal, w / 2, h / 2], np.float64),
                prior_focal_length=True)
            db.add_image(img_fn.name, cam_id)

    db.commit()
    db.close()

    # feature extraction
    cmd = [
        colmap_path, 'feature_extractor', '--database_path',
        f'{project_dir}/database.db', '--image_path', f'{image_dir}'
    ]
    print(' '.join(cmd))
    subprocess.run(cmd, check=True)

    # feature matching
    cmd = [
        colmap_path, 'exhaustive_matcher', '--database_path',
        f'{project_dir}/database.db'
    ]
    print(' '.join(cmd))
    subprocess.run(cmd, check=True)

    # SfM
    Path(f'{project_dir}/sparse').mkdir(exist_ok=True, parents=True)
    cmd = [
        colmap_path, 'mapper', '--database_path', f'{project_dir}/database.db',
        '--image_path', f'{image_dir}', '--output_path',
        f'{project_dir}/sparse'
    ]
    print(' '.join(cmd))
    subprocess.run(cmd, check=True)

    # dense reconstruction
    Path(f'{project_dir}/dense').mkdir(exist_ok=True, parents=True)
    cmd = [
        colmap_path, 'image_undistorter', '--image_path', f'{image_dir}',
        '--input_path', f'{project_dir}/sparse/0', '--output_path',
        f'{project_dir}/dense'
    ]
    print(' '.join(cmd))
    subprocess.run(cmd, check=True)

    cmd = [
        colmap_path, 'patch_match_stereo', '--workspace_path',
        f'{project_dir}/dense'
    ]
    print(' '.join(cmd))
    subprocess.run(cmd, check=True)

    cmd = [
        str(colmap_path),
        'stereo_fusion',
        '--workspace_path',
        f'{project_dir}/dense',
        '--workspace_format',
        'COLMAP',
        '--input_type',
        'geometric',
        '--output_path',
        f'{project_dir}/points.ply',
        # '--StereoFusion.num_threads','14',
        '--StereoFusion.check_num_images',
        '5',
        # '--StereoFusion.cache_size','14',
        # '--StereoFusion.use_cache','1',
        # '--StereoFusion.max_image_size','1024'
    ]
    print(' '.join(cmd))
    subprocess.run(cmd, check=True)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--project_dir', type=str)
    parser.add_argument('--colmap', type=str)
    parser.add_argument('--same_camera',
                        action='store_true',
                        default=False,
                        dest='same_camera')
    args = parser.parse_args()
    image_dir = f'{args.project_dir}/images'
    run_sfm(image_dir,
            args.project_dir,
            args.same_camera,
            colmap_path=args.colmap)


if __name__ == "__main__":
    main()