Visualize quality and grasps

active_grasp/policy.py

@@ -34,7 +34,7 @@ class BasePolicy(Policy):
         msg = rospy.wait_for_message(info_topic, CameraInfo, rospy.Duration(2.0))
         self.intrinsic = from_camera_info_msg(msg)
         self.vgn = VGN(Path(rospy.get_param("vgn/model")))
-        self.score_fn = lambda g: g.pose.translation[2]
+        self.score_fn = lambda g: g.pose.translation[2]  # TODO
 
     def init_visualizer(self):
         self.visualizer = Visualizer(self.base_frame)
@@ -42,13 +42,12 @@ class BasePolicy(Policy):
     def activate(self, bbox):
         self.bbox = bbox
 
-        # Define the VGN task frame s.t. the bounding box is in its center
         self.center = 0.5 * (bbox.min + bbox.max)
         self.T_base_task = Transform.translation(self.center - np.full(3, 0.15))
         tf.broadcast(self.T_base_task, self.base_frame, self.task_frame)
-        rospy.sleep(0.1)  # wait for the transform to be published
+        rospy.sleep(1.0)  # wait for the transform to be published
 
-        N, self.T = 40, 10  # spatial and temporal resolution
+        N, self.T = 40, 10
         grid_shape = (N,) * 3
 
         self.tsdf = UniformTSDFVolume(0.3, N)
@@ -67,64 +66,62 @@ class BasePolicy(Policy):
     def integrate_img(self, img, extrinsic):
         self.viewpoints.append(extrinsic.inv())
         self.tsdf.integrate(img, self.intrinsic, extrinsic * self.T_base_task)
-        self.visualizer.scene_cloud(self.task_frame, self.tsdf.get_scene_cloud())
-        self.visualizer.path(self.viewpoints)
 
         if self.filter_grasps:
-            tsdf_grid = self.tsdf.get_grid()
+            out = self.vgn.predict(self.tsdf.get_grid())
-            out = self.vgn.predict(tsdf_grid)
             t = (len(self.viewpoints) - 1) % self.T
             self.qual_hist[t, ...] = out.qual
             self.rot_hist[t, ...] = out.rot
             self.width_hist[t, ...] = out.width
 
+            mean_qual = self.compute_mean_quality()
+            self.visualizer.quality(self.task_frame, self.tsdf.voxel_size, mean_qual)
+
+        self.visualizer.scene_cloud(self.task_frame, self.tsdf.get_scene_cloud())
+        self.visualizer.path(self.viewpoints)
+
     def compute_best_grasp(self):
         if self.filter_grasps:
-            T = len(self.viewpoints) if len(self.viewpoints) // self.T == 0 else self.T
+            qual = self.compute_mean_quality()
-            mask = self.qual_hist[:T, ...] > 0.0
+            index_list = select_local_maxima(qual, 0.9, 3)
-            # The next line prints a warning since some voxels have no grasp
+            grasps = [g for i in index_list if (g := self.select_mean_at(i))]
-            # predictions resulting in empty slices.
-            qual = np.mean(self.qual_hist[:T, ...], axis=0, where=mask)
-            qual = np.nan_to_num(qual, copy=False)
-            qual = threshold_quality(qual, 0.9)
-            index_list = select_local_maxima(qual, 3)
-
-            grasps = []
-            for (i, j, k) in index_list:
-                ts = np.flatnonzero(self.qual_hist[:T, i, j, k])
-                if len(ts) < 3:
-                    continue
-                oris = Rotation.from_quat([self.rot_hist[t, :, i, j, k] for t in ts])
-                ori = oris.mean()
-                # TODO check variance as well
-                pos = np.array([i, j, k], dtype=np.float64)
-                width = self.width_hist[ts, i, j, k].mean()
-                quality = self.qual_hist[ts, i, j, k].mean()
-                grasps.append(Grasp(Transform(ori, pos), width, quality))
         else:
-            tsdf_grid = self.tsdf.get_grid()
+            out = self.vgn.predict(self.tsdf.get_grid())
-            out = self.vgn.predict(tsdf_grid)
+            qual = out.qual
-            qual = threshold_quality(out.qual, 0.9)
+            index_list = select_local_maxima(qual, 0.9, 3)
-            index_list = select_local_maxima(qual, 3)
             grasps = [select_at(out, i) for i in index_list]
 
         grasps = [from_voxel_coordinates(g, self.tsdf.voxel_size) for g in grasps]
-        grasps = self.transform_grasps_to_base_frame(grasps)
+        grasps = self.transform_and_reject(grasps)
-        grasps = self.select_grasps_on_target_object(grasps)
         grasps = sort_grasps(grasps, self.score_fn)
 
+        self.visualizer.quality(self.task_frame, self.tsdf.voxel_size, qual)
+        self.visualizer.grasps(grasps)
+
         return grasps[0] if len(grasps) > 0 else None
 
-    def transform_grasps_to_base_frame(self, grasps):
+    def compute_mean_quality(self):
-        for grasp in grasps:
+        qual = np.mean(self.qual_hist, axis=0, where=self.qual_hist > 0.0)
-            grasp.pose = self.T_base_task * grasp.pose
+        return np.nan_to_num(qual, copy=False)  # mean of empty slices returns nan
-        return grasps
 
-    def select_grasps_on_target_object(self, grasps):
+    def select_mean_at(self, index):
+        i, j, k = index
+        ts = np.flatnonzero(self.qual_hist[:, i, j, k])
+        if len(ts) < 3:
+            return
+        ori = Rotation.from_quat([self.rot_hist[t, :, i, j, k] for t in ts])
+        pos = np.array([i, j, k], dtype=np.float64)
+        width = self.width_hist[ts, i, j, k].mean()
+        qual = self.qual_hist[ts, i, j, k].mean()
+        return Grasp(Transform(ori.mean(), pos), width, qual)
+
+    def transform_and_reject(self, grasps):
         result = []
         for grasp in grasps:
-            tip = grasp.pose.rotation.apply([0, 0, 0.05]) + grasp.pose.translation
+            pose = self.T_base_task * grasp.pose
+            tip = pose.rotation.apply([0, 0, 0.05]) + pose.translation
             if self.bbox.is_inside(tip):
+                grasp.pose = pose
                 result.append(grasp)
         return result
 

active_grasp/visualization.py

@@ -1,3 +1,4 @@
+from geometry_msgs.msg import PoseArray
 import numpy as np
 import rospy
 
@@ -11,10 +12,17 @@ class Visualizer:
         self.frame = frame
         self.marker_pub = rospy.Publisher(topic, MarkerArray, queue_size=1)
         self.scene_cloud_pub = rospy.Publisher("scene_cloud", PointCloud2, queue_size=1)
+        self.quality_pub = rospy.Publisher("quality", PointCloud2, queue_size=1)
+        self.grasps_pub = rospy.Publisher("grasps", PoseArray, queue_size=1)
 
     def clear(self):
-        marker = Marker(action=Marker.DELETEALL)
+        self.draw([Marker(action=Marker.DELETEALL)])
-        self.draw([marker])
+        msg = to_cloud_msg(self.frame, np.array([]))
+        self.scene_cloud_pub.publish(msg)
+        self.quality_pub.publish(msg)
+        msg = PoseArray()
+        msg.header.frame_id = self.frame
+        self.grasps_pub.publish(msg)
 
     def bbox(self, bbox):
         pose = Transform.translation((bbox.min + bbox.max) / 2.0)
@@ -23,10 +31,6 @@ class Visualizer:
         marker = create_cube_marker(self.frame, pose, scale, color, ns="bbox")
         self.draw([marker])
 
-    def scene_cloud(self, frame, cloud):
-        msg = to_cloud_msg(frame, np.asarray(cloud.points))
-        self.scene_cloud_pub.publish(msg)
-
     def path(self, poses):
         color = np.r_[31, 119, 180] / 255.0
         points = [p.translation for p in poses]
@@ -52,3 +56,18 @@ class Visualizer:
 
     def draw(self, markers):
         self.marker_pub.publish(MarkerArray(markers=markers))
+
+    def scene_cloud(self, frame, cloud):
+        msg = to_cloud_msg(frame, np.asarray(cloud.points))
+        self.scene_cloud_pub.publish(msg)
+
+    def quality(self, frame, voxel_size, quality):
+        points, values = grid_to_map_cloud(voxel_size, quality, threshold=0.8)
+        msg = to_cloud_msg(frame, points, intensities=values)
+        self.quality_pub.publish(msg)
+
+    def grasps(self, grasps):
+        msg = PoseArray()
+        msg.header.frame_id = self.frame
+        msg.poses = [to_pose_msg(grasp.pose) for grasp in grasps]
+        self.grasps_pub.publish(msg)

launch/active_grasp.rviz

@@ -3,9 +3,12 @@ Panels:
     Help Height: 78
     Name: Displays
     Property Tree Widget:
-      Expanded: ~
+      Expanded:
+        - /Grasps1/Status1
+        - /Markers1
+        - /Markers1/Namespaces1
       Splitter Ratio: 0.5
-    Tree Height: 471
+    Tree Height: 668
   - Class: rviz/Selection
     Name: Selection
   - Class: rviz/Tool Properties
@@ -163,8 +166,6 @@ Visualization Manager:
           Value: false
         task:
           Value: true
-        world:
-          Value: false
       Marker Alpha: 1
       Marker Scale: 0.5
       Name: TF
@@ -188,7 +189,6 @@ Visualization Manager:
                               {}
                             panda_rightfinger:
                               {}
-        world:
           task:
             {}
       Update Interval: 0
@@ -221,6 +221,53 @@ Visualization Manager:
       Use Fixed Frame: true
       Use rainbow: true
       Value: true
+    - Alpha: 0.20000000298023224
+      Autocompute Intensity Bounds: false
+      Autocompute Value Bounds:
+        Max Value: 10
+        Min Value: -10
+        Value: true
+      Axis: Z
+      Channel Name: intensity
+      Class: rviz/PointCloud2
+      Color: 255; 255; 255
+      Color Transformer: Intensity
+      Decay Time: 0
+      Enabled: true
+      Invert Rainbow: false
+      Max Color: 138; 226; 52
+      Max Intensity: 1
+      Min Color: 239; 41; 41
+      Min Intensity: 0
+      Name: Quality
+      Position Transformer: XYZ
+      Queue Size: 10
+      Selectable: true
+      Size (Pixels): 3
+      Size (m): 0.007499999832361937
+      Style: Boxes
+      Topic: /quality
+      Unreliable: false
+      Use Fixed Frame: true
+      Use rainbow: false
+      Value: true
+    - Alpha: 1
+      Arrow Length: 0.30000001192092896
+      Axes Length: 0.05000000074505806
+      Axes Radius: 0.009999999776482582
+      Class: rviz/PoseArray
+      Color: 255; 25; 0
+      Enabled: true
+      Head Length: 0.07000000029802322
+      Head Radius: 0.029999999329447746
+      Name: Grasps
+      Queue Size: 10
+      Shaft Length: 0.23000000417232513
+      Shaft Radius: 0.009999999776482582
+      Shape: Axes
+      Topic: /grasps
+      Unreliable: false
+      Value: true
     - Alpha: 1
       Axes Length: 0.05000000074505806
       Axes Radius: 0.009999999776482582
@@ -242,7 +289,7 @@ Visualization Manager:
       Marker Topic: visualization_marker_array
       Name: Markers
       Namespaces:
-        bbox: true
+        bbox: false
         path: true
       Queue Size: 100
       Value: true
@@ -274,7 +321,7 @@ Visualization Manager:
   Views:
     Current:
       Class: rviz/Orbit
-      Distance: 1.2179546356201172
+      Distance: 1.2591626644134521
       Enable Stereo Rendering:
         Stereo Eye Separation: 0.05999999865889549
         Stereo Focal Distance: 1
@@ -282,25 +329,25 @@ Visualization Manager:
         Value: false
       Field of View: 0.7853981852531433
       Focal Point:
-        X: 0.2475447803735733
+        X: 0.2026486098766327
-        Y: 0.03526053577661514
+        Y: -0.0014560874551534653
-        Z: 0.4393550157546997
+        Z: 0.3729882836341858
       Focal Shape Fixed Size: true
       Focal Shape Size: 0.05000000074505806
       Invert Z Axis: false
       Name: Current View
       Near Clip Distance: 0.009999999776482582
-      Pitch: 0.2147984653711319
+      Pitch: 0.299798846244812
       Target Frame: <Fixed Frame>
-      Yaw: 5.383471488952637
+      Yaw: 5.213493824005127
     Saved: ~
 Window Geometry:
   Displays:
     collapsed: false
-  Height: 768
+  Height: 965
   Hide Left Dock: false
   Hide Right Dock: true
-  QMainWindow State: 000000ff00000000fd00000004000000000000015600000262fc0200000008fb0000001200530065006c0065006300740069006f006e00000001e10000009b0000005c00fffffffb0000001e0054006f006f006c002000500072006f007000650072007400690065007302000001ed000001df00000185000000a3fb000000120056006900650077007300200054006f006f02000001df000002110000018500000122fb000000200054006f006f006c002000500072006f0070006500720074006900650073003203000002880000011d000002210000017afb000000100044006900730070006c006100790073010000003d00000262000000c900fffffffb0000002000730065006c0065006300740069006f006e00200062007500660066006500720200000138000000aa0000023a00000294fb00000014005700690064006500530074006500720065006f02000000e6000000d2000003ee0000030bfb0000000c004b0069006e0065006300740200000186000001060000030c00000261000000010000010f000002b0fc0200000003fb0000001e0054006f006f006c002000500072006f00700065007200740069006500730100000041000000780000000000000000fb0000000a00560069006500770073000000003d000002b0000000a400fffffffb0000001200530065006c0065006300740069006f006e010000025a000000b200000000000000000000000200000490000000a9fc0100000001fb0000000a00560069006500770073030000004e00000080000002e10000019700000003000003e50000003efc0100000002fb0000000800540069006d00650100000000000003e5000002eb00fffffffb0000000800540069006d00650100000000000004500000000000000000000002890000026200000004000000040000000800000008fc00000002000000020000000000000002000000010000000a0054006f006f006c00730100000000ffffffff0000000000000000
+  QMainWindow State: 000000ff00000000fd00000004000000000000015600000327fc0200000008fb0000001200530065006c0065006300740069006f006e00000001e10000009b0000005c00fffffffb0000001e0054006f006f006c002000500072006f007000650072007400690065007302000001ed000001df00000185000000a3fb000000120056006900650077007300200054006f006f02000001df000002110000018500000122fb000000200054006f006f006c002000500072006f0070006500720074006900650073003203000002880000011d000002210000017afb000000100044006900730070006c006100790073010000003d00000327000000c900fffffffb0000002000730065006c0065006300740069006f006e00200062007500660066006500720200000138000000aa0000023a00000294fb00000014005700690064006500530074006500720065006f02000000e6000000d2000003ee0000030bfb0000000c004b0069006e0065006300740200000186000001060000030c00000261000000010000010f000002b0fc0200000003fb0000001e0054006f006f006c002000500072006f00700065007200740069006500730100000041000000780000000000000000fb0000000a00560069006500770073000000003d000002b0000000a400fffffffb0000001200530065006c0065006300740069006f006e010000025a000000b200000000000000000000000200000490000000a9fc0100000001fb0000000a00560069006500770073030000004e00000080000002e10000019700000003000004c50000003efc0100000002fb0000000800540069006d00650100000000000004c5000002eb00fffffffb0000000800540069006d00650100000000000004500000000000000000000003690000032700000004000000040000000800000008fc00000002000000020000000000000002000000010000000a0054006f006f006c00730100000000ffffffff0000000000000000
   Selection:
     collapsed: false
   Time:
@@ -309,6 +356,6 @@ Window Geometry:
     collapsed: false
   Views:
     collapsed: true
-  Width: 997
+  Width: 1221
-  X: 876
+  X: 1113
-  Y: 127
+  Y: 263