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 31# This script converts a VisualSfM reconstruction file to a PLY point cloud.32 33import argparse34 35import numpy as np36 37 38def parse_args():39    parser = argparse.ArgumentParser()40    parser.add_argument("--nvm_path", required=True)41    parser.add_argument("--ply_path", required=True)42    parser.add_argument("--normalize", type=bool, default=True)43    parser.add_argument("--normalize_p0", type=float, default=0.2)44    parser.add_argument("--normalize_p1", type=float, default=0.8)45    parser.add_argument("--min_track_length", type=int, default=3)46    args = parser.parse_args()47    return args48 49 50def main():51    args = parse_args()52 53    with open(args.nvm_path, "r") as fid:54        fid.readline()55        fid.readline()56        num_images = int(fid.readline())57 58        for i in range(num_images + 1):59            fid.readline()60 61        num_points = int(fid.readline())62 63        xyz = np.zeros((num_points, 3), dtype=np.float64)64        rgb = np.zeros((num_points, 3), dtype=np.uint16)65        track_lengths = np.zeros((num_points,), dtype=np.uint32)66 67        for i in range(num_points):68            if i % 1000 == 0:69                print("Reading point", i, "/", num_points)70            elems = fid.readline().split()71            xyz[i] = map(float, elems[0:3])72            rgb[i] = map(int, elems[3:6])73            track_lengths[i] = int(elems[6])74 75    mask = track_lengths >= args.min_track_length76    xyz = xyz[mask]77    rgb = rgb[mask]78 79    if args.normalize:80        sorted_x = np.sort(xyz[:, 0])81        sorted_y = np.sort(xyz[:, 1])82        sorted_z = np.sort(xyz[:, 2])83 84        num_coords = sorted_x.size85        min_coord = int(args.normalize_p0 * num_coords)86        max_coord = int(args.normalize_p1 * num_coords)87        mean_coords = xyz.mean(0)88 89        bbox_min = np.array(90            [sorted_x[min_coord], sorted_y[min_coord], sorted_z[min_coord]]91        )92        bbox_max = np.array(93            [sorted_x[max_coord], sorted_y[max_coord], sorted_z[max_coord]]94        )95 96        extent = np.linalg.norm(bbox_max - bbox_min)97        scale = 10.0 / extent98 99        xyz -= mean_coords100        xyz *= scale101 102    with open(args.ply_path, "w") as fid:103        fid.write("ply\n")104        fid.write("format ascii 1.0\n")105        fid.write("element vertex %d\n" % xyz.shape[0])106        fid.write("property float x\n")107        fid.write("property float y\n")108        fid.write("property float z\n")109        fid.write("property float nx\n")110        fid.write("property float ny\n")111        fid.write("property float nz\n")112        fid.write("property uchar diffuse_red\n")113        fid.write("property uchar diffuse_green\n")114        fid.write("property uchar diffuse_blue\n")115        fid.write("end_header\n")116        for i in range(xyz.shape[0]):117            if i % 1000 == 0:118                print("Writing point", i, "/", xyz.shape[0])119            fid.write(120                "%f %f %f 0 0 0 %d %d %d\n"121                % (122                    xyz[i, 0],123                    xyz[i, 1],124                    xyz[i, 2],125                    rgb[i, 0],126                    rgb[i, 1],127                    rgb[i, 2],128                )129            )130 131 132if __name__ == "__main__":133    main()134