import os import sys import tqdm import json import scipy import trimesh import argparse import numpy as np from qpsolvers import solve_qp from trimesh import transformations as tra # Add meshcat imports for visualization from grasp_gen.utils.meshcat_utils import ( create_visualizer, visualize_pointcloud, visualize_mesh, visualize_grasp, get_color_from_score, ) def make_parser(): parser = argparse.ArgumentParser( formatter_class=argparse.ArgumentDefaultsHelpFormatter ) parser.add_argument( "--object", default="meshes/metal_sheet.obj", help="Object mesh." ) parser.add_argument( "--root-dir", default="", help="Root directory for object mesh path. This will not be stored in the output.json.", ) parser.add_argument( "--gripper", required=False, help="JSON or Yaml file describing the vacuum gripper.", ) parser.add_argument( "--grasps", default=None, help="If specified, will evaluate grasps only (no sampling).", ) parser.add_argument( "--num", default=1000, type=int, help="Number of suction grasps." ) parser.add_argument( "--scale", type=float, default=1.0, required=False, help="Scale of object mesh." ) parser.add_argument( "--num-disturbances", default=10, type=int, help="Number of random disturbance samples.", ) parser.add_argument( "--qp-solver", default="clarabel", choices=("clarabel", "cvxopt", "daqp", "ecos", "osqp", "scs"), help="Name of the QP solver to use.", ) parser.add_argument( "--output", default=None, help="Name of the JSON file to write results to. Compatible with acronym-pipeline format.", ) parser.add_argument("--no-viz", action="store_true", help="No visualization.") parser.add_argument( "--random-seed", type=int, default=None, help="Random seed. Defaults to None.", ) parser.add_argument( "--random-seed-eval", type=int, default=None, help="Random seed for evaluation. Defaults to None.", ) return parser def colorized_points(points, quality): return trimesh.PointCloud(vertices=points, colors=colorize(quality)) def colorize(quality): colors = np.zeros(shape=(len(quality), 4), dtype=np.uint8) colors[:, 3] = 255 # colors[:, 1] = np.array(quality) * 255 # colors[:, 0] = 255 - np.array(quality) * 255 for i, qual in enumerate(quality): colors[i, :3] = color_interpolation(qual) return colors def color_interpolation(input): red = np.array([255, 0, 0]) yellow = np.array([255, 255, 0]) green = np.array([0, 255, 0]) if input < 0.5: return (1.0 - 2.0 * input) * red + 2.0 * input * yellow else: return (2.0 - 2.0 * input) * yellow + (2.0 * input - 1.0) * green def colorize_for_meshcat(quality): """Convert quality scores to colors for meshcat visualization.""" colors = np.zeros(shape=(len(quality), 3), dtype=np.uint8) for i, qual in enumerate(quality): colors[i, :3] = color_interpolation(qual) return colors def skew(x): """Create skew-symmetric matrix from vector for cross-product via matrix multiplication.""" return np.array([[0, -x[2], x[1]], [x[2], 0, -x[0]], [-x[1], x[0], 0]]) def adjoint_transform(frame_a, frame_b): """Used to transform wrenches or twists between different reference frames. E.g.: adj_ab = adjoint_transform(a, b) can be used to transform this wrench: f_a = np.dot(adj_ab.T, f_b) see e.g. p. 62 of MLS: 'A Mathematical Introduction to Robotic Manipulation' """ transform_ab = np.matmul(frame_a, trimesh.transformations.inverse_matrix(frame_b)) rot_ab = transform_ab[:3, :3] tra_ab = transform_ab[:3, 3] adjoint = np.zeros((6, 6)) adjoint[:3, :3] = rot_ab adjoint[3:, 3:] = rot_ab adjoint[:3, 3:] = np.matmul(skew(tra_ab), rot_ab) return adjoint def sunflower_radius(k, n, b): if k > n - b: return 1.0 else: return np.sqrt(k - 0.5) / np.sqrt(n - (b + 1) / 2) def sunflower(n, radius, alpha=2.0, geodesic=False): golden_ratio = (1 + np.sqrt(5)) / 2 points = [] angle_stride = 360 * golden_ratio if geodesic else 2 * np.pi / golden_ratio**2 b = round(alpha * np.sqrt(n)) # number of boundary points for k in range(1, n + 1): r = sunflower_radius(k, n, b) * radius theta = k * angle_stride points.append((r * np.cos(theta), r * np.sin(theta), 0.0)) return np.array(points) class SuctionCupArray(object): def __init__( self, poses, num_sides, radius, height, spring_max_length_change=0.1, suction_force=250.0, material_kappa=0.005, friction_mu=0.5, interpolated_perimeter_vertices=0, standoff_distance=0.001, num_suction_points_for_hole_check=50, collision_mesh_fname=None, ): # pose of apex (?) (height w.r.t. base) # number of base vertices # that form an n-gon with circumradius r and # location of base vertices (opening angle of right pyramid) self.poses = np.array(poses) self.num_suction_cups = len(poses) self.num_sides = num_sides self.radius = radius self.height = height self.standoff_distance = standoff_distance # take scale into account ! self.V = suction_force # newton suction force self.kappa = ( material_kappa # elastic limit / yield strength of the suction cup material ) self.mu = friction_mu self._create_pyramid( self.num_sides, self.radius, self.height, interpolated_perimeter_vertices=interpolated_perimeter_vertices, ) self.num_suction_points_for_hole_check = num_suction_points_for_hole_check self.suction_points = sunflower( n=num_suction_points_for_hole_check, radius=self.radius, alpha=2.0, geodesic=False, ) # 10% according to: "Deformation constraints in a mass-spring model to describe rigid cloth behaviour" by X. Provot et al. # assumption is that all springs are equal - and they are currently in their steady state self.spring_max_length_change = spring_max_length_change perimeter_spring_lengths, flexion_spring_lengths, cone_spring_lengths = ( self.get_spring_lengths(obj_mesh=None, suction_ring_transform=None) ) spring_length = perimeter_spring_lengths[0] self.perimeter_spring_length_max = spring_length * ( 1.0 + spring_max_length_change ) self.perimeter_spring_length_min = spring_length * ( 1.0 - spring_max_length_change ) spring_length = cone_spring_lengths[0] self.cone_spring_length_max = spring_length * (1.0 + spring_max_length_change) self.cone_spring_length_min = spring_length * (1.0 - spring_max_length_change) spring_length = flexion_spring_lengths[0] self.flexion_spring_length_max = spring_length * ( 1.0 + spring_max_length_change ) self.flexion_spring_length_min = spring_length * ( 1.0 - spring_max_length_change ) self.collision_mesh = None if collision_mesh_fname is not None: self.collision_mesh = trimesh.load(collision_mesh_fname, force="mesh") self.collision_mesh_fname = collision_mesh_fname def _create_pyramid( self, sides, radius, height, rotation=0.0, interpolated_perimeter_vertices=0 ): # generate base polygon one_segment = np.pi * 2 / sides vertices = [] total_number_of_vertices = sides + sides * interpolated_perimeter_vertices + 1 mesh_vertex_indices = [] faces = [] perimeter_spring_paths = [] # edges we need to use to calculate path lengths perimeter_springs = [] cone_springs = [] flexion_springs = [] apex_index = total_number_of_vertices - 1 for i in range(sides): point = ( np.sin(one_segment * i + rotation) * radius, np.cos(one_segment * i + rotation) * radius, 0.0, ) vertices.append(point) mesh_vertex_indices.append(len(vertices) - 1) # append interpolated points along segment if interpolated_perimeter_vertices > 0: point2 = ( np.sin(one_segment * (i + 1) + rotation) * radius, np.cos(one_segment * (i + 1) + rotation) * radius, 0.0, ) pts_x = np.linspace( point[0], point2[0], interpolated_perimeter_vertices + 1, endpoint=False, ) pts_y = np.linspace( point[1], point2[1], interpolated_perimeter_vertices + 1, endpoint=False, ) pts_z = np.linspace( point[2], point2[2], interpolated_perimeter_vertices + 1, endpoint=False, ) tmp = [] for j in range(interpolated_perimeter_vertices): tmp.append([len(vertices) - 1, len(vertices)]) vertices.append((pts_x[j + 1], pts_y[j + 1], pts_z[j + 1])) tmp.append( [len(vertices) - 1, len(vertices) % (total_number_of_vertices - 1)] ) perimeter_spring_paths.append(tmp) first_vertex = i * (1 + interpolated_perimeter_vertices) second_vertex = apex_index third_vertex = (first_vertex + 1 + interpolated_perimeter_vertices) % ( total_number_of_vertices - 1 ) faces.append([i, sides, (i + 1) % sides]) perimeter_springs.append([first_vertex, third_vertex]) cone_springs.append([first_vertex, second_vertex]) flexion_springs.append( [ first_vertex, (third_vertex + 1 + interpolated_perimeter_vertices) % (total_number_of_vertices - 1), ] ) # add apex vertices.append((0, 0, height)) vertices = np.array(vertices) mesh_vertex_indices.append(len(vertices) - 1) self.vertices = vertices self.perimeter_spring_paths = np.array(perimeter_spring_paths) self.perimeter_springs = np.array(perimeter_springs) self.interpolated_perimeter_vertices = interpolated_perimeter_vertices self.cone_springs = np.array(cone_springs) self.flexion_springs = np.array(flexion_springs) self.mesh_vertex_indices = np.array(mesh_vertex_indices) self.faces = np.array(faces) # self.mesh = trimesh.Trimesh(vertices=vertices[mesh_vertex_indices], faces=faces) def sample_grasps(self, obj_mesh, num_grasps): # First sample which suction cup from the array to chose suction_cup_indices = np.random.randint( low=0, high=self.num_suction_cups, size=num_grasps ) suction_cup_poses = self.poses[suction_cup_indices] # Sample where to attach the suction cup on the object mesh point_on_surface, face_index = trimesh.sample.sample_surface( obj_mesh, num_grasps ) approach_vector = obj_mesh.face_normals[face_index] # get transformation from point and approach vector grasp_transforms = [] for p, v, T in zip(point_on_surface, approach_vector, suction_cup_poses): contact_transform = trimesh.geometry.plane_transform(p, v) grasp_transform = ( T @ tra.rotation_matrix( angle=np.random.uniform(-np.pi, np.pi), direction=[0, 0, 1] ) @ trimesh.transformations.translation_matrix( [0, 0, -self.standoff_distance] ) @ contact_transform ) # TODO: maybe this can be avoided grasp_transform = trimesh.transformations.inverse_matrix(grasp_transform) grasp_transforms.append(grasp_transform) new_points_on_surface = np.array([ct[:3, 3] for ct in grasp_transforms]) return new_points_on_surface, approach_vector, np.array(grasp_transforms) def get_spring_lengths(self, obj_mesh, suction_ring_transform): # intersect all rays emerging from points on suction cup ring # (shifted towards apex) in approach_vector direction # # - Project n-gon vertices onto object surface (along approach vector) # - Calculate apex distance to base such that average distance == height if obj_mesh is None: transformed_vertices = self.vertices else: if trimesh.ray.has_embree: intersector = trimesh.ray.ray_pyembree.RayMeshIntersector( obj_mesh, scale_to_box=True ) else: intersector = trimesh.ray.ray_triangle.RayMeshIntersector(obj_mesh) transformed_vertices = trimesh.transform_points( self.vertices, suction_ring_transform, translate=True ) ray_origins = transformed_vertices[:-1] # transformed_apex = transformed_vertices[-1] suction_ring, index_ray, _ = intersector.intersects_location( ray_origins, np.tile(-suction_ring_transform[:3, 2], (len(ray_origins), 1)), multiple_hits=False, ) if len(suction_ring) != len(ray_origins): raise RuntimeError("no suction") inv_index_ray = np.empty(index_ray.size, dtype=np.int32) for i in np.arange(index_ray.size): inv_index_ray[index_ray[i]] = i # suction_ring = np.copy(suction_ring[inv_index_ray]) suction_ring = suction_ring[inv_index_ray] # calculate apex apex_param = min( 0, np.sum( (suction_ring - suction_ring_transform[:3, 3]).dot( suction_ring_transform[:3, 2] ) ) / self.num_sides - self.height, ) apex = ( suction_ring_transform[:3, 3] - apex_param * suction_ring_transform[:3, 2] ) transformed_vertices = np.vstack([suction_ring, apex]) if self.interpolated_perimeter_vertices > 0: perimeter_spring_lengths = np.sum( [ [ np.linalg.norm( transformed_vertices[v[0]] - transformed_vertices[v[1]] ) for v in p ] for p in self.perimeter_spring_paths ], axis=1, ) else: perimeter_spring_lengths = np.linalg.norm( transformed_vertices[self.perimeter_springs[:, 0]] - transformed_vertices[self.perimeter_springs[:, 1]], axis=1, ) flexion_spring_lengths = np.linalg.norm( transformed_vertices[self.flexion_springs[:, 0]] - transformed_vertices[self.flexion_springs[:, 1]], axis=1, ) cone_spring_lengths = np.linalg.norm( transformed_vertices[self.cone_springs[:, 0]] - transformed_vertices[self.cone_springs[:, 1]], axis=1, ) return perimeter_spring_lengths, flexion_spring_lengths, cone_spring_lengths def is_sealed(self, obj_mesh, grasp_transform): # From: https://arxiv.org/pdf/1709.06670.pdf # # Calculate energy of suction cup using a spring-mass model # # # (a) The cone faces of the suction cup do not collide with the object # during approach or in the contact configuration. sealed = [True] * self.num_suction_cups suction_cup_collision = False if suction_cup_collision: return [False] * self.num_suction_cups # # Check for object holes. # # (b) The object surface has no holes within the contact ring traced # out by C's perimeter springs. holes_present = False if trimesh.ray.has_embree: intersector = trimesh.ray.ray_pyembree.RayMeshIntersector( obj_mesh, scale_to_box=True ) else: intersector = trimesh.ray.ray_triangle.RayMeshIntersector(obj_mesh) # Check all three suction cup seals for holes for i, suction_cup_T in enumerate(self.poses): # sample points in suction cup bottom suction_ring_transform = grasp_transform @ suction_cup_T suction_points_transformed = trimesh.transform_points( self.suction_points, suction_ring_transform, translate=True ) collisions, _, _ = intersector.intersects_location( suction_points_transformed, np.tile(-suction_ring_transform[:3, 2], (len(self.suction_points), 1)), multiple_hits=False, ) # trimesh.Scene((trimesh.points.PointCloud(suction_points), obj_mesh)).show() holes_present = len(suction_points_transformed) != len(collisions) if holes_present: # s = trimesh.Scene([obj_mesh]) # s.add_geometry(transform=self._last_suction_transform, geometry=dbg_mesh) # s.show() sealed[i] = False if not any(sealed): return sealed # # Calculate Spring Energy # # (c) The energy required in each spring to maintain the contact # configuration is below a real-valued threshold modeling the # maximum deformation of the suction cup material against the # object surface. for i, suction_cup_T in enumerate(self.poses): if sealed[i]: try: ( perimeter_spring_lengths, flexion_spring_lengths, cone_spring_lengths, ) = self.get_spring_lengths( obj_mesh=obj_mesh, suction_ring_transform=grasp_transform @ suction_cup_T, ) if any( perimeter_spring_lengths > self.perimeter_spring_length_max ) or any( perimeter_spring_lengths < self.perimeter_spring_length_min ): # logging.warn("Perimenter spring length violation") sealed[i] = False if any(cone_spring_lengths > self.cone_spring_length_max) or any( cone_spring_lengths < self.cone_spring_length_min ): # logging.warn("Cone spring length violation") sealed[i] = False if any( flexion_spring_lengths > self.flexion_spring_length_max ) or any(flexion_spring_lengths < self.flexion_spring_length_min): # logging.warn("Flexion spring length violation") sealed[i] = False except RuntimeError: sealed[i] = False return sealed def is_wrench_resistant(self, contact_frames, w, qp_solver="clarabel"): """Given a contact_frame and disturbance wrench, return whether this grasp can resist the disurbance.""" object_frame = np.eye(4) contact_wrench_basis = np.eye(6) grasp_map = [] num_fingers = len(contact_frames) for contact_frame in contact_frames: adjoint = adjoint_transform(object_frame, contact_frame) # the grasp map maps each finger's contact forces/torques # into the common object's coordinate frame # It is of dimensioniality 6 x (num_fingers * 6) grasp_map.append(np.matmul(adjoint.T, contact_wrench_basis)) grasp_map = np.hstack(grasp_map) # generate quadratic program of form: # # minimize # (1/2) * x.T * P * x + q.T * x # # subject to # G * x <= h # A * x == b # # Our objective: # ||grasp_map x - w||_2^2 = x^T grasp_map^T grasp_map x - 2 w^T grasp_map x + w^T w # That means: # P = 2 * grasp_map^T grasp_map # q = -2 * w^T grasp_map P = 2.0 * np.matmul(grasp_map.T, grasp_map) q = -2.0 * np.dot(w.T, grasp_map) # Note: f_N = f_z + V # Our constraints: # (friction) # sqrt(3) | f_x | <= mu f_N # sqrt(3) | f_y | <= mu f_N # sqrt(3) | t_z | <= r mu f_N # (material) # sqrt(2) | t_x | <= pi radius kappa # sqrt(2) | t_y | <= pi radius kappa # (suction) # f_z >= -V sqrt3 = np.sqrt(3.0) sqrt2 = np.sqrt(2.0) G = np.array( [ [sqrt3, 0, -self.mu, 0, 0, 0], [-sqrt3, 0, -self.mu, 0, 0, 0], [0, sqrt3, -self.mu, 0, 0, 0], [0, -sqrt3, -self.mu, 0, 0, 0], [0, 0, -self.mu * self.radius, 0, 0, sqrt3], [0, 0, -self.mu * self.radius, 0, 0, -sqrt3], [0, 0, 0, sqrt2, 0, 0], [0, 0, 0, -sqrt2, 0, 0], [0, 0, 0, 0, sqrt2, 0], [0, 0, 0, 0, -sqrt2, 0], [0, 0, -1.0, 0, 0, 0], ] ) G = scipy.linalg.block_diag(*([G] * num_fingers)) h = np.array( [ self.mu * self.V, self.mu * self.V, self.mu * self.V, self.mu * self.V, self.radius * self.mu * self.V, self.radius * self.mu * self.V, np.pi * self.radius * self.kappa, np.pi * self.radius * self.kappa, np.pi * self.radius * self.kappa, np.pi * self.radius * self.kappa, self.V, ] ) h = np.tile(h, num_fingers) x = solve_qp( P=scipy.sparse.csc_matrix(P), q=q, G=scipy.sparse.csc_matrix(G), h=h, solver=qp_solver, ) # other options: ['cvxopt', 'daqp', 'ecos', 'osqp', 'scs'] # ['clarabel', 'daqp', 'ecos', 'osqp', 'scs'] # # Seem to work: osqp, scs, clarabel # # Run-time for 2000 grasps: # clarabel: 31s (without np.ndarray) # clarabel: 30s (with scipy.sparse.csc_matrix) <--- # osqp: 34s # scs: 67s if x is None: return False, None res = np.linalg.norm(grasp_map.dot(x) - w) ** 2 # print("QP solution: res = {}".format(res)) return np.isclose(res, 0.0, atol=1e-3), x def evaluate_grasps( self, obj_mesh, points_on_surface, approach_vectors, grasp_transforms, num_disturbances=10, qp_solver="clarabel", tqdm_disable=True, ): """Check seal formation and wrench resistance of suction grasps.""" # check sealing sealed = [] for tf in tqdm.tqdm(grasp_transforms, disable=tqdm_disable): seal = self.is_sealed(obj_mesh, grasp_transform=tf) sealed.append(seal) sealed = np.array(sealed) # sealing operation; here: if a single cup is sealed it's fine # sealed = sealed.sum(axis=1).clip(0, 1) # here: all suction cups need to be sealed sealed = sealed.all(axis=1) # print(f"Number of sealed grasps: {sum(sealed)} out of {len(sealed)}") # Maybe move this into is_sealed(), see (a) # check collisions collision_mesh_gripper = self.collision_mesh collision_mgr_gripper, collision_mgr_gripper_node_names = ( trimesh.collision.scene_to_collision(collision_mesh_gripper.scene()) ) collision_mgr_gripper_node_names = list(collision_mgr_gripper_node_names.keys()) assert len(collision_mgr_gripper_node_names) <= 1 collision_mgr, _ = trimesh.collision.scene_to_collision(obj_mesh.scene()) grasp_in_collision = [] # create random disturbances disturbances = np.random.normal(size=(num_disturbances, 3)) disturbances = disturbances / np.linalg.norm(disturbances, axis=1).reshape( -1, 1 ) grasp_success = [] for tf, is_sealed in tqdm.tqdm( zip(grasp_transforms, sealed), total=len(grasp_transforms), disable=tqdm_disable, ): # contact_frame = trimesh.geometry.plane_transform(point_on_surf, approach_v) # contact_frame = trimesh.transformations.inverse_matrix(contact_frame) # contact_frame = np.matmul( # contact_frame, trimesh.transformations.euler_matrix(np.pi, 0, 0) # ) if not is_sealed: grasp_success.append(0.0) grasp_in_collision.append(False) else: # check collisions collision_mgr_gripper.set_transform( name=collision_mgr_gripper_node_names[0], transform=tf ) in_collision = collision_mgr.in_collision_other(collision_mgr_gripper) grasp_in_collision.append(in_collision) suction_contacts = [p @ tf for p in self.poses] succ = [] for dist in disturbances: disturbance = np.hstack([dist, [0, 0, 0]]) # TODO transform to suction ring transform suc, _ = self.is_wrench_resistant( suction_contacts, disturbance, qp_solver=qp_solver ) succ.append(suc) grasp_success.append(sum(succ) / len(succ)) return ( points_on_surface, approach_vectors, grasp_transforms, sealed, np.array(grasp_success), np.array(grasp_in_collision), ) def as_dict(self): return { "poses": self.poses.tolist(), "num_sides": self.num_sides, "radius": self.radius, "height": self.height, "suction_force": self.V, "material_kappa": self.kappa, "friction_mu": self.mu, "interpolated_perimeter_vertices": self.interpolated_perimeter_vertices, "spring_max_length_change": self.spring_max_length_change, "standoff_distance": self.standoff_distance, "num_suction_points_for_hole_check": self.num_suction_points_for_hole_check, "collision_mesh_fname": self.collision_mesh_fname, } @classmethod def from_file(cls, fname, key=None): if fname.endswith(".json"): import json data = json.load(open(fname, "r")) elif fname.endswith(".yaml"): import yaml data = yaml.safe_load(open(fname, "r")) if key is not None: data = data[key] return cls( poses=data["poses"], num_sides=data["num_sides"], radius=data["radius"], height=data["height"], suction_force=data["suction_force"], material_kappa=data["material_kappa"], friction_mu=data["friction_mu"], interpolated_perimeter_vertices=data["interpolated_perimeter_vertices"], spring_max_length_change=data["spring_max_length_change"], standoff_distance=data["standoff_distance"], num_suction_points_for_hole_check=data["num_suction_points_for_hole_check"], collision_mesh_fname=data["collision_mesh_fname"], ) # make function so I can call from fastapi service def main_func(args): if args.output is not None: if os.path.isfile(args.output) and os.access(args.output, os.R_OK): print(f"File {args.output} already exists. Skipping.") sys.exit(0) if args.random_seed: np.random.seed(args.random_seed) if args.gripper: suction_gripper = SuctionCupArray.from_file(fname=args.gripper) else: suction_gripper = SuctionCupArray( poses=np.array([np.eye(4)]), num_sides=15, radius=0.025, height=0.06, friction_mu=0.5, material_kappa=0.005, collision_mesh_fname=os.path.join(args.root_dir, args.object), ) if not args.grasps: # load mesh fname_object_mesh = os.path.join(args.root_dir, args.object) obj_mesh = trimesh.load(fname_object_mesh) if args.scale != 1.0: obj_mesh.apply_scale(args.scale) # fix object's CoM obj_center_mass = obj_mesh.center_mass obj_mesh.apply_translation(-obj_center_mass) print( f"Dimensions of {fname_object_mesh}: {obj_mesh.extents}, CoM: {obj_center_mass}" ) # sample grasps import time t0 = time.time() points_on_surface, approach_vectors, grasp_transforms = ( suction_gripper.sample_grasps(obj_mesh=obj_mesh, num_grasps=args.num) ) print(f"Sampling {args.num} grasps took {time.time() - t0}s") else: # Use grasps from input file -- for now assume JSON grasps = json.load(open(args.grasps, "r")) args.object = grasps["object"]["file"] args.scale = grasps["object"]["scale"] fname_object_mesh = os.path.join(args.root_dir, grasps["object"]["file"]) obj_mesh = trimesh.load(fname_object_mesh) if grasps["object"].get("scale", 1.0) != 1.0: obj_mesh.apply_scale(grasps["object"]["scale"]) # fix object's CoM obj_center_mass = obj_mesh.center_mass obj_mesh.apply_translation(-obj_center_mass) grasp_transforms = np.array(grasps["grasps"]["transforms"]) grasp_transforms[:, :3, 3] -= obj_center_mass approach_vectors = np.array([None] * len(grasp_transforms)) points_on_surface = np.copy(grasp_transforms[:, :3, 3]) # evaluate if args.random_seed_eval: np.random.seed(args.random_seed_eval) import time t0 = time.time() points, approach_vectors, contact_transforms, sealed, success, in_collision = ( suction_gripper.evaluate_grasps( obj_mesh=obj_mesh, points_on_surface=points_on_surface, approach_vectors=approach_vectors, grasp_transforms=grasp_transforms, num_disturbances=args.num_disturbances, qp_solver=args.qp_solver, ) ) print(f"Evaluating {args.num} grasps took {time.time() - t0}s") if args.output is not None: if False: # Filter only successful ones -- not used # won't work with in_collision None success_idx = sealed.nonzero() success_idx = np.intersect1d(success_idx, (~in_collision).nonzero()) object_in_gripper = [True] * len(success_idx) else: success_idx = np.arange(len(contact_transforms)) object_in_gripper = np.logical_and(sealed, ~(in_collision.astype(bool))) # Transform data according to mesh.center_mass suction_points = points if points[0] is not None: suction_points += obj_center_mass output_transforms = np.copy(contact_transforms) output_transforms[:, :3, 3] += obj_center_mass # ACRONYM dataset file format output_dict = { "object": { "file": args.object, # this will be without root_dir "scale": args.scale, }, "gripper": suction_gripper.as_dict(), "grasps": { "transforms": output_transforms[success_idx].tolist(), "suction_points": suction_points[success_idx].tolist(), "approach_vectors": approach_vectors[success_idx].tolist(), "object_in_gripper": object_in_gripper[success_idx].tolist(), "is_sealed": sealed[success_idx].tolist(), "score": success[success_idx].tolist(), }, } print( f"Writing result to ({sum(object_in_gripper)}/{len(success_idx)} successes): {args.output}" ) with open(args.output, "w") as f: json.dump(output_dict, f) if not args.no_viz: # Create meshcat visualizer vis = create_visualizer() gripper_name = "single_suction_cup_30mm" # Visualize object mesh visualize_mesh(vis, "object_mesh", obj_mesh, color=[169, 169, 169]) # Visualize sealed configurations print(f"Visualizing sealed configurations ({sum(sealed)}/{len(sealed)})") sealed_colors = colorize_for_meshcat(sealed.astype(float)) visualize_pointcloud(vis, "sealed_points", points, sealed_colors, size=0.005) # Visualize configurations not in collision print( f"Visualizing configurations not in collision ({sum(~in_collision)}/{len(in_collision)})" ) collision_colors = colorize_for_meshcat((~in_collision).astype(float)) visualize_pointcloud( vis, "collision_free_points", points, collision_colors, size=0.005 ) # Visualize quality of configurations print( f"Visualizing quality of configurations (configs greater zero: {sum(success > 0)}/{len(success)}). Max: {np.max(success)}" ) quality_colors = colorize_for_meshcat(success) visualize_pointcloud(vis, "quality_points", points, quality_colors, size=0.005) # Visualize best configurations (top 10) print("Visualizing 'best' configurations (top 10).") combined_criteria = sealed * success # Get indices of top 10 grasps top_indices = np.argsort(combined_criteria)[-10:][ ::-1 ] # Sort descending and take top 10 # Visualize the top 10 grasps for i, idx in enumerate(top_indices): grasp_transform = contact_transforms[idx] grasp_score = combined_criteria[idx] # Use different colors for different ranks (red for best, orange for others) if i == 0: color = [255, 0, 0] # Red for best else: color = [255, 165, 0] # Orange for others # Visualize the grasp visualize_grasp( vis, f"best_grasps/grasp_{i:02d}", grasp_transform, color=color, gripper_name=gripper_name, linewidth=0.6, ) # # Visualize gripper mesh at this position # if suction_gripper.collision_mesh is not None: # visualize_mesh( # vis, # f"best_grippers/gripper_{i:02d}", # suction_gripper.collision_mesh, # color=color, # transform=grasp_transform # ) print(f" Grasp {i+1}: Score = {grasp_score:.3f}") print( f"Visualized {len([idx for idx in top_indices if combined_criteria[idx] > 0])} grasps with positive scores." ) print("Visualization complete. Check meshcat browser window.") input("Press Enter to continue...") if __name__ == "__main__": parser = make_parser() args = parser.parse_args() main_func(args)