Spaces:
Sleeping
Sleeping
| 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) | |
| # load camera intrinsics from the COLMAP reconstruction | |
| 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() | |
| # load image, depth map, and normal map | |
| 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. | |
| # create 3D points | |
| #depth_map = np.minimum(depth_map, 100.) | |
| points3D = np.dstack(camera.get_image_grid() + [depth_map]) | |
| points3D[:,:,:2] *= depth_map[:,:,np.newaxis] | |
| # save | |
| 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) | |