From 7acaed11a0e740994aab73be4e83f7f4c3fd7f64 Mon Sep 17 00:00:00 2001
From: Victor Rong <47129420+notagoose@users.noreply.github.com>
Date: Thu, 13 Jul 2023 15:01:53 -0700
Subject: [PATCH] Reorient point cloud normals for better surface
 reconstruction (#2220)

Implemented normal reorientation
---
 nerfstudio/exporter/exporter_utils.py | 17 +++++++++++++++++
 nerfstudio/scripts/exporter.py        |  6 ++++++
 2 files changed, 23 insertions(+)

diff --git a/nerfstudio/exporter/exporter_utils.py b/nerfstudio/exporter/exporter_utils.py
index 68211b5c3..21ea545b9 100644
--- a/nerfstudio/exporter/exporter_utils.py
+++ b/nerfstudio/exporter/exporter_utils.py
@@ -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
 
 
diff --git a/nerfstudio/scripts/exporter.py b/nerfstudio/scripts/exporter.py
index 153ab4e4b..d5ccd82bb 100644
--- a/nerfstudio/scripts/exporter.py
+++ b/nerfstudio/scripts/exporter.py
@@ -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,