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 30import argparse31 32import numpy as np33import open3d34from read_write_model import qvec2rotmat, read_model35 36 37class Model:38    def __init__(self):39        self.cameras = []40        self.images = []41        self.points3D = []42        self.__vis = None43 44    def read_model(self, path, ext=""):45        self.cameras, self.images, self.points3D = read_model(path, ext)46 47    def add_points(self, min_track_len=3, remove_statistical_outlier=True):48        pcd = open3d.geometry.PointCloud()49 50        xyz = []51        rgb = []52        for point3D in self.points3D.values():53            track_len = len(point3D.point2D_idxs)54            if track_len < min_track_len:55                continue56            xyz.append(point3D.xyz)57            rgb.append(point3D.rgb / 255)58 59        pcd.points = open3d.utility.Vector3dVector(xyz)60        pcd.colors = open3d.utility.Vector3dVector(rgb)61 62        # remove obvious outliers63        if remove_statistical_outlier:64            [pcd, _] = pcd.remove_statistical_outlier(65                nb_neighbors=20, std_ratio=2.066            )67 68        # open3d.visualization.draw_geometries([pcd])69        self.__vis.add_geometry(pcd)70        self.__vis.poll_events()71        self.__vis.update_renderer()72 73    def add_cameras(self, scale=1):74        frames = []75        for img in self.images.values():76            # rotation77            R = qvec2rotmat(img.qvec)78 79            # translation80            t = img.tvec81 82            # invert83            t = -R.T @ t84            R = R.T85 86            # intrinsics87            cam = self.cameras[img.camera_id]88 89            if cam.model in ("SIMPLE_PINHOLE", "SIMPLE_RADIAL", "RADIAL"):90                fx = fy = cam.params[0]91                cx = cam.params[1]92                cy = cam.params[2]93            elif cam.model in (94                "PINHOLE",95                "OPENCV",96                "OPENCV_FISHEYE",97                "FULL_OPENCV",98            ):99                fx = cam.params[0]100                fy = cam.params[1]101                cx = cam.params[2]102                cy = cam.params[3]103            else:104                raise Exception("Camera model not supported")105 106            # intrinsics107            K = np.identity(3)108            K[0, 0] = fx109            K[1, 1] = fy110            K[0, 2] = cx111            K[1, 2] = cy112 113            # create axis, plane and pyramed geometries that will be drawn114            cam_model = draw_camera(K, R, t, cam.width, cam.height, scale)115            frames.extend(cam_model)116 117        # add geometries to visualizer118        for i in frames:119            self.__vis.add_geometry(i)120 121    def create_window(self):122        self.__vis = open3d.visualization.Visualizer()123        self.__vis.create_window()124 125    def show(self):126        self.__vis.poll_events()127        self.__vis.update_renderer()128        self.__vis.run()129        self.__vis.destroy_window()130 131 132def draw_camera(K, R, t, w, h, scale=1, color=[0.8, 0.2, 0.8]):133    """Create axis, plane and pyramed geometries in Open3D format.134    :param K: calibration matrix (camera intrinsics)135    :param R: rotation matrix136    :param t: translation137    :param w: image width138    :param h: image height139    :param scale: camera model scale140    :param color: color of the image plane and pyramid lines141    :return: camera model geometries (axis, plane and pyramid)142    """143 144    # intrinsics145    K = K.copy() / scale146    Kinv = np.linalg.inv(K)147 148    # 4x4 transformation149    T = np.column_stack((R, t))150    T = np.vstack((T, (0, 0, 0, 1)))151 152    # axis153    axis = open3d.geometry.TriangleMesh.create_coordinate_frame(154        size=0.5 * scale155    )156    axis.transform(T)157 158    # points in pixel159    points_pixel = [160        [0, 0, 0],161        [0, 0, 1],162        [w, 0, 1],163        [0, h, 1],164        [w, h, 1],165    ]166 167    # pixel to camera coordinate system168    points = [Kinv @ p for p in points_pixel]169 170    # image plane171    width = abs(points[1][0]) + abs(points[3][0])172    height = abs(points[1][1]) + abs(points[3][1])173    plane = open3d.geometry.TriangleMesh.create_box(width, height, depth=1e-6)174    plane.paint_uniform_color(color)175    plane.translate([points[1][0], points[1][1], scale])176    plane.transform(T)177 178    # pyramid179    points_in_world = [(R @ p + t) for p in points]180    lines = [181        [0, 1],182        [0, 2],183        [0, 3],184        [0, 4],185    ]186    colors = [color for i in range(len(lines))]187    line_set = open3d.geometry.LineSet(188        points=open3d.utility.Vector3dVector(points_in_world),189        lines=open3d.utility.Vector2iVector(lines),190    )191    line_set.colors = open3d.utility.Vector3dVector(colors)192 193    # return as list in Open3D format194    return [axis, plane, line_set]195 196 197def parse_args():198    parser = argparse.ArgumentParser(199        description="Visualize COLMAP binary and text models"200    )201    parser.add_argument(202        "--input_model", required=True, help="path to input model folder"203    )204    parser.add_argument(205        "--input_format",206        choices=[".bin", ".txt"],207        help="input model format",208        default="",209    )210    args = parser.parse_args()211    return args212 213 214def main():215    args = parse_args()216 217    # read COLMAP model218    model = Model()219    model.read_model(args.input_model, ext=args.input_format)220 221    print("num_cameras:", len(model.cameras))222    print("num_images:", len(model.images))223    print("num_points3D:", len(model.points3D))224 225    # display using Open3D visualization tools226    model.create_window()227    model.add_points()228    model.add_cameras(scale=0.25)229    model.show()230 231 232if __name__ == "__main__":233    main()234