finish task_gen

task_gen_dependencies/layout_object.py

@@ -0,0 +1,74 @@
+import os
+
+import numpy as np
+import trimesh
+from scipy.interpolate import RegularGridInterpolator
+from scipy.spatial.transform import Rotation as R
+
+from task_gen_dependencies.object import OmniObject
+from task_gen_dependencies.sdf import compute_sdf_from_obj_surface
+from task_gen_dependencies.transform_utils import farthest_point_sampling, random_point, get_bott_up_point
+
+
+def load_and_prepare_mesh(obj_path, up_axis, scale=1.0):
+    if not os.path.exists(obj_path):
+        return None
+    mesh = trimesh.load(obj_path, force="mesh")
+    if 'z' in up_axis:
+        align_rotation = R.from_euler('xyz', [0, 180, 0], degrees=True).as_matrix()
+    elif 'y' in up_axis:
+        align_rotation = R.from_euler('xyz', [-90, 180, 0], degrees=True).as_matrix()
+    elif 'x' in up_axis:
+        align_rotation = R.from_euler('xyz', [0, 0, 90], degrees=True).as_matrix()
+    else:
+        align_rotation = R.from_euler('xyz', [-90, 180, 0], degrees=True).as_matrix()
+    
+
+    align_transform = np.eye(4)
+    align_transform[:3, :3] = align_rotation
+    mesh.apply_transform(align_transform)
+    mesh.apply_scale(scale)
+    return mesh
+
+
+
+def setup_sdf(mesh):
+    _, sdf_voxels = compute_sdf_from_obj_surface(mesh)
+    # create callable sdf function with interpolation
+
+    min_corner = mesh.bounds[0]
+    max_corner = mesh.bounds[1]
+    x = np.linspace(min_corner[0], max_corner[0], sdf_voxels.shape[0])
+    y = np.linspace(min_corner[1], max_corner[1], sdf_voxels.shape[1])
+    z = np.linspace(min_corner[2], max_corner[2], sdf_voxels.shape[2])
+    sdf_func = RegularGridInterpolator((x, y, z), sdf_voxels, bounds_error=False, fill_value=0)
+    return sdf_func
+
+
+class LayoutObject(OmniObject):
+    def __init__(self, obj_info, use_sdf=False, N_collision_points=60, **kwargs):
+        super().__init__(name=obj_info['object_id'], **kwargs)
+
+        obj_dir = obj_info["obj_path"]
+        up_axis = obj_info["upAxis"]
+        if len(up_axis) ==0:
+            up_axis = ['y']
+        mesh_scale = obj_info.get("scale", 0.001)*1000
+        self.mesh = load_and_prepare_mesh(obj_dir, up_axis, mesh_scale)
+        if use_sdf:
+            self.sdf = setup_sdf(self.mesh)
+        
+        if self.mesh is not None:
+            mesh_points, _ =  trimesh.sample.sample_surface(self.mesh, 2000) # 表面采样
+            if mesh_points.shape[0] > N_collision_points:
+                self.collision_points = farthest_point_sampling(mesh_points, N_collision_points) # 碰撞检测点
+
+            self.anchor_points = {}
+            self.anchor_points['top'] = get_bott_up_point(mesh_points, 1.5, descending=False)
+            self.anchor_points['buttom'] = get_bott_up_point(mesh_points, 1.5, descending=True)
+
+            self.anchor_points['top'] = random_point(self.anchor_points['top'], 3)[np.newaxis, :]
+            self.anchor_points['buttom'] = random_point(self.anchor_points['buttom'], 3)[np.newaxis, :]
+
+        self.size = self.mesh.extents.copy()
+        self.up_axis = up_axis[0]