add min_cam_table_included_degree and random_view

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)

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

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