diff --git a/.gitignore b/.gitignore index 5d381cc..b3a6109 100644 --- a/.gitignore +++ b/.gitignore @@ -3,7 +3,7 @@ __pycache__/ *.py[cod] *$py.class - +temp* # C extensions *.so diff --git a/data_generator.py b/data_generator.py index 434ff0c..28735a9 100644 --- a/data_generator.py +++ b/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) diff --git a/test/tempdir/depth/tmp_L.png b/test/tempdir/depth/tmp_L.png deleted file mode 100644 index 8506846..0000000 Binary files a/test/tempdir/depth/tmp_L.png and /dev/null differ diff --git a/test/tempdir/depth/tmp_R.png b/test/tempdir/depth/tmp_R.png deleted file mode 100644 index 1fee75a..0000000 Binary files a/test/tempdir/depth/tmp_R.png and /dev/null differ diff --git a/test/tempdir/mask/tmp_L.png b/test/tempdir/mask/tmp_L.png deleted file mode 100644 index aba9186..0000000 Binary files a/test/tempdir/mask/tmp_L.png and /dev/null differ diff --git a/test/tempdir/mask/tmp_R.png b/test/tempdir/mask/tmp_R.png deleted file mode 100644 index c1e34d0..0000000 Binary files a/test/tempdir/mask/tmp_R.png and /dev/null differ diff --git a/test/tempdir/params.json b/test/tempdir/params.json deleted file mode 100644 index 5d1cf97..0000000 --- a/test/tempdir/params.json +++ /dev/null @@ -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" -} \ No newline at end of file diff --git a/utils/blender_util.py b/utils/blender_util.py index 47d016a..080670c 100644 --- a/utils/blender_util.py +++ b/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) \ No newline at end of file + 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) \ No newline at end of file diff --git a/utils/pose.py b/utils/pose.py index a9373b2..63cf6ab 100644 --- a/utils/pose.py +++ b/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 \ No newline at end of file diff --git a/utils/pts.py b/utils/pts.py new file mode 100644 index 0000000..860734d --- /dev/null +++ b/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] + \ No newline at end of file diff --git a/utils/view_sample_util.py b/utils/view_sample_util.py index c6836d7..d7b34c2 100644 --- a/utils/view_sample_util.py +++ b/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