optimize code structure

utils/blender_util.py

@@ -0,0 +1,391 @@
+import os
+import json
+import bpy
+import time
+import gc
+import numpy as np
+import mathutils
+
+
+class BlenderUtils:
+
+    TABLE_NAME: str = "table"
+    CAMERA_NAME: str = "Camera"
+    CAMERA_RIGHT_NAME: str = "CameraRight"
+    CAMERA_OBJECT_NAME: str = "CameraObject"
+    DISPLAY_TABLE_NAME: str = "display_table"
+    MESH_FILE_NAME: str = "mesh.obj"
+
+    @staticmethod
+    def get_obj_path(obj_dir, name):
+        return os.path.join(obj_dir, name, BlenderUtils.MESH_FILE_NAME)
+
+    @staticmethod
+    def load_obj(name, mesh_path, scale=1):
+        print(mesh_path)
+        bpy.ops.wm.obj_import(filepath=mesh_path)
+        loaded_object = bpy.context.selected_objects[-1]
+        loaded_object.name = name
+        loaded_object.data.name = name
+        loaded_object.scale = (scale, scale, scale)
+        bpy.ops.rigidbody.object_add()
+        return loaded_object
+
+    @staticmethod
+    def get_obj(name):
+        return bpy.data.objects.get(name)
+
+    @staticmethod
+    def get_obj_pose(name):
+        obj = BlenderUtils.get_obj(name)
+        return np.asarray(obj.matrix_world)
+
+    @staticmethod
+    def add_plane(name, location, orientation, size=10):
+        bpy.ops.mesh.primitive_plane_add(size=size, location=location)
+        plane = bpy.context.selected_objects[-1]
+        plane.name = name
+        plane.rotation_euler = orientation
+        bpy.ops.rigidbody.object_add()
+        bpy.context.object.rigid_body.type = "PASSIVE"
+
+    @staticmethod
+    def add_table(table_model_path):
+        table = BlenderUtils.load_obj(
+            BlenderUtils.TABLE_NAME, table_model_path, scale=0.01
+        )
+        bpy.ops.rigidbody.object_add()
+        bpy.context.object.rigid_body.type = "PASSIVE"
+
+        mat = bpy.data.materials.new(name="TableYellowMaterial")
+        mat.diffuse_color = (1.0, 1.0, 0.0, 1.0)
+        if len(table.data.materials) > 0:
+            table.data.materials[0] = mat
+        else:
+            table.data.materials.append(mat)
+
+    @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.display.shading.show_xray = False
+        bpy.context.scene.display.shading.use_dof = False
+        bpy.context.scene.display.render_aa = "OFF"
+        bpy.context.scene.view_settings.view_transform = "Standard"
+
+        bpy.context.scene.eevee.use_ssr = False  # 关闭屏幕空间反射
+        bpy.context.scene.eevee.use_bloom = False  # 关闭辉光
+        bpy.context.scene.eevee.use_gtao = False  # 关闭环境光遮蔽
+        bpy.context.scene.eevee.use_soft_shadows = False  # 关闭软阴影
+        bpy.context.scene.eevee.use_shadows = False  # 关闭所有阴影
+        bpy.context.scene.world.use_nodes = False  # 如果你不需要环境光，关闭环境节点
+
+        # bpy.context.scene.eevee.use_sss = False            # 关闭次表面散射
+
+        # 2. 设置最低的采样数
+        bpy.context.scene.eevee.taa_render_samples = 1
+        bpy.context.scene.eevee.taa_samples = 1
+        BlenderUtils.init_light_and_camera(
+            init_light_and_camera_config, binocular_vision
+        )
+
+        BlenderUtils.add_plane("plane_floor", location=(0, 0, 0), orientation=(0, 0, 0))
+        BlenderUtils.add_plane("plane_ceil", location=(0, 0, 10), orientation=(0, 0, 0))
+        BlenderUtils.add_plane(
+            "plane_wall_1", location=(5, 0, 5), orientation=(0, np.pi / 2, 0)
+        )
+        BlenderUtils.add_plane(
+            "plane_wall_2", location=(-5, 0, 5), orientation=(0, np.pi / 2, 0)
+        )
+        BlenderUtils.add_plane(
+            "plane_wall_3", location=(0, 5, 5), orientation=(np.pi / 2, 0, 0)
+        )
+        BlenderUtils.add_plane(
+            "plane_wall_4", location=(0, -5, 5), orientation=(np.pi / 2, 0, 0)
+        )
+
+        BlenderUtils.add_table(table_model_path)
+
+    @staticmethod
+    def set_light_params(light, config):
+        light.location = config["location"]
+        light.rotation_euler = config["orientation"]
+        if light.type == "SUN":
+            light.data.energy = config["power"]
+        elif light.type == "POINT":
+            light.data.energy = config["power"]
+
+    @staticmethod
+    def set_camera_params(camera, config, binocular_vision):
+
+        camera_object = bpy.data.objects.new(BlenderUtils.CAMERA_OBJECT_NAME, None)
+        bpy.context.collection.objects.link(camera_object)
+        cameras = [bpy.data.objects.get("Camera")]
+        camera.location = [0, 0, 0]
+        camera.rotation_euler = [0, 0, 0]
+        camera.parent = camera_object
+        if binocular_vision:
+            left_camera = cameras[0]
+            right_camera = left_camera.copy()
+            right_camera.name = BlenderUtils.CAMERA_RIGHT_NAME
+            right_camera.data = left_camera.data.copy()
+            right_camera.data.name = BlenderUtils.CAMERA_RIGHT_NAME
+            bpy.context.collection.objects.link(right_camera)
+            right_camera.parent = camera_object
+            right_camera.location = [config["eye_distance"] / 2, 0, 0]
+            left_camera.location = [-config["eye_distance"] / 2, 0, 0]
+            binocular_angle = config["eye_angle"]
+            half_angle = np.radians(binocular_angle / 2)
+
+            left_camera.rotation_euler[1] = -half_angle
+            right_camera.rotation_euler[1] = half_angle
+            cameras.append(right_camera)
+
+        for camera in cameras:
+            camera.data.clip_start = config["near_plane"]
+            camera.data.clip_end = config["far_plane"]
+
+            bpy.context.scene.render.resolution_x = config["resolution"][0]
+            bpy.context.scene.render.resolution_y = config["resolution"][1]
+            sensor_height = 24.0
+            focal_length = sensor_height / (
+                2 * np.tan(np.radians(config["fov_vertical"]) / 2)
+            )
+            camera.data.lens = focal_length
+            camera.data.sensor_width = (
+                sensor_height * config["resolution"][0] / config["resolution"][1]
+            )
+            camera.data.sensor_height = sensor_height
+
+    @staticmethod
+    def init_light_and_camera(init_light_and_camera_config, binocular_vision):
+
+        camera = BlenderUtils.get_obj(BlenderUtils.CAMERA_NAME)
+        BlenderUtils.set_camera_params(
+            camera,
+            init_light_and_camera_config[BlenderUtils.CAMERA_NAME],
+            binocular_vision,
+        )
+
+    @staticmethod
+    def get_obj_diag(name):
+        obj = BlenderUtils.get_obj(name)
+        return np.linalg.norm(obj.dimensions)
+
+    @staticmethod
+    def matrix_to_blender_pose(matrix):
+        location = matrix[:3, 3]
+        rotation_matrix = matrix[:3, :3]
+        rotation_matrix_blender = mathutils.Matrix(rotation_matrix.tolist())
+        rotation_euler = rotation_matrix_blender.to_euler()
+        return location, rotation_euler
+
+    @staticmethod
+    def set_camera_at(pose):
+        camera = BlenderUtils.get_obj(BlenderUtils.CAMERA_OBJECT_NAME)
+        location, rotation_euler = BlenderUtils.matrix_to_blender_pose(pose)
+
+        camera.location = location
+        camera.rotation_euler = rotation_euler
+
+    @staticmethod
+    def get_object_bottom_z(obj):
+        vertices = [v.co for v in obj.data.vertices]
+        vertices_world = [obj.matrix_world @ v for v in vertices]
+        min_z = min([v.z for v in vertices_world])
+        return min_z
+
+    @staticmethod
+    def render_normal_and_depth(
+        output_dir, file_name, binocular_vision=False, target_object=None
+    ):
+        target_cameras = [BlenderUtils.CAMERA_NAME]
+        if binocular_vision:
+            target_cameras.append(BlenderUtils.CAMERA_RIGHT_NAME)
+            
+        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"
+            scene = bpy.context.scene
+            scene.render.filepath = ""
+
+            mask_dir = os.path.join(output_dir, "normal")
+            if not os.path.exists(mask_dir):
+                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"
+            scene.render.resolution_percentage = 100
+            scene.render.use_overwrite = False
+            scene.render.use_file_extension = False
+            scene.render.use_placeholder = False
+            scene.use_nodes = True
+            tree = scene.node_tree
+
+            for node in tree.nodes:
+                tree.nodes.remove(node)
+
+            rl = tree.nodes.new("CompositorNodeRLayers")
+
+            map_range = tree.nodes.new("CompositorNodeMapRange")
+            map_range.inputs["From Min"].default_value = 0.01
+            map_range.inputs["From Max"].default_value = 5
+            map_range.inputs["To Min"].default_value = 0
+            map_range.inputs["To Max"].default_value = 1
+            tree.links.new(rl.outputs["Depth"], map_range.inputs[0])
+
+            output_depth = tree.nodes.new("CompositorNodeOutputFile")
+
+            depth_dir = os.path.join(output_dir, "depth")
+            if not os.path.exists(depth_dir):
+                os.makedirs(depth_dir)
+            output_depth.base_path = depth_dir
+            output_depth.file_slots[0].path = f"{file_name}_{cam_suffix}.####"
+            output_depth.format.file_format = "PNG"
+            output_depth.format.color_mode = "BW"
+            output_depth.format.color_depth = "16"
+            tree.links.new(map_range.outputs[0], output_depth.inputs[0])
+            bpy.ops.render.render(write_still=True)
+
+        msg = "success"
+        return msg
+        
+    @staticmethod
+    def render_mask(
+        output_dir, file_name, binocular_vision=False, target_object=None
+    ):
+        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"
+            scene = bpy.context.scene
+            scene.render.filepath = ""
+
+            mask_dir = os.path.join(output_dir, "mask")
+            if not os.path.exists(mask_dir):
+                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"
+            scene.render.resolution_percentage = 100
+            scene.render.use_overwrite = False
+            scene.render.use_file_extension = False
+            scene.render.use_placeholder = False
+
+            
+            bpy.ops.render.render(write_still=True)
+
+        msg = "success"
+        return msg
+
+    @staticmethod
+    def save_cam_params(scene_dir, idx, binocular_vision=False):
+        camera = BlenderUtils.get_obj(BlenderUtils.CAMERA_NAME)
+        extrinsic = np.array(camera.matrix_world)
+        cam_data = camera.data
+        focal_length = cam_data.lens
+        sensor_width = cam_data.sensor_width
+        sensor_height = cam_data.sensor_height
+        resolution_x = bpy.context.scene.render.resolution_x
+        resolution_y = bpy.context.scene.render.resolution_y
+        intrinsic = np.zeros((3, 3))
+        intrinsic[0, 0] = focal_length * resolution_x / sensor_width  # fx
+        intrinsic[1, 1] = focal_length * resolution_y / sensor_height  # fy
+        intrinsic[0, 2] = resolution_x / 2.0  # cx
+        intrinsic[1, 2] = resolution_y / 2.0  # cy
+        intrinsic[2, 2] = 1.0
+        cam_object = BlenderUtils.get_obj(BlenderUtils.CAMERA_OBJECT_NAME)
+        extrinsic_cam_object = np.array(cam_object.matrix_world)
+        data = {
+            "extrinsic": extrinsic.tolist(),
+            "extrinsic_cam_object": extrinsic_cam_object.tolist(),
+            "intrinsic": intrinsic.tolist(),
+            "far_plane": camera.data.clip_end,
+            "near_plane": camera.data.clip_start,
+        }
+        if binocular_vision:
+            right_camera = BlenderUtils.get_obj(BlenderUtils.CAMERA_RIGHT_NAME)
+            extrinsic_right = np.array(right_camera.matrix_world)
+            print("result:", extrinsic_right)
+
+            data["extrinsic_R"] = extrinsic_right.tolist()
+
+        cam_params_dir = os.path.join(scene_dir, "camera_params")
+        if not os.path.exists(cam_params_dir):
+            os.makedirs(cam_params_dir)
+        cam_params_path = os.path.join(cam_params_dir, f"{idx}.json")
+        with open(cam_params_path, "w") as f:
+            json.dump(data, f, indent=4)
+
+    @staticmethod
+    def reset_objects_and_platform():
+        all_objects = bpy.data.objects
+        keep_objects = {
+            "plane_floor",
+            "plane_ceil",
+            "plane_wall_1",
+            "plane_wall_2",
+            "plane_wall_3",
+            "plane_wall_4",
+        }
+        keep_objects.add(BlenderUtils.CAMERA_OBJECT_NAME)
+        keep_objects.add(BlenderUtils.CAMERA_NAME)
+        keep_objects.add(BlenderUtils.CAMERA_RIGHT_NAME)
+        keep_objects.add(BlenderUtils.TABLE_NAME)
+
+        for obj in all_objects:
+            if obj.name not in keep_objects:
+                bpy.data.objects.remove(obj, do_unlink=True)
+
+        for block in bpy.data.meshes:
+            if block.users == 0:
+                bpy.data.meshes.remove(block)
+        for block in bpy.data.materials:
+            if block.users == 0:
+                bpy.data.materials.remove(block)
+        for block in bpy.data.images:
+            if block.users == 0:
+                bpy.data.images.remove(block)
+
+        gc.collect()
+        bpy.context.scene.frame_set(0)
+
+    @staticmethod
+    def save_scene_info(scene_root_dir, display_table_config, target_name):
+        all_objects = bpy.data.objects
+        no_save_objects = {
+            "plane_floor",
+            "plane_ceil",
+            "plane_wall_1",
+            "plane_wall_2",
+            "plane_wall_3",
+            "plane_wall_4",
+        }
+        no_save_objects.add(BlenderUtils.CAMERA_OBJECT_NAME)
+        no_save_objects.add(BlenderUtils.CAMERA_NAME)
+        no_save_objects.add(BlenderUtils.CAMERA_RIGHT_NAME)
+        no_save_objects.add(BlenderUtils.TABLE_NAME)
+        scene_info = {}
+        for obj in all_objects:
+            if (
+                obj.name not in no_save_objects
+                and obj.name != BlenderUtils.DISPLAY_TABLE_NAME
+            ):
+                obj_info = {
+                    "location": list(obj.location),
+                    "rotation_euler": list(obj.rotation_euler),
+                    "scale": list(obj.scale),
+                }
+                scene_info[obj.name] = obj_info
+        scene_info[BlenderUtils.DISPLAY_TABLE_NAME] = display_table_config
+        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)

utils/material_util.py

@@ -0,0 +1,96 @@
+import bpy
+
+class MaterialUtil:
+
+    ''' --------- Basic --------- '''
+    @staticmethod
+    def change_object_material(obj, mat):
+        if obj.data.materials:
+            obj.data.materials[0] = mat
+        else:
+            obj.data.materials.append(mat)
+    
+    ''' ------- Materials ------- '''
+    @staticmethod
+    def create_normal_material():
+        normal_mat = bpy.data.materials.new(name="NormalMaterial")
+        normal_mat.use_nodes = True
+
+        nodes = normal_mat.node_tree.nodes
+        links = normal_mat.node_tree.links
+
+        nodes.clear()
+
+        geometry_node = nodes.new(type="ShaderNodeNewGeometry")
+        vector_transform_node = nodes.new(type="ShaderNodeVectorTransform")
+        separate_xyz_node = nodes.new(type="ShaderNodeSeparateXYZ")
+        multiply_node_x = nodes.new(type="ShaderNodeMath")
+        multiply_node_y = nodes.new(type="ShaderNodeMath")
+        multiply_node_z = nodes.new(type="ShaderNodeMath")
+        combine_xyz_node = nodes.new(type="ShaderNodeCombineXYZ")
+        light_path_node = nodes.new(type="ShaderNodeLightPath")
+        emission_node_1 = nodes.new(type="ShaderNodeEmission")
+        emission_node_2 = nodes.new(type="ShaderNodeEmission")
+        mix_shader_node_1 = nodes.new(type="ShaderNodeMixShader")
+        mix_shader_node_2 = nodes.new(type="ShaderNodeMixShader")
+        material_output_node = nodes.new(type="ShaderNodeOutputMaterial")
+
+        vector_transform_node.vector_type = 'VECTOR'
+        vector_transform_node.convert_from = 'WORLD'
+        vector_transform_node.convert_to = 'CAMERA'
+
+        multiply_node_x.operation = 'MULTIPLY'
+        multiply_node_x.inputs[1].default_value = 1.0
+
+        multiply_node_y.operation = 'MULTIPLY'
+        multiply_node_y.inputs[1].default_value = 1.0
+
+        multiply_node_z.operation = 'MULTIPLY'
+        multiply_node_z.inputs[1].default_value = -1.0
+
+        emission_node_1.inputs['Strength'].default_value = 1.0
+        emission_node_2.inputs['Strength'].default_value = 1.0
+
+        mix_shader_node_2.inputs['Fac'].default_value = 0.5
+
+        links.new(geometry_node.outputs['Normal'], vector_transform_node.inputs['Vector'])
+        links.new(vector_transform_node.outputs['Vector'], separate_xyz_node.inputs['Vector'])
+        links.new(separate_xyz_node.outputs['X'], multiply_node_x.inputs[0])
+        links.new(separate_xyz_node.outputs['Y'], multiply_node_y.inputs[0])
+        links.new(separate_xyz_node.outputs['Z'], multiply_node_z.inputs[0])
+        links.new(multiply_node_x.outputs['Value'], combine_xyz_node.inputs['X'])
+        links.new(multiply_node_y.outputs['Value'], combine_xyz_node.inputs['Y'])
+        links.new(multiply_node_z.outputs['Value'], combine_xyz_node.inputs['Z'])
+        links.new(combine_xyz_node.outputs['Vector'], emission_node_1.inputs['Color'])
+        links.new(light_path_node.outputs['Is Camera Ray'], mix_shader_node_1.inputs['Fac'])
+        links.new(emission_node_1.outputs['Emission'], mix_shader_node_1.inputs[2])
+        links.new(mix_shader_node_1.outputs['Shader'], mix_shader_node_2.inputs[1])
+        links.new(emission_node_2.outputs['Emission'], mix_shader_node_2.inputs[2])
+        links.new(mix_shader_node_2.outputs['Shader'], material_output_node.inputs['Surface'])
+        return normal_mat
+    
+    @staticmethod
+    def create_mask_material(color=(1.0, 1.0, 1.0)): 
+        mask_mat = bpy.data.materials.new(name="MaskMaterial")
+        mask_mat.use_nodes = True
+
+        nodes = mask_mat.node_tree.nodes
+        links = mask_mat.node_tree.links
+
+        nodes.clear()
+        emission_node = nodes.new(type="ShaderNodeEmission")
+        emission_node.inputs['Color'].default_value = (*color, 1.0) 
+        emission_node.inputs['Strength'].default_value = 1.0
+        material_output_node = nodes.new(type="ShaderNodeOutputMaterial")
+        links.new(emission_node.outputs['Emission'], material_output_node.inputs['Surface'])
+
+        return mask_mat
+
+    
+
+# -------- debug --------
+if __name__ == "__main__":
+    cube = bpy.data.objects.get("Cube")
+    normal_mat = MaterialUtil.create_normal_material()
+    MaterialUtil.change_object_material(cube, normal_mat)
+    

utils/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

utils/view_sample_util.py

@@ -0,0 +1,168 @@
+
+import numpy as np
+import bmesh
+from collections import defaultdict
+from scipy.spatial.transform import Rotation as R
+from utils.pose import PoseUtil
+import random
+
+class ViewSampleUtil:
+    @staticmethod
+    def farthest_point_sampling(points, num_samples):
+        num_points = points.shape[0]
+        if num_samples >= num_points:
+            return points, np.arange(num_points)
+        sampled_indices = np.zeros(num_samples, dtype=int)
+        sampled_indices[0] = np.random.randint(num_points)
+        min_distances = np.full(num_points, np.inf)
+        for i in range(1, num_samples):
+            current_point = points[sampled_indices[i - 1]]
+            dist_to_current_point = np.linalg.norm(points - current_point, axis=1)
+            min_distances = np.minimum(min_distances, dist_to_current_point)
+            sampled_indices[i] = np.argmax(min_distances)
+        downsampled_points = points[sampled_indices]
+        return downsampled_points, sampled_indices
+    
+    @staticmethod
+    def voxel_downsample(points, voxel_size):
+        voxel_grid = defaultdict(list)
+        for i, point in enumerate(points):
+            voxel_index = tuple((point // voxel_size).astype(int))
+            voxel_grid[voxel_index].append(i)
+
+        downsampled_points = []
+        downsampled_indices = []
+        for indices in voxel_grid.values():
+            selected_index = indices[0]
+            downsampled_points.append(points[selected_index])
+            downsampled_indices.append(selected_index)
+
+        return np.array(downsampled_points), downsampled_indices
+
+    @staticmethod
+    def sample_view_data(obj, distance_range:tuple = (0.25,0.5), voxel_size:float = 0.005, max_views: int = 1, pertube_repeat:int = 1) -> dict:
+        view_data = {
+            "look_at_points": [],
+            "cam_positions": [],
+        }
+        mesh = obj.data
+        bm = bmesh.new()
+        bm.from_mesh(mesh)
+        bm.verts.ensure_lookup_table()
+        bm.faces.ensure_lookup_table()
+        bm.normal_update()
+        
+        look_at_points = []
+        cam_positions = []
+        normals = []
+        for v in bm.verts:
+            look_at_point = np.array(v.co)
+            
+            view_data["look_at_points"].append(look_at_point)
+            normal = np.zeros(3)
+            for loop in v.link_loops:
+                normal += np.array(loop.calc_normal())
+            normal /= len(v.link_loops)
+            normal = normal / np.linalg.norm(normal)
+            if np.isnan(normal).any():
+                continue
+            if np.dot(normal, look_at_point) < 0:
+                normal = -normal
+            normals.append(normal)
+            
+            for _ in range(pertube_repeat):
+                perturb_angle = np.radians(np.random.uniform(0, 30))
+                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()
+                perturbed_normal = np.dot(rotation_matrix, normal)
+                middle_distance = (distance_range[0] + distance_range[1]) / 2
+                perturbed_distance = random.uniform(middle_distance-0.05, middle_distance+0.05)
+                cam_position = look_at_point + perturbed_distance * perturbed_normal
+                look_at_points.append(look_at_point)
+                cam_positions.append(cam_position)
+            
+            
+        bm.free() 
+        look_at_points = np.array(look_at_points)
+        cam_positions = np.array(cam_positions)
+        voxel_downsampled_look_at_points, selected_indices = ViewSampleUtil.voxel_downsample(look_at_points, voxel_size)
+        voxel_downsampled_cam_positions = cam_positions[selected_indices]
+        voxel_downsampled_normals = np.array(normals)[selected_indices]
+
+        fps_downsampled_look_at_points, selected_indices = ViewSampleUtil.farthest_point_sampling(voxel_downsampled_look_at_points, max_views*2)
+        fps_downsampled_cam_positions = voxel_downsampled_cam_positions[selected_indices]
+
+        view_data["look_at_points"] = fps_downsampled_look_at_points.tolist()
+        view_data["cam_positions"] = fps_downsampled_cam_positions.tolist()
+        view_data["normals"] = voxel_downsampled_normals
+        view_data["voxel_down_sampled_points"] = voxel_downsampled_look_at_points 
+        return view_data
+    
+    @staticmethod
+    def get_world_points_and_normals(view_data: dict, obj_world_pose: np.ndarray) -> tuple:
+        world_points = []
+        world_normals = []
+        for voxel_down_sampled_points, normal in zip(view_data["voxel_down_sampled_points"], view_data["normals"]):
+            voxel_down_sampled_points_world = obj_world_pose @ np.append(voxel_down_sampled_points, 1.0)
+            normal_world = obj_world_pose[:3, :3] @ normal
+            world_points.append(voxel_down_sampled_points_world[:3])
+            world_normals.append(normal_world)
+        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, 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"]):
+            look_at_point_world = obj_world_pose @ np.append(look_at_point, 1.0)
+            cam_position_world = obj_world_pose @ np.append(cam_position, 1.0)
+            if look_at_point_world[2] < min_height_z:
+                min_height_z = look_at_point_world[2]
+            look_at_point_world = look_at_point_world[:3]  
+            cam_position_world = cam_position_world[:3]    
+            
+            forward_vector = cam_position_world - look_at_point_world
+            forward_vector /= np.linalg.norm(forward_vector)
+            
+            up_vector = np.array([0, 0, 1])
+
+            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
+
+            cam_pose = np.eye(4)
+            cam_pose[:3, :3] = rotation_matrix
+            cam_pose[:3, 3] = cam_position_world
+            cam_poses.append(cam_pose)
+
+        filtered_cam_poses = []
+        for cam_pose in cam_poses:
+            if cam_pose[2, 3] > min_height_z:
+                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]
+            
+        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, 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, min_cam_table_included_degree, random_view_ratio)
+        view_data["voxel_down_sampled_points"], view_data["normals"] =  ViewSampleUtil.get_world_points_and_normals(view_data, obj_world_pose)
+        return view_data
+