fix normal

.gitignore

@@ -3,7 +3,7 @@
 __pycache__/
 *.py[cod]
 *$py.class
-
+temp*
 # C extensions
 *.so
 

data_generator.py

@@ -102,7 +102,6 @@ class DataGenerator:
         
         bpy.ops.rigidbody.object_add()
         bpy.context.object.rigid_body.type = 'PASSIVE'
-        bpy.ops.object.shade_auto_smooth()
         
         MaterialUtil.change_object_material(platform, MaterialUtil.create_mask_material(color=(1.0, 0, 0)))
 
@@ -114,6 +113,7 @@ class DataGenerator:
         return platform
 
     def put_display_object(self, name):
+        
         config = self.random_config["display_object"]
 
         x = random.uniform(config["min_x"], config["max_x"])
@@ -133,10 +133,8 @@ class DataGenerator:
         platform_bbox = self.platform.bound_box
         platform_bbox_world = [self.platform.matrix_world @ mathutils.Vector(corner) for corner in platform_bbox]
         platform_top_z = max([v.z for v in platform_bbox_world])
-        
         obj_mesh_path = BlenderUtils.get_obj_path(self.obj_dir, name)
         obj = BlenderUtils.load_obj(name, obj_mesh_path)
-        
         obj_bottom_z = BlenderUtils.get_object_bottom_z(obj)
         offset_z = obj_bottom_z
         
@@ -205,6 +203,7 @@ class DataGenerator:
         BlenderUtils.save_scene_info(scene_dir, self.display_table_config, object_name)
         
         MaterialUtil.change_object_material(self.target_obj, MaterialUtil.create_normal_material())
+        
         for i, cam_pose in enumerate(view_data["cam_poses"]):
             BlenderUtils.set_camera_at(cam_pose)
             BlenderUtils.render_normal_and_depth(scene_dir, f"{i}", binocular_vision=self.binocular_vision, target_object = self.target_obj)
@@ -219,6 +218,8 @@ class DataGenerator:
                 file_path = os.path.join(depth_dir, depth_file)
                 new_file_path = os.path.join(depth_dir, f"{name}.png")
                 os.rename(file_path,new_file_path)
+        BlenderUtils.save_blend(scene_dir)
+        exit(0)
         
         return True
     
@@ -266,6 +267,7 @@ class DataGenerator:
         if diag > self.max_diag or diag < self.min_diag:
             self.add_log(f"The diagonal size of the object <{object_name}>(size: {round(diag,3)}) does not meet the requirements.", "error")
             return "diag_error"
+        
         return self.simulate_scene(diag=diag)
 
     

test/tempdir/params.json

@@ -1,29 +0,0 @@
-{
-    "cam_pose": [
-        [
-            -0.8127143979072571,
-            -0.3794165253639221,
-            0.4421972334384918,
-            0.2740877568721771
-        ],
-        [
-            0.5826622247695923,
-            -0.5292212963104248,
-            0.616789698600769,
-            0.46910107135772705
-        ],
-        [
-            0.0,
-            0.7589254975318909,
-            0.6511774659156799,
-            1.2532192468643188
-        ],
-        [
-            0.0,
-            0.0,
-            0.0,
-            1.0
-        ]
-    ],
-    "scene_path": "/media/hofee/data/project/python/nbv_reconstruction/sample_for_training/scenes/google_scan-backpack_0288"
-}

utils/blender_util.py

@@ -66,7 +66,7 @@ class BlenderUtils:
 
     @staticmethod
     def setup_scene(init_light_and_camera_config, table_model_path, binocular_vision):
-        bpy.context.scene.render.engine = "BLENDER_EEVEE_NEXT"
+        bpy.context.scene.render.engine = "BLENDER_EEVEE"
         bpy.context.scene.display.shading.show_xray = False
         bpy.context.scene.display.shading.use_dof = False
         bpy.context.scene.display.render_aa = "OFF"
@@ -198,6 +198,8 @@ class BlenderUtils:
     def render_normal_and_depth(
         output_dir, file_name, binocular_vision=False, target_object=None
     ):
+        # use pass z
+        bpy.context.scene.view_layers["ViewLayer"].use_pass_z = True
         target_cameras = [BlenderUtils.CAMERA_NAME]
         if binocular_vision:
             target_cameras.append(BlenderUtils.CAMERA_RIGHT_NAME)
@@ -213,9 +215,14 @@ class BlenderUtils:
                 os.makedirs(mask_dir)
 
             scene.render.filepath = os.path.join(
-                output_dir, mask_dir, f"{file_name}_{cam_suffix}.png"
-            )
-            scene.render.image_settings.color_depth = "8"
+                output_dir, mask_dir, f"{file_name}_{cam_suffix}.exr"
+            ) 
+            
+            scene.render.image_settings.file_format = "OPEN_EXR"
+            scene.render.image_settings.color_mode = "RGB"
+            bpy.context.scene.view_settings.view_transform = "Raw"
+            scene.render.image_settings.color_depth = "16"
+            bpy.context.scene.render.filter_size = 1.5
             scene.render.resolution_percentage = 100
             scene.render.use_overwrite = False
             scene.render.use_file_extension = False
@@ -258,8 +265,7 @@ class BlenderUtils:
         target_cameras = [BlenderUtils.CAMERA_NAME]
         if binocular_vision:
             target_cameras.append(BlenderUtils.CAMERA_RIGHT_NAME)
-        # use pass z
-        bpy.context.scene.view_layers["ViewLayer"].use_pass_z = True
+        
         for cam_name in target_cameras:
             bpy.context.scene.camera = BlenderUtils.get_obj(cam_name)
             cam_suffix = "L" if cam_name == BlenderUtils.CAMERA_NAME else "R"
@@ -388,4 +394,9 @@ class BlenderUtils:
         scene_info["target_name"] = target_name
         scene_info_path = os.path.join(scene_root_dir, "scene_info.json")
         with open(scene_info_path, "w") as outfile:
-            json.dump(scene_info, outfile)
+            json.dump(scene_info, outfile)
+            
+    @staticmethod
+    def save_blend(scene_root_dir):
+        blend_path = os.path.join(scene_root_dir, "scene.blend")
+        bpy.ops.wm.save_as_mainfile(filepath=blend_path)

utils/pose.py

@@ -149,3 +149,18 @@ class PoseUtil:
         if debug:
             print("uniform scale:", scale)
         return scale
+
+    @staticmethod
+    def rotation_matrix_from_axis_angle(axis, angle):
+        cos_angle = np.cos(angle)
+        sin_angle = np.sin(angle)
+        one_minus_cos = 1 - cos_angle
+
+        x, y, z = axis
+        rotation_matrix = np.array([
+            [cos_angle + x*x*one_minus_cos, x*y*one_minus_cos - z*sin_angle, x*z*one_minus_cos + y*sin_angle],
+            [y*x*one_minus_cos + z*sin_angle, cos_angle + y*y*one_minus_cos, y*z*one_minus_cos - x*sin_angle],
+            [z*x*one_minus_cos - y*sin_angle, z*y*one_minus_cos + x*sin_angle, cos_angle + z*z*one_minus_cos]
+        ])
+
+        return rotation_matrix

utils/pts.py

@@ -0,0 +1,83 @@
+import numpy as np
+
+
+class PtsUtil:
+    
+    @staticmethod
+    def random_downsample_point_cloud(point_cloud, num_points, require_idx=False):
+        if point_cloud.shape[0] == 0:
+            if require_idx:
+                return point_cloud, np.array([])
+            return point_cloud
+        idx = np.random.choice(len(point_cloud), num_points, replace=True)
+        if require_idx:
+            return point_cloud[idx], idx
+        return point_cloud[idx]
+    
+    @staticmethod
+    def fps_downsample_point_cloud(point_cloud, num_points, require_idx=False):
+        N = point_cloud.shape[0]
+        mask = np.zeros(N, dtype=bool)
+
+        sampled_indices = np.zeros(num_points, dtype=int)
+        sampled_indices[0] = np.random.randint(0, N)
+        distances = np.linalg.norm(point_cloud - point_cloud[sampled_indices[0]], axis=1)
+        for i in range(1, num_points):
+            farthest_index = np.argmax(distances)
+            sampled_indices[i] = farthest_index
+            mask[farthest_index] = True
+
+            new_distances = np.linalg.norm(point_cloud - point_cloud[farthest_index], axis=1)
+            distances = np.minimum(distances, new_distances)
+        
+        sampled_points = point_cloud[sampled_indices]
+        if require_idx:
+            return sampled_points, sampled_indices
+        return sampled_points
+    
+    @staticmethod
+    def voxelize_points(points, voxel_size):
+        voxel_indices = np.floor(points / voxel_size).astype(np.int32)
+        unique_voxels = np.unique(voxel_indices, axis=0, return_inverse=True)
+        return unique_voxels
+
+    @staticmethod
+    def transform_point_cloud(points, pose_mat):
+        points_h = np.concatenate([points, np.ones((points.shape[0], 1))], axis=1)
+        points_h = np.dot(pose_mat, points_h.T).T
+        return points_h[:, :3]
+    
+    @staticmethod
+    def get_overlapping_points(point_cloud_L, point_cloud_R, voxel_size=0.005, require_idx=False):
+        voxels_L, indices_L = PtsUtil.voxelize_points(point_cloud_L, voxel_size)
+        voxels_R, _ = PtsUtil.voxelize_points(point_cloud_R, voxel_size)
+
+        voxel_indices_L = voxels_L.view([("", voxels_L.dtype)] * 3)
+        voxel_indices_R = voxels_R.view([("", voxels_R.dtype)] * 3)
+        overlapping_voxels = np.intersect1d(voxel_indices_L, voxel_indices_R)
+        mask_L = np.isin(
+            indices_L, np.where(np.isin(voxel_indices_L, overlapping_voxels))[0]
+        )
+        overlapping_points = point_cloud_L[mask_L]
+        if require_idx:
+            return overlapping_points, mask_L
+        return overlapping_points
+    
+    @staticmethod
+    def filter_points(points, normals, cam_pose, theta=45, z_range=(0.2, 0.45)):
+        
+        """ filter with normal """ 
+        normals_normalized = normals / np.linalg.norm(normals, axis=1, keepdims=True)
+        cos_theta = np.dot(normals_normalized, np.array([0, 0, 1]))
+        theta_rad = np.deg2rad(theta)
+        idx = cos_theta > np.cos(theta_rad)
+        filtered_sampled_points = points[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]
+        

utils/view_sample_util.py

@@ -4,6 +4,7 @@ import bmesh
 from collections import defaultdict
 from scipy.spatial.transform import Rotation as R
 from utils.pose import PoseUtil
+from utils.pts import PtsUtil
 import random
 
 class ViewSampleUtil:
@@ -71,7 +72,7 @@ class ViewSampleUtil:
             normals.append(normal)
             
             for _ in range(pertube_repeat):
-                perturb_angle = np.radians(np.random.uniform(0, 30))
+                perturb_angle = np.radians(np.random.uniform(0, 10))
                 perturb_axis = np.random.normal(size=3)
                 perturb_axis /= np.linalg.norm(perturb_axis)
                 rotation_matrix = R.from_rotvec(perturb_angle * perturb_axis).as_matrix()
@@ -127,12 +128,27 @@ class ViewSampleUtil:
             
             up_vector = np.array([0, 0, 1])
 
-            right_vector = np.cross(up_vector, forward_vector)
-            right_vector /= np.linalg.norm(right_vector)
+            dot_product = np.dot(forward_vector, up_vector)
+            angle = np.degrees(np.arccos(dot_product))
+            print(angle)
+            if angle < 110:
+                target_angle = np.random.uniform(110, 150)
+                angle_difference = np.radians(target_angle - angle)
 
-            corrected_up_vector = np.cross(forward_vector, right_vector)
-            rotation_matrix = np.array([right_vector, corrected_up_vector, forward_vector]).T
+                rotation_axis = np.cross(forward_vector, up_vector)
+                rotation_axis /= np.linalg.norm(rotation_axis)
+                rotation_matrix = PoseUtil.rotation_matrix_from_axis_angle(rotation_axis, angle_difference)
+                new_cam_position_world = np.dot(rotation_matrix, cam_position_world - look_at_point_world) + look_at_point_world
+                cam_position_world = new_cam_position_world
+                forward_vector = cam_position_world - look_at_point_world
+                forward_vector /= np.linalg.norm(forward_vector)
+                right_vector = np.cross(up_vector, forward_vector)
+                right_vector /= np.linalg.norm(right_vector)
 
+                corrected_up_vector = np.cross(forward_vector, right_vector)
+                rotation_matrix = np.array([right_vector, corrected_up_vector, forward_vector]).T
+            else:
+                rotation_matrix = np.array([right_vector, up_vector, forward_vector]).T
             cam_pose = np.eye(4)
             cam_pose[:3, :3] = rotation_matrix
             cam_pose[:3, 3] = cam_position_world
@@ -146,16 +162,17 @@ class ViewSampleUtil:
                 cos_angle = np.dot(direction_vector, horizontal_normal) / (np.linalg.norm(direction_vector) * np.linalg.norm(horizontal_normal))
                 angle = np.arccos(np.clip(cos_angle, -1.0, 1.0))  
                 angle_degree = np.degrees(angle)
-                if angle_degree < 90 - min_cam_table_included_degree:
+                if angle_degree < 90 + min_cam_table_included_degree:
                     filtered_cam_poses.append(cam_pose)
                 if random.random() < random_view_ratio:
                     pertube_pose = PoseUtil.get_uniform_pose([0.1, 0.1, 0.1], [3, 3, 3], 0, 180, "cm")
                     filtered_cam_poses.append(pertube_pose @ cam_pose)
                     
         if len(filtered_cam_poses) > max_views:
-            indices = np.random.choice(len(filtered_cam_poses), max_views, replace=False)
+            cam_points = np.array([cam_pose[:3, 3] for cam_pose in filtered_cam_poses])
+            _, indices = PtsUtil.fps_downsample_point_cloud(cam_points, max_views, require_idx=True)
             filtered_cam_poses = [filtered_cam_poses[i] for i in indices]
-            
+        
         return np.array(filtered_cam_poses)
         
     @staticmethod