From 36e9034c9fb02ce4cf863a2b690f2382de223d95 Mon Sep 17 00:00:00 2001 From: Riley Date: Wed, 25 Oct 2023 16:41:43 -0700 Subject: [PATCH] Ply dataparser (#2557) * ply position and color loading * removed scripts dir * nerfacto bugfix in viewer * add world scale export option for pointcloud --------- Co-authored-by: Justin Kerr --- .../data/dataparsers/nerfstudio_dataparser.py | 31 +++++++++++++++++++ nerfstudio/scripts/exporter.py | 12 +++++++ nerfstudio/viewer_beta/export_panel.py | 2 ++ .../viewer_beta/render_state_machine.py | 2 +- 4 files changed, 46 insertions(+), 1 deletion(-) diff --git a/nerfstudio/data/dataparsers/nerfstudio_dataparser.py b/nerfstudio/data/dataparsers/nerfstudio_dataparser.py index eaeda0bce..9fbb950c6 100644 --- a/nerfstudio/data/dataparsers/nerfstudio_dataparser.py +++ b/nerfstudio/data/dataparsers/nerfstudio_dataparser.py @@ -20,6 +20,7 @@ from pathlib import Path from typing import Literal, Optional, Type import numpy as np +import open3d as o3d import torch from PIL import Image @@ -309,6 +310,12 @@ class Nerfstudio(DataParser): applied_scale = float(meta["applied_scale"]) scale_factor *= applied_scale + # Load 3D points + metadata = {} + if "ply_file_path" in meta: + ply_file_path = data_dir / meta["ply_file_path"] + metadata.update(self._load_3D_points(ply_file_path, transform_matrix, scale_factor)) + dataparser_outputs = DataparserOutputs( image_filenames=image_filenames, cameras=cameras, @@ -319,10 +326,34 @@ class Nerfstudio(DataParser): metadata={ "depth_filenames": depth_filenames if len(depth_filenames) > 0 else None, "depth_unit_scale_factor": self.config.depth_unit_scale_factor, + **metadata, }, ) return dataparser_outputs + def _load_3D_points(self, ply_file_path: Path, transform_matrix: torch.Tensor, scale_factor: float): + pcd = o3d.io.read_point_cloud(str(ply_file_path)) + + points3D = torch.from_numpy(np.asarray(pcd.points, dtype=np.float32)) + points3D = ( + torch.cat( + ( + points3D, + torch.ones_like(points3D[..., :1]), + ), + -1, + ) + @ transform_matrix.T + ) + points3D *= scale_factor + points3D_rgb = torch.from_numpy((np.asarray(pcd.colors) * 255).astype(np.uint8)) + + out = { + "points3D_xyz": points3D, + "points3D_rgb": points3D_rgb, + } + return out + def _get_fname(self, filepath: Path, data_dir: Path, downsample_folder_prefix="images_") -> Path: """Get the filename of the image file. downsample_folder_prefix can be used to point to auxiliary image data, e.g. masks diff --git a/nerfstudio/scripts/exporter.py b/nerfstudio/scripts/exporter.py index 64adc5c9d..821dd952e 100644 --- a/nerfstudio/scripts/exporter.py +++ b/nerfstudio/scripts/exporter.py @@ -121,6 +121,9 @@ class ExportPointCloud(Exporter): """Number of rays to evaluate per batch. Decrease if you run out of memory.""" std_ratio: float = 10.0 """Threshold based on STD of the average distances across the point cloud to remove outliers.""" + save_world_frame: bool = True + """If true, saves in the frame of the transform.json file, if false saves in the frame of the scaled + dataparser transform""" def main(self) -> None: """Export point cloud.""" @@ -157,6 +160,15 @@ class ExportPointCloud(Exporter): crop_obb=crop_obb, std_ratio=self.std_ratio, ) + if self.save_world_frame: + # apply the inverse dataparser transform to the point cloud + points = np.asarray(pcd.points) + poses = np.eye(4,dtype=np.float32)[None,...].repeat(points.shape[0], axis=0)[:,:3,:] + poses[:,:3,3] = points + poses = pipeline.datamanager.train_dataparser_outputs.transform_poses_to_original_space(torch.from_numpy(poses)) + points = poses[:,:3,3].numpy() + pcd.points = o3d.utility.Vector3dVector(points) + torch.cuda.empty_cache() CONSOLE.print(f"[bold green]:white_check_mark: Generated {pcd}") diff --git a/nerfstudio/viewer_beta/export_panel.py b/nerfstudio/viewer_beta/export_panel.py index 425e10a19..5bb0fc3b3 100644 --- a/nerfstudio/viewer_beta/export_panel.py +++ b/nerfstudio/viewer_beta/export_panel.py @@ -81,6 +81,7 @@ def populate_point_cloud_tab( ) -> None: server.add_gui_markdown("Render depth, project to an oriented point cloud, and filter. ") num_points = server.add_gui_number("# Points", initial_value=1_000_000, min=1, max=None, step=1) + world_frame = server.add_gui_checkbox("Save in world frame", True, hint="Save the point cloud in the transforms.json frame, rather than scaled scene frame") remove_outliers = server.add_gui_checkbox("Remove outliers", True) normals = server.add_gui_dropdown( "Normals", @@ -104,6 +105,7 @@ def populate_point_cloud_tab( f"--remove-outliers {remove_outliers.value}", f"--normal-method {normals.value}", f"--use_bounding_box {control_panel.crop_viewport}", + f"--save-world-frame {world_frame.value}", get_crop_string(control_panel.crop_obb), ] ) diff --git a/nerfstudio/viewer_beta/render_state_machine.py b/nerfstudio/viewer_beta/render_state_machine.py index 0c0bd66ca..88098a59a 100644 --- a/nerfstudio/viewer_beta/render_state_machine.py +++ b/nerfstudio/viewer_beta/render_state_machine.py @@ -179,7 +179,7 @@ class RenderStateMachine(threading.Thread): pass else: # convert to z_depth if depth compositing is enabled - R = camera.camera_to_worlds[0:3, 0:3].T + R = camera.camera_to_worlds[0,0:3, 0:3].T pts = camera_ray_bundle.directions * outputs["depth"] pts = (R @ (pts.view(-1, 3).T)).T.view(*camera_ray_bundle.directions.shape) outputs["gl_z_buf_depth"] = -pts[..., 2:3] # negative z axis is the coordinate convention