83 lines
3.3 KiB
Python
83 lines
3.3 KiB
Python
import numpy as np
|
|
import open3d as o3d
|
|
from task_gen_dependencies.transform_utils import transform_points
|
|
|
|
def compute_sdf_from_obj_surface(mesh, resolution=2): # 2mm
|
|
if mesh is None or len(mesh.vertices) == 0:
|
|
raise ValueError("Invalid mesh provided.")
|
|
bounds_min, bounds_max = mesh.bounds # 网格点上面的这个有点问题。。。
|
|
vertices = o3d.core.Tensor(mesh.vertices, dtype=o3d.core.Dtype.Float32)
|
|
triangles = o3d.core.Tensor(mesh.faces, dtype=o3d.core.Dtype.UInt32)
|
|
scene = o3d.t.geometry.RaycastingScene()
|
|
scene.add_triangles(vertices, triangles)
|
|
shape = np.ceil((bounds_max - bounds_min) / resolution).astype(int)
|
|
|
|
grid = np.mgrid[
|
|
bounds_min[0]:bounds_max[0]:complex(0, shape[0]),
|
|
bounds_min[1]:bounds_max[1]:complex(0, shape[1]),
|
|
bounds_min[2]:bounds_max[2]:complex(0, shape[2])
|
|
]
|
|
grid = grid.reshape(3, -1).T
|
|
|
|
sdf_voxels = scene.compute_signed_distance(grid.astype(np.float32))
|
|
|
|
sdf_voxels = -sdf_voxels.cpu().numpy()
|
|
sdf_voxels = sdf_voxels.reshape(shape)
|
|
|
|
return grid, sdf_voxels
|
|
|
|
|
|
def compute_sdf_from_obj(obj_mes, bounds_max, bounds_min, resolution=2): # 2mm
|
|
# 读取 OBJ 文件并转换为 Trimesh 对象
|
|
|
|
# 将 Trimesh 对象转换为 Open3D TriangleMesh
|
|
vertices = o3d.core.Tensor(obj_mes.vertices, dtype=o3d.core.Dtype.Float32)
|
|
triangles = o3d.core.Tensor(obj_mes.faces, dtype=o3d.core.Dtype.UInt32)
|
|
|
|
# 创建 RaycastingScene
|
|
scene = o3d.t.geometry.RaycastingScene()
|
|
_ = scene.add_triangles(vertices, triangles) # 物体都添加计算
|
|
|
|
# 创建 3D 网格以进行采样
|
|
shape = np.ceil((np.array(bounds_max) - np.array(bounds_min)) / resolution).astype(int)
|
|
steps = (np.array(bounds_max) - np.array(bounds_min)) / shape
|
|
grid = np.mgrid[
|
|
bounds_min[0]:bounds_max[0]:steps[0],
|
|
bounds_min[1]:bounds_max[1]:steps[1],
|
|
bounds_min[2]:bounds_max[2]:steps[2]
|
|
]
|
|
grid = grid.reshape(3, -1).T
|
|
|
|
# 计算 SDF
|
|
sdf_voxels = scene.compute_signed_distance(grid.astype(np.float32))
|
|
|
|
# 转换为 NumPy 数组并调整形状
|
|
sdf_voxels = sdf_voxels.cpu().numpy()
|
|
sdf_voxels = -sdf_voxels # 翻转符号
|
|
sdf_voxels = sdf_voxels.reshape(shape)
|
|
|
|
return grid, sdf_voxels
|
|
|
|
|
|
# cost function
|
|
def get_distance_with_sdf(collision_points, pose, sdf_func):
|
|
transformed_pcs = transform_points(collision_points, pose) # 基于不同pose下point 坐标
|
|
transformed_pcs_flatten = transformed_pcs.reshape(-1, 3) # [num_poses * num_points, 3]
|
|
signed_distance = sdf_func(transformed_pcs_flatten) # [num_poses * num_points]
|
|
return signed_distance
|
|
|
|
|
|
def calculate_collision_cost_sdf(transformed_pcs_active, passive_pose, sdf_func, threshold):
|
|
assert passive_pose.shape == (4, 4)
|
|
|
|
# Convert transformed points to the passive object's coordinate system
|
|
inv_passive_pose = np.linalg.inv(passive_pose)
|
|
transformed_pcs_in_passive = transform_points(transformed_pcs_active[0], inv_passive_pose[None])
|
|
# Flatten the points for SDF function
|
|
transformed_pcs_flatten = transformed_pcs_in_passive.reshape(-1, 3)
|
|
# # Calculate signed distances
|
|
signed_distance = sdf_func(transformed_pcs_flatten)
|
|
# Penalize if any point is closer than the minimum distance
|
|
collision_cost = -np.sum(np.maximum(0, threshold - signed_distance))
|
|
|
|
return collision_cost |