Draw colored camera views in rviz

active_grasp/nbv.py

@@ -30,13 +30,15 @@ class NextBestView(BasePolicy):
 
         # Generate viewpoints
         views = self.generate_viewpoints()
-        self.visualizer.views(views)
 
         # Evaluate viewpoints
         gains = [self.compute_ig(v) for v in views]
         costs = [self.compute_cost(v) for v in views]
         utilities = gains / np.sum(gains) - costs / np.sum(costs)
 
+        # Visualize
+        self.visualizer.views(self.intrinsic, views, utilities)
+
         # Determine next-best-view
         nbv = views[np.argmax(utilities)]
 

active_grasp/policy.py

@@ -43,6 +43,7 @@ class BasePolicy(Policy):
 
         self.center = 0.5 * (bbox.min + bbox.max)
         self.T_base_task = Transform.translation(self.center - np.full(3, 0.15))
+        self.T_task_base = self.T_base_task.inv()
         tf.broadcast(self.T_base_task, self.base_frame, self.task_frame)
         rospy.sleep(1.0)  # wait for the transform to be published
 

active_grasp/visualization.py

@@ -1,4 +1,5 @@
 from geometry_msgs.msg import PoseArray
+import matplotlib.colors
 import numpy as np
 import rospy
 
@@ -6,6 +7,8 @@ from robot_helpers.ros.rviz import *
 from robot_helpers.spatial import Transform
 from vgn.utils import *
 
+cmap = matplotlib.colors.LinearSegmentedColormap.from_list("RedGreen", ["r", "g"])
+
 
 class Visualizer:
     def __init__(self, frame, topic="visualization_marker_array"):
@@ -13,7 +16,6 @@ class Visualizer:
         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.views_pub = rospy.Publisher("views", PoseArray, queue_size=1)
         self.grasps_pub = rospy.Publisher("grasps", PoseArray, queue_size=1)
 
     def clear(self):
@@ -85,8 +87,64 @@ class Visualizer:
         msg = to_cloud_msg(frame, np.asarray(cloud.points))
         self.scene_cloud_pub.publish(msg)
 
-    def views(self, views):
+    def views(self, intrinsic, views, values):
-        msg = PoseArray()
+        vmin, vmax = min(values), max(values)
-        msg.header.frame_id = self.frame
+        scale = [0.002, 0.0, 0.0]
-        msg.poses = [to_pose_msg(view) for view in views]
+        near, far = 0.0, 0.02
-        self.views_pub.publish(msg)
+        markers = []
+        for i, (view, value) in enumerate(zip(views, values)):
+            color = cmap((value - vmin) / (vmax - vmin))
+            markers.append(
+                _create_cam_view_marker(
+                    self.frame,
+                    view,
+                    scale,
+                    color,
+                    intrinsic,
+                    near,
+                    far,
+                    ns="views",
+                    id=i,
+                )
+            )
+        self.draw(markers)
+
+
+def _create_cam_view_marker(
+    frame, pose, scale, color, intrinsic, near, far, ns="", id=0
+):
+    marker = create_marker(Marker.LINE_LIST, frame, pose, scale, color, ns, id)
+    x_n = near * intrinsic.width / (2.0 * intrinsic.fx)
+    y_n = near * intrinsic.height / (2.0 * intrinsic.fy)
+    z_n = near
+    x_f = far * intrinsic.width / (2.0 * intrinsic.fx)
+    y_f = far * intrinsic.height / (2.0 * intrinsic.fy)
+    z_f = far
+    points = [
+        [x_n, y_n, z_n],
+        [-x_n, y_n, z_n],
+        [-x_n, y_n, z_n],
+        [-x_n, -y_n, z_n],
+        [-x_n, -y_n, z_n],
+        [x_n, -y_n, z_n],
+        [x_n, -y_n, z_n],
+        [x_n, y_n, z_n],
+        [x_f, y_f, z_f],
+        [-x_f, y_f, z_f],
+        [-x_f, y_f, z_f],
+        [-x_f, -y_f, z_f],
+        [-x_f, -y_f, z_f],
+        [x_f, -y_f, z_f],
+        [x_f, -y_f, z_f],
+        [x_f, y_f, z_f],
+        [x_n, y_n, z_n],
+        [x_f, y_f, z_f],
+        [-x_n, y_n, z_n],
+        [-x_f, y_f, z_f],
+        [-x_n, -y_n, z_n],
+        [-x_f, -y_f, z_f],
+        [x_n, -y_n, z_n],
+        [x_f, -y_f, z_f],
+    ]
+    marker.points = [to_point_msg(p) for p in points]
+    return marker