| import sys |
| sys.path.append("..") |
|
|
| import imageio |
| import numpy as np |
| import os |
|
|
| from plyfile import PlyData, PlyElement |
| from pycolmap import SceneManager |
| from scipy.ndimage.interpolation import zoom |
|
|
|
|
| |
|
|
| def main(args): |
| suffix = ".photometric.bin" if args.photometric else ".geometric.bin" |
|
|
| image_file = os.path.join(args.dense_folder, "images", args.image_filename) |
| depth_file = os.path.join( |
| args.dense_folder, args.stereo_folder, "depth_maps", |
| args.image_filename + suffix) |
| if args.save_normals: |
| normals_file = os.path.join( |
| args.dense_folder, args.stereo_folder, "normal_maps", |
| args.image_filename + suffix) |
|
|
| |
| scene_manager = SceneManager(os.path.join(args.dense_folder, "sparse")) |
| scene_manager.load_cameras() |
| scene_manager.load_images() |
|
|
| image_id, image = scene_manager.get_image_from_name(args.image_filename) |
| camera = scene_manager.cameras[image.camera_id] |
| rotation_camera_from_world = image.R() |
| camera_center = image.C() |
|
|
| |
| image = imageio.imread(image_file) |
|
|
| with open(depth_file, "rb") as fid: |
| w = int("".join(iter(lambda: fid.read(1), "&"))) |
| h = int("".join(iter(lambda: fid.read(1), "&"))) |
| c = int("".join(iter(lambda: fid.read(1), "&"))) |
| depth_map = np.fromfile(fid, np.float32).reshape(h, w) |
| if (h, w) != image.shape[:2]: |
| depth_map = zoom( |
| depth_map, |
| (float(image.shape[0]) / h, float(image.shape[1]) / w), |
| order=0) |
|
|
| if args.save_normals: |
| with open(normals_file, "rb") as fid: |
| w = int("".join(iter(lambda: fid.read(1), "&"))) |
| h = int("".join(iter(lambda: fid.read(1), "&"))) |
| c = int("".join(iter(lambda: fid.read(1), "&"))) |
| normals = np.fromfile( |
| fid, np.float32).reshape(c, h, w).transpose([1, 2, 0]) |
| if (h, w) != image.shape[:2]: |
| normals = zoom( |
| normals, |
| (float(image.shape[0]) / h, float(image.shape[1]) / w, 1.), |
| order=0) |
|
|
| if args.min_depth is not None: |
| depth_map[depth_map < args.min_depth] = 0. |
| if args.max_depth is not None: |
| depth_map[depth_map > args.max_depth] = 0. |
|
|
| |
| |
| points3D = np.dstack(camera.get_image_grid() + [depth_map]) |
| points3D[:,:,:2] *= depth_map[:,:,np.newaxis] |
|
|
| |
| points3D = points3D.astype(np.float32).reshape(-1, 3) |
| if args.save_normals: |
| normals = normals.astype(np.float32).reshape(-1, 3) |
| image = image.reshape(-1, 3) |
| if image.dtype != np.uint8: |
| if image.max() <= 1: |
| image = (image * 255.).astype(np.uint8) |
| else: |
| image = image.astype(np.uint8) |
|
|
| if args.world_space: |
| points3D = points3D.dot(rotation_camera_from_world) + camera_center |
| if args.save_normals: |
| normals = normals.dot(rotation_camera_from_world) |
|
|
| if args.save_normals: |
| vertices = np.rec.fromarrays( |
| tuple(points3D.T) + tuple(normals.T) + tuple(image.T), |
| names="x,y,z,nx,ny,nz,red,green,blue") |
| else: |
| vertices = np.rec.fromarrays( |
| tuple(points3D.T) + tuple(image.T), names="x,y,z,red,green,blue") |
| vertices = PlyElement.describe(vertices, "vertex") |
| PlyData([vertices]).write(args.output_filename) |
|
|
|
|
| |
|
|
| if __name__ == "__main__": |
| import argparse |
|
|
| parser = argparse.ArgumentParser( |
| formatter_class=argparse.ArgumentDefaultsHelpFormatter) |
|
|
| parser.add_argument("dense_folder", type=str) |
| parser.add_argument("image_filename", type=str) |
| parser.add_argument("output_filename", type=str) |
|
|
| parser.add_argument( |
| "--photometric", default=False, action="store_true", |
| help="use photometric depthmap instead of geometric") |
|
|
| parser.add_argument( |
| "--world_space", default=False, action="store_true", |
| help="apply the camera->world extrinsic transformation to the result") |
|
|
| parser.add_argument( |
| "--save_normals", default=False, action="store_true", |
| help="load the estimated normal map and save as part of the PLY") |
|
|
| parser.add_argument( |
| "--stereo_folder", type=str, default="stereo", |
| help="folder in the dense workspace containing depth and normal maps") |
|
|
| parser.add_argument( |
| "--min_depth", type=float, default=None, |
| help="set pixels with depth less than this value to zero depth") |
|
|
| parser.add_argument( |
| "--max_depth", type=float, default=None, |
| help="set pixels with depth greater than this value to zero depth") |
|
|
| args = parser.parse_args() |
|
|
| main(args) |
|
|
