EasyMocap/easymocap/mytools/colmap_wrapper.py

'''
  @ Date: 2022-06-20 15:03:50
  @ Author: Qing Shuai
  @ Mail: s_q@zju.edu.cn
  @ LastEditors: Qing Shuai
  @ LastEditTime: 2022-09-02 22:38:25
  @ FilePath: /EasyMocapPublic/easymocap/mytools/colmap_wrapper.py
'''

import shutil
import sys
import os
import sqlite3
import numpy as np
from os.path import join
import cv2
from .debug_utils import mkdir, run_cmd, log, mywarn
from .colmap_structure import Camera, Image, CAMERA_MODEL_NAMES
from .colmap_structure import rotmat2qvec
from .colmap_structure import read_points3d_binary

IS_PYTHON3 = sys.version_info[0] >= 3

MAX_IMAGE_ID = 2**31 - 1

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 < {}),
    FOREIGN KEY(camera_id) REFERENCES cameras(camera_id))
""".format(MAX_IMAGE_ID)

CREATE_TWO_VIEW_GEOMETRIES_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_TWO_VIEW_GEOMETRIES_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,
    qvec BLOB,
    tvec 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_IMAGES_TABLE,
    CREATE_KEYPOINTS_TABLE,
    CREATE_DESCRIPTORS_TABLE,
    CREATE_MATCHES_TABLE,
    CREATE_TWO_VIEW_GEOMETRIES_TABLE,
    CREATE_NAME_INDEX
])

def image_ids_to_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 pair_id_to_image_ids(pair_id):
    image_id2 = pair_id % MAX_IMAGE_ID
    image_id1 = (pair_id - image_id2) // MAX_IMAGE_ID
    return image_id1, image_id2

def array_to_blob(array):
    if IS_PYTHON3:
        return array.tobytes()
    else:
        return np.getbuffer(array)

def blob_to_array(blob, dtype, shape=(-1,)):
    if blob is None:
        return np.empty((0, 2), dtype=dtype)
    if IS_PYTHON3:
        return np.frombuffer(blob, dtype=dtype).reshape(*shape)
    else:
        return np.frombuffer(blob, dtype=dtype).reshape(*shape)


class COLMAPDatabase(sqlite3.Connection):

    @staticmethod
    def connect(database_path):
        return sqlite3.connect(database_path, factory=COLMAPDatabase)


    def __init__(self, *args, **kwargs):
        super(COLMAPDatabase, self).__init__(*args, **kwargs)

        self.create_tables = lambda: self.executescript(CREATE_ALL)
        self.create_cameras_table = \
            lambda: self.executescript(CREATE_CAMERAS_TABLE)
        self.create_descriptors_table = \
            lambda: self.executescript(CREATE_DESCRIPTORS_TABLE)
        self.create_images_table = \
            lambda: self.executescript(CREATE_IMAGES_TABLE)
        self.create_two_view_geometries_table = \
            lambda: self.executescript(CREATE_TWO_VIEW_GEOMETRIES_TABLE)
        self.create_keypoints_table = \
            lambda: self.executescript(CREATE_KEYPOINTS_TABLE)
        self.create_matches_table = \
            lambda: self.executescript(CREATE_MATCHES_TABLE)
        self.create_name_index = lambda: self.executescript(CREATE_NAME_INDEX)

    def add_camera(self, model, width, height, params,
                   prior_focal_length=False, camera_id=None):
        params = np.asarray(params, np.float64)
        cursor = self.execute(
            "INSERT INTO cameras VALUES (?, ?, ?, ?, ?, ?)",
            (camera_id, model, width, height, array_to_blob(params),
             prior_focal_length))
        return cursor.lastrowid

    def add_image(self, name, camera_id,
                  prior_q=np.full(4, np.NaN), prior_t=np.full(3, np.NaN), image_id=None):
        cursor = self.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]))
        return cursor.lastrowid

    def add_keypoints(self, image_id, keypoints):
        assert(len(keypoints.shape) == 2)
        assert(keypoints.shape[1] in [2, 4, 6])

        keypoints = np.asarray(keypoints, np.float32)
        self.execute(
            "INSERT INTO keypoints VALUES (?, ?, ?, ?)",
            (image_id,) + keypoints.shape + (array_to_blob(keypoints),))

    def add_descriptors(self, image_id, descriptors):
        descriptors = np.ascontiguousarray(descriptors, np.uint8)
        self.execute(
            "INSERT INTO descriptors VALUES (?, ?, ?, ?)",
            (image_id,) + descriptors.shape + (array_to_blob(descriptors),))

    def add_matches(self, 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 = image_ids_to_pair_id(image_id1, image_id2)
        matches = np.asarray(matches, np.uint32)
        self.execute(
            "INSERT INTO matches VALUES (?, ?, ?, ?)",
            (pair_id,) + matches.shape + (array_to_blob(matches),))

    def add_two_view_geometry(self, image_id1, image_id2, matches, extra, config=2):
        assert(len(matches.shape) == 2)
        assert(matches.shape[1] == 2)

        if image_id1 > image_id2:
            matches = matches[:,::-1]

        pair_id = image_ids_to_pair_id(image_id1, image_id2)
        matches = np.asarray(matches, np.uint32)
        if 'qvec' in extra.keys():
            self.execute(
                "INSERT INTO two_view_geometries VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
                (pair_id,) + matches.shape + (array_to_blob(matches), config,
                array_to_blob(extra['F']), array_to_blob(extra['E']), array_to_blob(extra['H']),
                array_to_blob(extra['qvec']), array_to_blob(extra['tvec'])))
        else:
            self.execute(
                "INSERT INTO two_view_geometries VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
                (pair_id,) + matches.shape + (array_to_blob(matches), config,
                array_to_blob(extra['F']), array_to_blob(extra['E']), array_to_blob(extra['H'])))
    
    def read_images(self):
        image_id_to_name, name_to_image_id = {}, {}
        image_results = self.execute("SELECT * FROM images")
        for result in image_results:
            image_id, name, camera_id, q0, q1, q2, q3, t0, t1, t2 = result
            image_id_to_name[image_id] = name
            name_to_image_id[name] = image_id
        return image_id_to_name, name_to_image_id

    def read_keypoints(self, mapping=None):
        image_id_to_keypoints = {}
        keypoints_results = self.execute("SELECT * FROM keypoints")
        for keypoints_result in keypoints_results:
            image_id, rows, cols, keypoints = keypoints_result
            keypoints = blob_to_array(keypoints, np.float32, (rows, cols))
            if mapping is None:
                image_id_to_keypoints[image_id] = keypoints
            else:
                image_id_to_keypoints[mapping[image_id]] = keypoints
        return image_id_to_keypoints
    
    def read_descriptors(self, mapping=None):
        image_id_to_descriptors = {}
        descriptors_results = self.execute("SELECT * FROM descriptors")
        for descriptors_result in descriptors_results:
            image_id, rows, cols, keypoints = descriptors_result
            keypoints = blob_to_array(keypoints, np.uint8, (rows, cols))
            if mapping is None:
                image_id_to_descriptors[image_id] = keypoints
            else:
                image_id_to_descriptors[mapping[image_id]] = keypoints
        return image_id_to_descriptors

    def read_matches(self, mapping=None):
        matches_results = self.execute("SELECT * FROM matches")
        matches = {}
        for matches_result in matches_results:
            pair_id, rows, cols, match = matches_result
            image_id0, image_id1 = pair_id_to_image_ids(pair_id)
            if rows == 0:
                continue
            match = blob_to_array(match, dtype=np.uint32, shape=(rows, cols))
            if mapping is not None:
                image_id0 = mapping[image_id0]
                image_id1 = mapping[image_id1]
            matches[(image_id0, image_id1)] = match
        return matches
    
    def read_two_view_geometry(self, mapping=None):
        geometry = self.execute("SELECT * FROM two_view_geometries")
        geometries = {}
        for _data in geometry:
            if len(_data) == 10:
                pair_id, rows, cols, data, config, F, E, H, qvec, tvec = _data
                extra = {
                    'F': F,
                    'E': E,
                    'H': H,
                    'qvec': qvec,
                    'tvec': tvec
                }
            elif len(_data) == 8:
                pair_id, rows, cols, data, config, F, E, H = _data
                extra = {
                    'F': F,
                    'E': E,
                    'H': H,
                }
            for key, val in extra.items():
                extra[key] = blob_to_array(val, dtype=np.float64)
            image_id0, image_id1 = pair_id_to_image_ids(pair_id)
            match = blob_to_array(data, dtype=np.uint32, shape=(rows, cols))
            if rows == 0:continue
            if mapping is not None:
                image_id0 = mapping[image_id0]
                image_id1 = mapping[image_id1]
            geometries[(image_id0, image_id1)] = {'matches': match, 'extra': extra, 'config': config}
        return geometries

def create_empty_db(database_path):
    if os.path.exists(database_path):
        mywarn('Removing old database: {}'.format(database_path))
        os.remove(database_path)
    print('Creating an empty database...')
    db = COLMAPDatabase.connect(database_path)
    db.create_tables()
    db.commit()
    db.close()

def create_cameras(db, cameras, subs, width, height, share_intri=True):
    model = 'OPENCV'
    if share_intri:
        cam_id = 1
        K = cameras[subs[0]]['K']
        D = cameras[subs[0]]['dist'].reshape(1, 5)
        fx, fy, cx, cy, k1, k2, p1, p2, k3, k4, k5, k6 = K[0, 0], K[1, 1], K[0, 2], K[1, 2], D[0, 0], D[0, 1], D[0, 2], D[0, 3], D[0, 4], 0, 0, 0
        
        params = [fx, fy, cx, cy, k1, k2, p1, p2]
        # params = [fx, fy, cx, cy, 0, 0, 0, 0]
        camera = Camera(
            id=cam_id,
            model=model,
            width=width,
            height=height,
            params=params
        )
        cameras_colmap = {cam_id: camera}
        cameras_map = {sub:cam_id for sub in subs}
        # 
        db.add_camera(CAMERA_MODEL_NAMES[model].model_id, width, height, params,
                   prior_focal_length=False, camera_id=cam_id)
    else:
        raise NotImplementedError
    return cameras_colmap, cameras_map

def create_images(db, cameras, cameras_map, image_names):
    subs = sorted(list(image_names.keys()))
    images = {}
    for sub, image_name in image_names.items():
        img_id = subs.index(sub) + 1
        R = cameras[sub]['R']
        T = cameras[sub]['T']
        qvec = rotmat2qvec(R)
        tvec = T.T[0]
        image = Image(
            id=img_id,
            qvec=qvec,
            tvec=tvec,
            camera_id=cameras_map[sub],
            name=os.path.basename(image_name),
            xys=[],
            point3D_ids=[]
        )
        images[img_id] = image
        db.add_image(image.name, camera_id=image.camera_id,
            prior_q=image.qvec, prior_t=image.tvec, image_id=img_id)
    return images

def copy_images(data, out, nf=0, copy_func=shutil.copyfile, mask='mask', add_mask=True):
    subs = sorted(os.listdir(join(data, 'images')))
    image_names = {}
    for sub in subs:
        srcname = join(data, 'images', sub, '{:06d}.jpg'.format(nf))
        if not os.path.exists(srcname):
            mywarn('{} not exists, skip'.format(srcname))
            return False
        dstname = join(out, 'images', '{}.jpg'.format(sub))
        image_names[sub] = dstname
        if os.path.exists(dstname):
            continue
        os.makedirs(os.path.dirname(dstname), exist_ok=True)
        copy_func(srcname, dstname)
        mskname = join(data, mask, sub, '{:06d}.png'.format(nf))
        dstname = join(out, 'mask', '{}.jpg.png'.format(sub))
        if os.path.exists(mskname) and add_mask:
            os.makedirs(os.path.dirname(dstname), exist_ok=True)
            copy_func(mskname, dstname)
    return True, image_names

def colmap_feature_extract(colmap, path, share_camera, add_mask, gpu=False,
    share_camera_per_folder=False):
    '''
struct SiftMatchingOptions {
  // Number of threads for feature matching and geometric verification.
  int num_threads = -1;

  // Whether to use the GPU for feature matching.
  bool use_gpu = true;

  // Index of the GPU used for feature matching. For multi-GPU matching,
  // you should separate multiple GPU indices by comma, e.g., "0,1,2,3".
  std::string gpu_index = "-1";

  // Maximum distance ratio between first and second best match.
  double max_ratio = 0.8;

  // Maximum distance to best match.
  double max_distance = 0.7;

  // Whether to enable cross checking in matching.
  bool cross_check = true;

  // Maximum number of matches.
  int max_num_matches = 32768;

  // Maximum epipolar error in pixels for geometric verification.
  double max_error = 4.0;

  // Confidence threshold for geometric verification.
  double confidence = 0.999;

  // Minimum/maximum number of RANSAC iterations. Note that this option
  // overrules the min_inlier_ratio option.
  int min_num_trials = 100;
  int max_num_trials = 10000;

  // A priori assumed minimum inlier ratio, which determines the maximum
  // number of iterations.
  double min_inlier_ratio = 0.25;

  // Minimum number of inliers for an image pair to be considered as
  // geometrically verified.
  int min_num_inliers = 15;

  // Whether to attempt to estimate multiple geometric models per image pair.
  bool multiple_models = false;

  // Whether to perform guided matching, if geometric verification succeeds.
  bool guided_matching = false;

  bool Check() const;
};
'''
    cmd = f'{colmap} feature_extractor --database_path {path}/database.db \
--image_path {path}/images \
--SiftExtraction.peak_threshold 0.006 \
--ImageReader.camera_model OPENCV \
'
    if share_camera and not share_camera_per_folder:
        cmd += ' --ImageReader.single_camera 1'
    elif share_camera_per_folder:
        cmd += ' --ImageReader.single_camera_per_folder 1'
    if gpu:
        cmd += ' --SiftExtraction.use_gpu 1'
        cmd += ' --SiftExtraction.gpu_index 0'
    if add_mask:
        cmd += f' --ImageReader.mask_path {path}/mask'
    cmd += f' >> {path}/log.txt'
    run_cmd(cmd)

def colmap_feature_match(colmap, path, gpu=False):
    cmd = f'{colmap} exhaustive_matcher --database_path {path}/database.db \
--SiftMatching.guided_matching 0 \
--SiftMatching.max_ratio 0.8 \
--SiftMatching.max_distance 0.5 \
--SiftMatching.cross_check 1 \
--SiftMatching.max_error 4 \
--SiftMatching.max_num_matches 32768 \
--SiftMatching.confidence 0.9999 \
--SiftMatching.max_num_trials 10000 \
--SiftMatching.min_inlier_ratio 0.25 \
--SiftMatching.min_num_inliers 30'
    if gpu:
        cmd += ' --SiftMatching.use_gpu 1'
        cmd += ' --SiftMatching.gpu_index 0'
    cmd += f' >> {path}/log.txt'
    run_cmd(cmd)

def colmap_ba(colmap, path, with_init=False):
    if with_init:
        cmd = f'{colmap} point_triangulator --database_path {path}/database.db \
--image_path {path}/images \
--input_path {path}/sparse/0 \
--output_path {path}/sparse/0 \
--Mapper.tri_merge_max_reproj_error 3 \
--Mapper.ignore_watermarks 1 \
--Mapper.filter_max_reproj_error 2 \
>> {path}/log.txt'
        run_cmd(cmd)
        cmd = f'{colmap} bundle_adjuster \
--input_path {path}/sparse/0 \
--output_path {path}/sparse/0 \
--BundleAdjustment.max_num_iterations 1000 \
>> {path}/log.txt'
        run_cmd(cmd)
        points3d = read_points3d_binary(join(path, 'sparse', '0', 'points3D.bin'))
        pids = list(points3d.keys())
        mean_error = np.mean([points3d[p].error for p in pids])
        log('Triangulate {} points, mean error: {:.2f} pixel'.format(len(pids), mean_error))
    else:
        mkdir(join(path, 'sparse'))
        cmd = f'{colmap} mapper --database_path {path}/database.db --image_path {path}/images --output_path {path}/sparse \
    --Mapper.ba_refine_principal_point 1 \
    --Mapper.ba_global_max_num_iterations 1000 \
    >> {path}/log.txt'
        run_cmd(cmd)
    

def colmap_dense(colmap, path):
    mkdir(join(path, 'dense'))
    cmd = f'{colmap} image_undistorter --image_path {path}/images --input_path {path}/sparse/0 --output_path {path}/dense --output_type COLMAP --max_image_size 2000'
    run_cmd(cmd)
    cmd = f'{colmap} patch_match_stereo \
--workspace_path {path}/dense \
--workspace_format COLMAP \
--PatchMatchStereo.geom_consistency true \
>> {path}/log.txt'

    run_cmd(cmd)        
    cmd = f'{colmap} stereo_fusion \
--workspace_path {path}/dense \
--workspace_format COLMAP \
--input_type geometric \
--output_path {path}/dense/fused.ply \
>> {path}/log.txt'
    run_cmd(cmd)