Minor

src/active_grasp/controller.py

@@ -6,6 +6,7 @@ import numpy as np
 import rospy
 from sensor_msgs.msg import Image
 import trimesh
+import time
 
 from .bbox import from_bbox_msg
 from .timer import Timer
@@ -133,11 +134,11 @@ class GraspController:
         return np.r_[linear, angular]
 
     def execute_grasp(self, grasp):
-        self.create_collision_scene(grasp)
+        self.create_collision_scene()
         T_base_grasp = self.postprocess(grasp.pose)
         self.gripper.move(0.08)
         self.moveit.goto(T_base_grasp * Transform.t([0, 0, -0.06]) * self.T_grasp_ee)
-        self.moveit.scene.remove_world_object(self.target_co_name)
+        self.remove_target_collision_object(grasp)
         rospy.sleep(0.5)  # Wait for the planning scene to be updated
         self.moveit.gotoL(T_base_grasp * self.T_grasp_ee)
         self.gripper.grasp()
@@ -146,7 +147,7 @@ class GraspController:
         success = self.gripper.read() > 0.002
         return "succeeded" if success else "failed"
 
-    def create_collision_scene(self, grasp):
+    def create_collision_scene(self):
         # Segment support surface
         cloud = self.policy.tsdf.get_scene_cloud()
         cloud = cloud.transform(self.policy.T_base_task.as_matrix())
@@ -154,34 +155,39 @@ class GraspController:
         support_cloud = cloud.select_by_index(inliers)
         cloud = cloud.select_by_index(inliers, invert=True)
 
-        # Cluster cloud, generate convex hulls and add collision objects to the scene
+        # Add collision object for the support
+        self.add_collision_mesh("support", compute_convex_hull(support_cloud))
+
+        # Cluster cloud
         labels = np.array(cloud.cluster_dbscan(eps=0.02, min_points=10))
-        self.target_co_name = None
+
-        min_dist = np.inf
+        # Generate convex collision objects for each segment
-        tip = grasp.pose.apply([0.0, 0.0, 0.05])
+        self.hulls = []
         for label in range(labels.max() + 1):
             segment = cloud.select_by_index(np.flatnonzero(labels == label))
             try:
                 hull = compute_convex_hull(segment)
                 name = f"object_{label}"
                 self.add_collision_mesh(name, hull)
-                dist = trimesh.proximity.closest_point_naive(hull, np.array([tip]))[1]
+                self.hulls.append(hull)
-                if dist < min_dist:
-                    self.target_co_name, min_dist = name, dist
             except:
                 # Qhull fails in some edge cases
                 pass
 
-        # Add collision object for the support
-        self.add_collision_mesh("support", compute_convex_hull(support_cloud))
-
-        rospy.sleep(1.0)
-
     def add_collision_mesh(self, name, mesh):
         frame, pose = self.base_frame, Transform.identity()
         co = create_collision_object_from_mesh(name, frame, pose, mesh)
         self.moveit.scene.add_object(co)
 
+    def remove_target_collision_object(self, grasp):
+        label, min_dist = None, np.inf
+        for i, hull in enumerate(self.hulls):
+            tip = grasp.pose.apply([0.0, 0.0, 0.05])
+            dist = trimesh.proximity.closest_point_naive(hull, np.array([tip]))[1]
+            if dist < min_dist:
+                label, min_dist = i, dist
+        self.moveit.scene.remove_world_object(f"object_{label}")
+
     def postprocess(self, T_base_grasp):
         rot = T_base_grasp.rotation
         if rot.as_matrix()[:, 0][0] < 0:  # Ensure that the camera is pointing forward

src/active_grasp/policy.py

@@ -78,13 +78,7 @@ class Policy:
     def update(self, img, x, q):
         raise NotImplementedError
 
-    def sort_grasps(self, grasps, qualities, q):
+    def select_best_grasp(self, grasps, qualities, q):
-        """
-        1. Transform grasp configurations into base_frame
-        2. Check whether the finger tips lie within the bounding box
-        3. Remove grasps for which no IK solution was found
-        4. Sort grasps according to score_fn
-        """
         filtered_grasps, scores = [], []
         for grasp, quality in zip(grasps, qualities):
             pose = self.T_base_task * grasp.pose
@@ -96,8 +90,7 @@ class Policy:
                 if q_grasp is not None:
                     filtered_grasps.append(grasp)
                     scores.append(self.score_fn(grasp, quality, q, q_grasp))
-        filtered_grasps, scores = np.asarray(filtered_grasps), np.asarray(scores)
+        i = np.argmax(scores)
-        i = np.argsort(-scores)
         return filtered_grasps[i], qualities[i], scores[i]
 
     def score_fn(self, grasp, quality, q, q_grasp):
@@ -116,14 +109,11 @@ class SingleViewPolicy(Policy):
 
             out = self.vgn.predict(tsdf_grid)
             self.vis.quality(self.task_frame, voxel_size, out.qual, 0.5)
-
             grasps, qualities = select_local_maxima(voxel_size, out, self.qual_thresh)
-            grasps, qualities, _ = self.sort_grasps(grasps, qualities, q)
 
             if len(grasps) > 0:
-                self.best_grasp = grasps[0]
+                self.best_grasp, qual, _ = self.select_best_grasp(grasps, qualities, q)
-                # self.vis.grasps(self.base_frame, grasps, qualities)
+                self.vis.grasp(self.base_frame, self.best_grasp, qual)
-                self.vis.grasp(self.base_frame, self.best_grasp, qualities[0])
 
             self.done = True
 
@@ -156,12 +146,10 @@ class MultiViewPolicy(Policy):
 
         with Timer("grasp_selection"):
             grasps, qualities = select_local_maxima(voxel_size, out, self.qual_thresh)
-            grasps, qualities, _ = self.sort_grasps(grasps, qualities, q)
 
         if len(grasps) > 0:
-            self.best_grasp = grasps[0]
+            self.best_grasp, quality, _ = self.select_best_grasp(grasps, qualities, q)
-            # self.vis.grasps(self.base_frame, grasps, qualities)
+            self.vis.grasp(self.base_frame, self.best_grasp, quality)
-            self.vis.grasp(self.base_frame, self.best_grasp, qualities[0])
         else:
             self.best_grasp = None
             self.vis.clear_grasp()