tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 10:34:29 +08:00
Code

Reorient point cloud normals for better surface reconstruction (#2220)

Browse Source 
Implemented normal reorientation
This commit is contained in:
Victor Rong 2023-07-13 15:01:53 -07:00 committed by GitHub
parent 554dec77f2
commit 7acaed11a0
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 23 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
nerfstudio/exporter/exporter_utils.py
View File

@ -85,6 +85,7 @@ def generate_point_cloud(
num_points: int = 1000000,
remove_outliers: bool = True,
estimate_normals: bool = False,
reorient_normals: bool = False,
rgb_output_name: str = "rgb",
depth_output_name: str = "depth",
normal_output_name: Optional[str] = None,
@ -99,6 +100,7 @@ def generate_point_cloud(
pipeline: Pipeline to evaluate with.
num_points: Number of points to generate. May result in less if outlier removal is used.
remove_outliers: Whether to remove outliers.
reorient_normals: Whether to re-orient the normals based on the view direction.
estimate_normals: Whether to estimate normals.
rgb_output_name: Name of the RGB output.
depth_output_name: Name of the depth output.
@ -122,6 +124,8 @@ def generate_point_cloud(
points = []
rgbs = []
normals = []
view_directions = []
with progress as progress_bar:
task = progress_bar.add_task("Generating Point Cloud", total=num_points)
while not progress_bar.finished:
@ -154,6 +158,7 @@ def generate_point_cloud(
), "Normal values from method output must be in [0, 1]"
normal = (normal * 2.0) - 1.0
point = ray_bundle.origins + ray_bundle.directions * depth
view_direction = ray_bundle.directions
if use_bounding_box:
comp_l = torch.tensor(bounding_box_min, device=point.device)
@ -164,16 +169,19 @@ def generate_point_cloud(
mask = torch.all(torch.concat([point > comp_l, point < comp_m], dim=-1), dim=-1)
point = point[mask]
rgb = rgb[mask]
view_direction = view_direction[mask]
if normal is not None:
normal = normal[mask]
points.append(point)
rgbs.append(rgb)
view_directions.append(view_direction)
if normal is not None:
normals.append(normal)
progress.advance(task, point.shape[0])
points = torch.cat(points, dim=0)
rgbs = torch.cat(rgbs, dim=0)
view_directions = torch.cat(view_directions, dim=0).cpu()
pcd = o3d.geometry.PointCloud()
pcd.points = o3d.utility.Vector3dVector(points.double().cpu().numpy())
@ -185,6 +193,8 @@ def generate_point_cloud(
pcd, ind = pcd.remove_statistical_outlier(nb_neighbors=20, std_ratio=std_ratio)
print("\033[A\033[A")
CONSOLE.print("[bold green]:white_check_mark: Cleaning Point Cloud")
if ind is not None:
view_directions = view_directions[ind]
# either estimate_normals or normal_output_name, not both
if estimate_normals:
@ -203,6 +213,13 @@ def generate_point_cloud(
normals = normals[ind]
pcd.normals = o3d.utility.Vector3dVector(normals.double().cpu().numpy())
# re-orient the normals
if reorient_normals:
normals = torch.from_numpy(np.array(pcd.normals)).float()
mask = torch.sum(view_directions * normals, dim=-1) > 0
normals[mask] *= -1
pcd.normals = o3d.utility.Vector3dVector(normals.double().cpu().numpy())
return pcd

6
nerfstudio/scripts/exporter.py
View File

@ -97,6 +97,8 @@ class ExportPointCloud(Exporter):
"""Number of points to generate. May result in less if outlier removal is used."""
remove_outliers: bool = True
"""Remove outliers from the point cloud."""
reorient_normals: bool = True
"""Reorient point cloud normals based on view direction."""
normal_method: Literal["open3d", "model_output"] = "model_output"
"""Method to estimate normals with."""
normal_output_name: str = "normals"
@ -138,6 +140,7 @@ class ExportPointCloud(Exporter):
pipeline=pipeline,
num_points=self.num_points,
remove_outliers=self.remove_outliers,
reorient_normals=self.reorient_normals,
estimate_normals=estimate_normals,
rgb_output_name=self.rgb_output_name,
depth_output_name=self.depth_output_name,
@ -243,6 +246,8 @@ class ExportPoissonMesh(Exporter):
"""Number of points to generate. May result in less if outlier removal is used."""
remove_outliers: bool = True
"""Remove outliers from the point cloud."""
reorient_normals: bool = True
"""Reorient point cloud normals based on view direction."""
depth_output_name: str = "depth"
"""Name of the depth output."""
rgb_output_name: str = "rgb"
@ -296,6 +301,7 @@ class ExportPoissonMesh(Exporter):
pipeline=pipeline,
num_points=self.num_points,
remove_outliers=self.remove_outliers,
reorient_normals=self.reorient_normals,
estimate_normals=estimate_normals,
rgb_output_name=self.rgb_output_name,
depth_output_name=self.depth_output_name,