beverley-gorry/colmap-vslamlab

1#!/usr/bin/env python2 3# Copyright (c), ETH Zurich and UNC Chapel Hill.4# All rights reserved.5#6# Redistribution and use in source and binary forms, with or without7# modification, are permitted provided that the following conditions are met:8#9#     * Redistributions of source code must retain the above copyright10#       notice, this list of conditions and the following disclaimer.11#12#     * Redistributions in binary form must reproduce the above copyright13#       notice, this list of conditions and the following disclaimer in the14#       documentation and/or other materials provided with the distribution.15#16#     * Neither the name of ETH Zurich and UNC Chapel Hill nor the names of17#       its contributors may be used to endorse or promote products derived18#       from this software without specific prior written permission.19#20# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"21# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE22# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE23# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE24# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR25# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF26# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS27# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN28# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)29# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE30# POSSIBILITY OF SUCH DAMAGE.31 32 33import collections34import os35 36import numpy as np37import pandas as pd38from pyntcloud import PyntCloud39from read_write_model import read_next_bytes, write_next_bytes40 41MeshPoint = collections.namedtuple(42    "MeshingPoint",43    ["position", "color", "normal", "num_visible_images", "visible_image_idxs"],44)45 46 47def read_fused(path_to_fused_ply, path_to_fused_ply_vis):48    """49    see: src/mvs/meshing.cc50        void ReadDenseReconstruction(const std::string& path51    """52    assert os.path.isfile(path_to_fused_ply)53    assert os.path.isfile(path_to_fused_ply_vis)54 55    point_cloud = PyntCloud.from_file(path_to_fused_ply)56    xyz_arr = point_cloud.points.loc[:, ["x", "y", "z"]].to_numpy()57    normal_arr = point_cloud.points.loc[:, ["nx", "ny", "nz"]].to_numpy()58    color_arr = point_cloud.points.loc[:, ["red", "green", "blue"]].to_numpy()59 60    with open(path_to_fused_ply_vis, "rb") as fid:61        num_points = read_next_bytes(fid, 8, "Q")[0]62        mesh_points = [0] * num_points63        for i in range(num_points):64            num_visible_images = read_next_bytes(fid, 4, "I")[0]65            visible_image_idxs = read_next_bytes(66                fid,67                num_bytes=4 * num_visible_images,68                format_char_sequence="I" * num_visible_images,69            )70            visible_image_idxs = np.array(tuple(map(int, visible_image_idxs)))71            mesh_point = MeshPoint(72                position=xyz_arr[i],73                color=color_arr[i],74                normal=normal_arr[i],75                num_visible_images=num_visible_images,76                visible_image_idxs=visible_image_idxs,77            )78            mesh_points[i] = mesh_point79        return mesh_points80 81 82def write_fused_ply(mesh_points, path_to_fused_ply):83    columns = ["x", "y", "z", "nx", "ny", "nz", "red", "green", "blue"]84    points_data_frame = pd.DataFrame(85        np.zeros((len(mesh_points), len(columns))), columns=columns86    )87 88    positions = np.asarray([point.position for point in mesh_points])89    normals = np.asarray([point.normal for point in mesh_points])90    colors = np.asarray([point.color for point in mesh_points])91 92    points_data_frame.loc[:, ["x", "y", "z"]] = positions93    points_data_frame.loc[:, ["nx", "ny", "nz"]] = normals94    points_data_frame.loc[:, ["red", "green", "blue"]] = colors95 96    points_data_frame = points_data_frame.astype(97        {98            "x": positions.dtype,99            "y": positions.dtype,100            "z": positions.dtype,101            "red": colors.dtype,102            "green": colors.dtype,103            "blue": colors.dtype,104            "nx": normals.dtype,105            "ny": normals.dtype,106            "nz": normals.dtype,107        }108    )109 110    point_cloud = PyntCloud(points_data_frame)111    point_cloud.to_file(path_to_fused_ply)112 113 114def write_fused_ply_vis(mesh_points, path_to_fused_ply_vis):115    """116    see: src/mvs/fusion.cc117        void WritePointsVisibility(const std::string& path, const std::vector<std::vector<int>>& points_visibility)118    """119    with open(path_to_fused_ply_vis, "wb") as fid:120        write_next_bytes(fid, len(mesh_points), "Q")121        for point in mesh_points:122            write_next_bytes(fid, point.num_visible_images, "I")123            format_char_sequence = "I" * point.num_visible_images124            write_next_bytes(125                fid, [*point.visible_image_idxs], format_char_sequence126            )127 128 129def write_fused(points, path_to_fused_ply, path_to_fused_ply_vis):130    write_fused_ply(points, path_to_fused_ply)131    write_fused_ply_vis(points, path_to_fused_ply_vis)132