diff --git a/active_grasp/nbv.py b/active_grasp/nbv.py index bf6ed45..07146fa 100644 --- a/active_grasp/nbv.py +++ b/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)] diff --git a/active_grasp/policy.py b/active_grasp/policy.py index 92be88d..178500a 100644 --- a/active_grasp/policy.py +++ b/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 diff --git a/active_grasp/visualization.py b/active_grasp/visualization.py index 4adbdca..c999463 100644 --- a/active_grasp/visualization.py +++ b/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): - msg = PoseArray() - msg.header.frame_id = self.frame - msg.poses = [to_pose_msg(view) for view in views] - self.views_pub.publish(msg) + def views(self, intrinsic, views, values): + vmin, vmax = min(values), max(values) + scale = [0.002, 0.0, 0.0] + near, far = 0.0, 0.02 + 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