From f94cb2c0d6993f43f76bd42c434f5069170f8893 Mon Sep 17 00:00:00 2001 From: hofee <64160135+GitHofee@users.noreply.github.com> Date: Sat, 28 Sep 2024 22:03:19 +0800 Subject: [PATCH] add min_cam_table_included_degree and random_view --- data_generator.py | 4 +- pose.py | 151 ++++++++++++++++++++++++++++++++++++++++++++ view_sample_util.py | 21 ++++-- 3 files changed, 170 insertions(+), 6 deletions(-) create mode 100644 pose.py diff --git a/data_generator.py b/data_generator.py index ff54d26..444041f 100644 --- a/data_generator.py +++ b/data_generator.py @@ -21,6 +21,8 @@ class DataGenerator: self.min_views = config["runner"]["generate"]["min_views"] self.min_diag = config["runner"]["generate"]["min_diag"] self.max_diag = config["runner"]["generate"]["max_diag"] + self.min_cam_table_included_degree = config["runner"]["generate"]["min_cam_table_included_degree"] + self.random_view_ratio = config["runner"]["generate"]["random_view_ratio"] self.binocular_vision = config["runner"]["generate"]["binocular_vision"] self.set_status_path = "http://localhost:5000/project/set_status" self.log_path = "http://localhost:5000/project/add_log" @@ -240,7 +242,7 @@ class DataGenerator: if not os.path.exists(scene_dir): os.makedirs(scene_dir) view_num = int(self.min_views + (diag - self.min_diag)/(self.max_diag - self.min_diag) * (self.max_views - self.min_views)) - view_data = ViewSampleUtil.sample_view_data_world_space(self.target_obj, distance_range=(0.2,0.4), voxel_size=0.005, max_views=view_num) + view_data = ViewSampleUtil.sample_view_data_world_space(self.target_obj, distance_range=(0.2,0.4), voxel_size=0.005, max_views=view_num, min_cam_table_included_degree = self.min_cam_table_included_degree, random_view_ratio = self.random_view_ratio ) object_points = np.array(view_data["voxel_down_sampled_points"]) normals = np.array(view_data["normals"]) points_normals = np.concatenate((object_points, normals), axis=1) diff --git a/pose.py b/pose.py new file mode 100644 index 0000000..ddaed4a --- /dev/null +++ b/pose.py @@ -0,0 +1,151 @@ +import numpy as np + +class PoseUtil: + ROTATION = 1 + TRANSLATION = 2 + SCALE = 3 + + @staticmethod + def get_uniform_translation(trans_m_min, trans_m_max, trans_unit, debug=False): + if isinstance(trans_m_min, list): + x_min, y_min, z_min = trans_m_min + x_max, y_max, z_max = trans_m_max + else: + x_min, y_min, z_min = trans_m_min, trans_m_min, trans_m_min + x_max, y_max, z_max = trans_m_max, trans_m_max, trans_m_max + + x = np.random.uniform(x_min, x_max) + y = np.random.uniform(y_min, y_max) + z = np.random.uniform(z_min, z_max) + translation = np.array([x, y, z]) + if trans_unit == "cm": + translation = translation / 100 + if debug: + print("uniform translation:", translation) + return translation + + @staticmethod + def get_uniform_rotation(rot_degree_min=0, rot_degree_max=180, debug=False): + axis = np.random.randn(3) + axis /= np.linalg.norm(axis) + theta = np.random.uniform( + rot_degree_min / 180 * np.pi, rot_degree_max / 180 * np.pi + ) + + K = np.array( + [[0, -axis[2], axis[1]], [axis[2], 0, -axis[0]], [-axis[1], axis[0], 0]] + ) + R = np.eye(3) + np.sin(theta) * K + (1 - np.cos(theta)) * (K @ K) + if debug: + print("uniform rotation:", theta * 180 / np.pi) + return R + + @staticmethod + def get_uniform_pose( + trans_min, trans_max, rot_min=0, rot_max=180, trans_unit="cm", debug=False + ): + translation = PoseUtil.get_uniform_translation( + trans_min, trans_max, trans_unit, debug + ) + rotation = PoseUtil.get_uniform_rotation(rot_min, rot_max, debug) + pose = np.eye(4) + pose[:3, :3] = rotation + pose[:3, 3] = translation + return pose + + @staticmethod + def get_n_uniform_pose( + trans_min, + trans_max, + rot_min=0, + rot_max=180, + n=1, + trans_unit="cm", + fix=None, + contain_canonical=True, + debug=False, + ): + if fix == PoseUtil.ROTATION: + translations = np.zeros((n, 3)) + for i in range(n): + translations[i] = PoseUtil.get_uniform_translation( + trans_min, trans_max, trans_unit, debug + ) + if contain_canonical: + translations[0] = np.zeros(3) + rotations = PoseUtil.get_uniform_rotation(rot_min, rot_max, debug) + elif fix == PoseUtil.TRANSLATION: + rotations = np.zeros((n, 3, 3)) + for i in range(n): + rotations[i] = PoseUtil.get_uniform_rotation(rot_min, rot_max, debug) + if contain_canonical: + rotations[0] = np.eye(3) + translations = PoseUtil.get_uniform_translation( + trans_min, trans_max, trans_unit, debug + ) + else: + translations = np.zeros((n, 3)) + rotations = np.zeros((n, 3, 3)) + for i in range(n): + translations[i] = PoseUtil.get_uniform_translation( + trans_min, trans_max, trans_unit, debug + ) + for i in range(n): + rotations[i] = PoseUtil.get_uniform_rotation(rot_min, rot_max, debug) + if contain_canonical: + translations[0] = np.zeros(3) + rotations[0] = np.eye(3) + + pose = np.eye(4, 4, k=0)[np.newaxis, :].repeat(n, axis=0) + pose[:, :3, :3] = rotations + pose[:, :3, 3] = translations + + return pose + + @staticmethod + def get_n_uniform_pose_batch( + trans_min, + trans_max, + rot_min=0, + rot_max=180, + n=1, + batch_size=1, + trans_unit="cm", + fix=None, + contain_canonical=False, + debug=False, + ): + + batch_poses = [] + for i in range(batch_size): + pose = PoseUtil.get_n_uniform_pose( + trans_min, + trans_max, + rot_min, + rot_max, + n, + trans_unit, + fix, + contain_canonical, + debug, + ) + batch_poses.append(pose) + pose_batch = np.stack(batch_poses, axis=0) + return pose_batch + + @staticmethod + def get_uniform_scale(scale_min, scale_max, debug=False): + if isinstance(scale_min, list): + x_min, y_min, z_min = scale_min + x_max, y_max, z_max = scale_max + else: + x_min, y_min, z_min = scale_min, scale_min, scale_min + x_max, y_max, z_max = scale_max, scale_max, scale_max + + x = np.random.uniform(x_min, x_max) + y = np.random.uniform(y_min, y_max) + z = np.random.uniform(z_min, z_max) + scale = np.array([x, y, z]) + if debug: + print("uniform scale:", scale) + return scale diff --git a/view_sample_util.py b/view_sample_util.py index 1deaa16..12f0c29 100644 --- a/view_sample_util.py +++ b/view_sample_util.py @@ -3,6 +3,8 @@ import numpy as np import bmesh from collections import defaultdict from scipy.spatial.transform import Rotation as R +from blender.pose import PoseUtil +import random class ViewSampleUtil: @staticmethod @@ -101,7 +103,7 @@ class ViewSampleUtil: return np.array(world_points), np.array(world_normals) @staticmethod - def get_cam_pose(view_data: dict, obj_world_pose: np.ndarray, max_views: int) -> np.ndarray: + def get_cam_pose(view_data: dict, obj_world_pose: np.ndarray, max_views: int, min_cam_table_included_degree: int, random_view_ratio: float) -> np.ndarray: cam_poses = [] min_height_z = 1000 for look_at_point, cam_position in zip(view_data["look_at_points"], view_data["cam_positions"]): @@ -134,8 +136,17 @@ class ViewSampleUtil: filtered_cam_poses = [] for cam_pose in cam_poses: if cam_pose[2, 3] > min_height_z: - filtered_cam_poses.append(cam_pose) - + direction_vector = cam_pose[:3, 2] + horizontal_normal = np.array([0, 0, 1]) + 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: + 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) filtered_cam_poses = [filtered_cam_poses[i] for i in indices] @@ -143,10 +154,10 @@ class ViewSampleUtil: return np.array(filtered_cam_poses) @staticmethod - def sample_view_data_world_space(obj, distance_range:tuple = (0.3,0.5), voxel_size:float = 0.005, max_views: int=1) -> dict: + def sample_view_data_world_space(obj, distance_range:tuple = (0.3,0.5), voxel_size:float = 0.005, max_views: int=1, min_cam_table_included_degree:int=20, random_view_ratio:float = 0.2) -> dict: obj_world_pose = np.asarray(obj.matrix_world) view_data = ViewSampleUtil.sample_view_data(obj, distance_range, voxel_size, max_views) - view_data["cam_poses"] = ViewSampleUtil.get_cam_pose(view_data, obj_world_pose, max_views) + view_data["cam_poses"] = ViewSampleUtil.get_cam_pose(view_data, obj_world_pose, max_views, min_cam_table_included_degree) view_data["voxel_down_sampled_points"], view_data["normals"] = ViewSampleUtil.get_world_points_and_normals(view_data, obj_world_pose) return view_data