From f43688023b9e6fe97bfd6adaaa1b4a8c8ff92a3b Mon Sep 17 00:00:00 2001
From: Michel Breyer <michel.p.breyer@gmail.com>
Date: Fri, 3 Sep 2021 11:50:14 +0200
Subject: [PATCH] Visualize best grasp

---
 active_grasp/nbv.py           |  2 +-
 active_grasp/policy.py        |  7 +++----
 active_grasp/visualization.py | 19 ++++++++++++++-----
 3 files changed, 18 insertions(+), 10 deletions(-)

diff --git a/active_grasp/nbv.py b/active_grasp/nbv.py
index 5eb3d51..5ae7895 100644
--- a/active_grasp/nbv.py
+++ b/active_grasp/nbv.py
@@ -24,7 +24,7 @@ class NextBestView(MultiViewPolicy):
         nbv, ig = views[i], gains[i]
 
         if self.check_stopping_criteria(ig):
-            self.vis.clear_views()
+            self.vis.clear_views(len(views))
             self.done = True
         else:
             return nbv.inv()  # the controller expects T_cam_base
diff --git a/active_grasp/policy.py b/active_grasp/policy.py
index d02c39d..74fc7ec 100644
--- a/active_grasp/policy.py
+++ b/active_grasp/policy.py
@@ -88,7 +88,7 @@ class SingleViewPolicy(Policy):
             out = self.vgn.predict(tsdf_grid)
             self.vis.quality(self.task_frame, voxel_size, out.qual)
 
-            grasps = select_grid(voxel_size, out, threshold=0.95)
+            grasps = select_grid(voxel_size, out, threshold=0.90)
             grasps, scores = self.sort_grasps(grasps)
 
             self.vis.grasps(self.base_frame, grasps, scores)
@@ -111,13 +111,12 @@ class MultiViewPolicy(Policy):
         tsdf_grid, voxel_size = self.tsdf.get_grid(), self.tsdf.voxel_size
         out = self.vgn.predict(tsdf_grid)
 
-        grasps = select_grid(voxel_size, out, threshold=0.95)
+        grasps = select_grid(voxel_size, out, threshold=0.90)
         grasps, scores = self.sort_grasps(grasps)
 
         if len(grasps) > 0:
             self.best_grasp = grasps[0]
-        else:
-            self.best_grasp = None
+            self.vis.best_grasp(self.base_frame, grasps[0], scores[0])
 
         self.vis.grasps(self.base_frame, grasps, scores)
 
diff --git a/active_grasp/visualization.py b/active_grasp/visualization.py
index ae2076c..5e5fb43 100644
--- a/active_grasp/visualization.py
+++ b/active_grasp/visualization.py
@@ -28,15 +28,19 @@ class Visualizer:
         self.quality_pub = rospy.Publisher("quality", PointCloud2, queue_size=1)
 
     def clear(self):
-        self.draw([Marker(action=Marker.DELETEALL)])
+        self.clear_markers()
         msg = to_cloud_msg("panda_link0", np.array([]))
         self.scene_cloud_pub.publish(msg)
         self.map_cloud_pub.publish(msg)
         self.quality_pub.publish(msg)
         rospy.sleep(0.1)
 
-    def clear_views(self):
-        self.draw([Marker(action=Marker.DELETE, ns="views")])
+    def clear_markers(self):
+        self.draw([Marker(action=Marker.DELETEALL)])
+
+    def clear_views(self, n):
+        markers = [Marker(action=Marker.DELETE, ns="views", id=i) for i in range(n)]
+        self.draw(markers)
 
     def draw(self, markers):
         self.marker_pub.publish(MarkerArray(markers=markers))
@@ -48,7 +52,12 @@ class Visualizer:
         marker = create_cube_marker(frame, pose, scale, color, ns="bbox")
         self.draw([marker])
 
-    def grasps(self, frame, grasps, scores, smin=0.9, smax=1.0, alpha=0.8):
+    def best_grasp(self, frame, grasp, score, smin=0.9, smax=1.0, alpha=1.0):
+        color = cmap((score - smin) / (smax - smin))
+        color = [color[0], color[1], color[2], alpha]
+        self.draw(create_grasp_markers(frame, grasp, color, "best_grasp", radius=0.01))
+
+    def grasps(self, frame, grasps, scores, smin=0.9, smax=1.0, alpha=0.6):
         if len(grasps) == 0:
             return
 
@@ -188,7 +197,7 @@ def create_grasp_markers(
     grasp,
     color,
     ns,
-    id,
+    id=0,
     finger_depth=0.05,
     radius=0.005,
 ):