init commit

workflows/simbox/tools/grasp/vis_grasp.py

@@ -0,0 +1,235 @@
+import json
+
+import numpy as np
+import open3d as o3d
+from scipy.spatial.transform import Rotation as R
+
+R1 = np.array([[0.0, 0.0, -1.0], [0.0, 1.0, 0.0], [1.0, 0.0, 0.0]])
+
+
+def create_mesh_box(width, height, depth, dx=0, dy=0, dz=0):
+    """Author: chenxi-wang
+    Create box instance with mesh representation.
+    """
+    box = o3d.geometry.TriangleMesh()
+    vertices = np.array(
+        [
+            [0, 0, 0],
+            [width, 0, 0],
+            [0, 0, depth],
+            [width, 0, depth],
+            [0, height, 0],
+            [width, height, 0],
+            [0, height, depth],
+            [width, height, depth],
+        ]
+    )
+    vertices[:, 0] += dx
+    vertices[:, 1] += dy
+    vertices[:, 2] += dz
+    triangles = np.array(
+        [
+            [4, 7, 5],
+            [4, 6, 7],
+            [0, 2, 4],
+            [2, 6, 4],
+            [0, 1, 2],
+            [1, 3, 2],
+            [1, 5, 7],
+            [1, 7, 3],
+            [2, 3, 7],
+            [2, 7, 6],
+            [0, 4, 1],
+            [1, 4, 5],
+        ]
+    )
+    box.vertices = o3d.utility.Vector3dVector(vertices)
+    box.triangles = o3d.utility.Vector3iVector(triangles)
+    return box
+
+
+def create_franka_gripper_o3d(
+    center,
+    rot_mat,
+    width,
+    depth,
+    depth_base=0.0584,
+    score=1,
+    color=None,
+    return_axis=False,
+):
+    """
+    Author: chenxi-wang
+
+    **Input:**
+
+    - center: numpy array of (3,), target point as gripper center
+
+    - R: numpy array of (3,3), rotation matrix of gripper
+
+    - width: float, gripper width
+
+    - score: float, grasp quality score
+
+    **Output:**
+
+    - open3d.geometry.TriangleMesh
+    """
+    height = 0.004
+    finger_width = 0.004
+    tail_length = 0.04
+    # depth_base = 0.0584
+
+    if color is not None:
+        color_r, color_g, color_b = color
+    else:
+        color_r = score  # red for high score
+        color_g = 0
+        color_b = 1 - score  # blue for low score
+
+    left = create_mesh_box(depth + depth_base + finger_width, finger_width, height)
+    right = create_mesh_box(depth + depth_base + finger_width, finger_width, height)
+    bottom = create_mesh_box(finger_width, width, height)
+    tail = create_mesh_box(tail_length, finger_width, height)
+
+    left_points = np.array(left.vertices)
+    left_triangles = np.array(left.triangles)
+    left_points[:, 0] -= depth_base + finger_width
+    left_points[:, 1] -= width / 2 + finger_width
+    left_points[:, 2] -= height / 2
+
+    right_points = np.array(right.vertices)
+    right_triangles = np.array(right.triangles) + 8
+    right_points[:, 0] -= depth_base + finger_width
+    right_points[:, 1] += width / 2
+    right_points[:, 2] -= height / 2
+
+    bottom_points = np.array(bottom.vertices)
+    bottom_triangles = np.array(bottom.triangles) + 16
+    bottom_points[:, 0] -= finger_width + depth_base
+    bottom_points[:, 1] -= width / 2
+    bottom_points[:, 2] -= height / 2
+
+    tail_points = np.array(tail.vertices)
+    tail_triangles = np.array(tail.triangles) + 24
+    tail_points[:, 0] -= tail_length + finger_width + depth_base
+    tail_points[:, 1] -= finger_width / 2
+    tail_points[:, 2] -= height / 2
+
+    vertices = np.concatenate([left_points, right_points, bottom_points, tail_points], axis=0)
+    # vertices = np.concatenate([left_points], axis=0)
+
+    # align axis for robotiq-2f-85 gripper
+    #
+    # R2 = np.array([[-1.0, 0.0, 0.0], [0.0, -1.0, 0.0], [0.0, 0.0, 1.0]])
+    vertices = np.dot(R1, vertices.T).T
+    # vertices = np.dot(R2, vertices.T).T
+
+    vertices = np.dot(rot_mat, vertices.T).T + center
+    triangles = np.concatenate([left_triangles, right_triangles, bottom_triangles, tail_triangles], axis=0)
+    colors = np.array([[color_r, color_g, color_b] for _ in range(len(vertices))])
+
+    gripper = o3d.geometry.TriangleMesh()
+    gripper.vertices = o3d.utility.Vector3dVector(vertices)
+    gripper.triangles = o3d.utility.Vector3iVector(triangles)
+    gripper.vertex_colors = o3d.utility.Vector3dVector(colors)
+
+    if return_axis:
+        axis = create_rotated_coordinate_frame(size=0.1, rot=rot_mat, t=center)
+        return gripper, axis
+    return gripper
+
+
+def create_rotated_coordinate_frame(size=1.0, rot=None, t=None):
+    """
+    Create a coordinate frame that can be rotated and translated.
+
+    Args:
+        size: frame size
+        rot: rotation matrix (3x3) or rotation angles [rx, ry, rz] in radians
+        t: translation vector [tx, ty, tz]
+    """
+    mesh_frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=size)
+
+    # 如果提供了旋转参数
+    if rot is not None:
+        if isinstance(rot, list) or (isinstance(rot, np.ndarray) and rot.size == 3):
+            rot = R.from_euler("xyz", rot).as_matrix()
+        mesh_frame.rotate(rot, center=(0, 0, 0))
+
+    if t is not None:
+        mesh_frame.translate(t)
+
+    return mesh_frame
+
+
+def get_grasp_candidates(json_path):
+    grasps = []
+    with open(json_path, "r", encoding="utf-8") as file:
+        data = json.load(file)
+    for _, value in data.items():
+        grasps += value
+    grasps = np.array(grasps)
+    return grasps
+
+
+def get_obj_pcd(mesh_path, unit="mm"):
+    mesh = o3d.io.read_triangle_mesh(mesh_path)
+    vertices = np.asarray(mesh.vertices)
+    if unit == "mm":
+        vertices = vertices / 1000.0
+    mesh.vertices = o3d.utility.Vector3dVector(vertices)
+    obj_pcd = mesh.sample_points_uniformly(number_of_points=200000)
+    return obj_pcd
+
+
+def create_coordinate_frame(origin=None, size=0.1):
+    if origin is None:
+        origin = [0, 0, 0]
+    frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=size)
+    frame.translate(origin)
+    return frame
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--obj_path", required=True, help="Path to the object mesh file")
+    parser.add_argument("--N", default=3000, type=int, help="Number of grasps to visualize")
+    parser.add_argument("--unit", choices=["mm", "m"], required=True, help="Unit of the mesh (mm or m)")
+    args = parser.parse_args()
+
+    obj_path = args.obj_path
+    N = args.N
+    scale_unit = args.unit
+
+    # Derive grasp path from obj_path
+    sparse_grasp_path = obj_path.replace("_obj.obj", "_grasp_sparse.npy")
+
+    sparse_grasp = np.load(sparse_grasp_path, allow_pickle=True)
+    grippers = []
+    print("widths :", sparse_grasp[:N, 1])
+    print("sparse_grasp.shape :", sparse_grasp.shape)
+    for grasp in sparse_grasp[:N]:
+        grasp_score = grasp[0]
+        grasp_width = grasp[1]
+        grasp_depth = grasp[3]
+        grasp_rot = grasp[4:13].reshape(3, 3) @ R1.T  # canonical gripper2object
+        grasp_center = grasp[13:16]
+        gripper_mesh = create_franka_gripper_o3d(
+            grasp_center,
+            grasp_rot,
+            grasp_width,
+            grasp_depth,
+            score=grasp_score,
+        )
+        grippers.append(gripper_mesh)
+
+    this_obj_pcd = get_obj_pcd(obj_path, unit=scale_unit)
+    this_obj_pcd.points = o3d.utility.Vector3dVector(np.array(this_obj_pcd.points))
+
+    obj_center = this_obj_pcd.get_center()
+    coordinate_frame = create_coordinate_frame(obj_center, size=0.05)
+
+    o3d.visualization.draw_geometries([this_obj_pcd, *grippers, coordinate_frame])