From e25f7b33340d517540f05dab563d1a56057c875c Mon Sep 17 00:00:00 2001
From: hofee <64160135+GitHofee@users.noreply.github.com>
Date: Wed, 23 Oct 2024 00:42:18 -0500
Subject: [PATCH] add save preprocessed normals

---
 preprocess/preprocessor.py | 9 ++++++++-
 utils/pts.py               | 6 +++---
 2 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/preprocess/preprocessor.py b/preprocess/preprocessor.py
index d8d951b..0fccce1 100644
--- a/preprocess/preprocessor.py
+++ b/preprocess/preprocessor.py
@@ -21,6 +21,12 @@ def save_target_points(root, scene, frame_idx, target_points: np.ndarray, file_t
     if not os.path.exists(os.path.join(root,scene, "pts")):
         os.makedirs(os.path.join(root,scene, "pts"))
     save_np_pts(pts_path, target_points, file_type)
+
+def save_target_normals(root, scene, frame_idx, target_normals: np.ndarray, file_type="txt"):
+    pts_path = os.path.join(root,scene, "normals", f"{frame_idx}.{file_type}")
+    if not os.path.exists(os.path.join(root,scene, "normals")):
+        os.makedirs(os.path.join(root,scene, "normals"))
+    save_np_pts(pts_path, target_normals, file_type)
     
 def save_scan_points_indices(root, scene, frame_idx, scan_points_indices: np.ndarray, file_type="txt"):
     indices_path = os.path.join(root,scene, "scan_points_indices", f"{frame_idx}.{file_type}")
@@ -135,7 +141,7 @@ def save_scene_data(root, scene, scene_idx=0, scene_total=1,file_type="txt"):
             has_points = target_points.shape[0] > 0
 
         if has_points:
-            target_points = PtsUtil.filter_points(
+            target_points, target_normals = PtsUtil.filter_points(
                 target_points, sampled_target_normal_L, cam_info["cam_to_world"], theta_limit = filter_degree, z_range=(min_z, max_z)
                 )
 
@@ -149,6 +155,7 @@ def save_scene_data(root, scene, scene_idx=0, scene_total=1,file_type="txt"):
             target_points = np.zeros((0, 3))
             
         save_target_points(root, scene, frame_id, target_points, file_type=file_type)
+        save_target_normals(root, scene, frame_id, target_normals, file_type=file_type)
         save_scan_points_indices(root, scene, frame_id, scan_points_indices, file_type=file_type)
     
     save_scan_points(root, scene, scan_points) # The "done" flag of scene preprocess
diff --git a/utils/pts.py b/utils/pts.py
index 30f8860..e5cb436 100644
--- a/utils/pts.py
+++ b/utils/pts.py
@@ -84,14 +84,14 @@ class PtsUtil:
         theta = np.arccos(cos_theta) * 180 / np.pi
         idx = theta < theta_limit
         filtered_sampled_points = points[idx]
-        
+        filtered_normals = normals[idx]
         
         """ filter with z range """
         points_cam = PtsUtil.transform_point_cloud(filtered_sampled_points, np.linalg.inv(cam_pose))
         idx = (points_cam[:, 2] > z_range[0]) & (points_cam[:, 2] < z_range[1])
         z_filtered_points = filtered_sampled_points[idx]
-        
-        return z_filtered_points[:, :3]
+        z_filtered_normals = filtered_normals[idx]
+        return z_filtered_points[:, :3], z_filtered_normals
     
     @staticmethod
     def point_to_hash(point, voxel_size):