beverley-gorry/colmap-vslamlab

1# Copyright (c), ETH Zurich and UNC Chapel Hill.2# All rights reserved.3#4# Redistribution and use in source and binary forms, with or without5# modification, are permitted provided that the following conditions are met:6#7#     * Redistributions of source code must retain the above copyright8#       notice, this list of conditions and the following disclaimer.9#10#     * Redistributions in binary form must reproduce the above copyright11#       notice, this list of conditions and the following disclaimer in the12#       documentation and/or other materials provided with the distribution.13#14#     * Neither the name of ETH Zurich and UNC Chapel Hill nor the names of15#       its contributors may be used to endorse or promote products derived16#       from this software without specific prior written permission.17#18# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"19# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE20# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE21# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE22# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR23# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF24# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS25# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN26# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)27# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE28# POSSIBILITY OF SUCH DAMAGE.29 30 31import argparse32 33import numpy as np34from database import COLMAPDatabase35 36if __name__ == "__main__":37    parser = argparse.ArgumentParser()38    parser.add_argument("--database_path", type=str, required=True)39    parser.add_argument("--is_cartesian", action="store_true")40    parser.add_argument("--cleanup", action="store_true")41    args = parser.parse_args()42 43    db = COLMAPDatabase.connect(args.database_path)44 45    pose_priors = {}46    rows = db.execute("SELECT * FROM images")47    for image_id, _, _, *cam_from_world_prior in rows:48        if not cam_from_world_prior:  # newer format database49            continue50        qvec = np.array(cam_from_world_prior[:4], dtype=float)51        tvec = np.array(cam_from_world_prior[4:], dtype=float)52        if np.isfinite(qvec).any():53            print(54                f"Warning: rotation prior for image {image_id} "55                "will be lost during migration."56            )57        if np.isfinite(tvec).any():58            pose_priors[image_id] = tvec59    print(f"Found location priors for {len(pose_priors)} images.")60 61    coordinate_systems = {"UNKNOWN": -1, "WGS84": 0, "CARTESIAN": 1}62    coordinate_system = coordinate_systems[63        "CARTESIAN" if args.is_cartesian else "WGS84"64    ]65    db.create_pose_priors_table()66    for image_id, position in pose_priors.items():67        (exists,) = db.execute(68            "SELECT COUNT(*) FROM pose_priors WHERE image_id = ?",69            (image_id,),70        ).fetchone()71        if exists:72            print(f"Location prior for {image_id} already exists, skipping.")73            continue74        db.add_pose_prior(image_id, position, coordinate_system)75 76    if args.cleanup:77        for col in ["qw", "qx", "qy", "qz", "tx", "ty", "tz"]:78            db.execute(f"ALTER TABLE images DROP COLUMN prior_{col}")79 80    db.commit()81