From 9f06bb27be0cafd66b88211c5d5fd61aac36bb9a Mon Sep 17 00:00:00 2001 From: Michel Breyer Date: Sun, 12 Sep 2021 16:23:10 +0200 Subject: [PATCH] Add C-space score function --- active_grasp/nbv.py | 4 ++-- active_grasp/policy.py | 33 +++++++++++++++++++++------------ active_grasp/visualization.py | 10 ++++++---- 3 files changed, 29 insertions(+), 18 deletions(-) diff --git a/active_grasp/nbv.py b/active_grasp/nbv.py index 778cac0..d649194 100644 --- a/active_grasp/nbv.py +++ b/active_grasp/nbv.py @@ -56,7 +56,7 @@ class NextBestView(MultiViewPolicy): def __init__(self): super().__init__() self.min_z_dist = rospy.get_param("~camera/min_z_dist") - self.max_views = 40 + self.max_views = 100 def activate(self, bbox, view_sphere): super().activate(bbox, view_sphere) @@ -78,7 +78,7 @@ class NextBestView(MultiViewPolicy): self.done = True else: with Timer("state_update"): - self.integrate(img, x) + self.integrate(img, x, q) views = self.view_candidates with Timer("ig_computation"): gains = [self.ig_fn(v) for v in views] diff --git a/active_grasp/policy.py b/active_grasp/policy.py index c83f170..e87d91b 100644 --- a/active_grasp/policy.py +++ b/active_grasp/policy.py @@ -20,6 +20,7 @@ class Policy: def load_parameters(self): self.base_frame = rospy.get_param("~base_frame_id") + self.T_grasp_ee = Transform.from_list(rospy.get_param("~ee_grasp_offset")).inv() self.cam_frame = rospy.get_param("~camera/frame_id") self.task_frame = "task" info_topic = rospy.get_param("~camera/info_topic") @@ -29,6 +30,7 @@ class Policy: def init_robot_model(self): self.model = KDLModel.from_parameter_server(self.base_frame, self.cam_frame) + self.ee_model = KDLModel.from_parameter_server(self.base_frame, "panda_link8") def init_visualizer(self): self.vis = Visualizer() @@ -65,7 +67,7 @@ class Policy: def update(self, img, x, q): raise NotImplementedError - def sort_grasps(self, in_grasps): + def sort_grasps(self, in_grasps, q): # Transforms grasps into base frame, checks whether they lie on the target, and sorts by their score grasps, scores = [], [] @@ -80,15 +82,17 @@ class Policy: tip = pose.rotation.apply([0, 0, 0.05]) + pose.translation if self.bbox.is_inside(tip): grasp.pose = pose - grasps.append(grasp) - scores.append(self.score_fn(grasp)) + q_grasp = self.ee_model.ik(q, pose * self.T_grasp_ee) + if q_grasp is not None: + grasps.append(grasp) + scores.append(self.score_fn(grasp, q, q_grasp)) grasps, scores = np.asarray(grasps), np.asarray(scores) indices = np.argsort(-scores) return grasps[indices], scores[indices] - def score_fn(self, grasp): - return grasp.quality + def score_fn(self, grasp, q, q_grasp): + return -np.linalg.norm(q - q_grasp) class SingleViewPolicy(Policy): @@ -105,10 +109,14 @@ class SingleViewPolicy(Policy): self.vis.quality(self.task_frame, voxel_size, out.qual, 0.5) grasps = select_grid(voxel_size, out, threshold=self.qual_threshold) - grasps, scores = self.sort_grasps(grasps) - self.vis.grasps(self.base_frame, grasps, scores) + grasps, scores = self.sort_grasps(grasps, q) + + if len(grasps) > 0: + smin, smax = np.min(scores), np.max(scores) + self.best_grasp = grasps[0] + self.vis.grasps(self.base_frame, grasps, scores, smin, smax) + self.vis.best_grasp(self.base_frame, grasps[0], scores[0], smin, smax) - self.best_grasp = grasps[0] if len(grasps) > 0 else None self.done = True @@ -118,7 +126,7 @@ class MultiViewPolicy(Policy): self.T = 10 # Window size of grasp prediction history self.qual_hist = np.zeros((self.T,) + (40,) * 3, np.float32) - def integrate(self, img, x): + def integrate(self, img, x, q): self.views.append(x) self.vis.path(self.base_frame, self.views) @@ -137,12 +145,13 @@ class MultiViewPolicy(Policy): with Timer("grasp_selection"): grasps = select_grid(voxel_size, out, threshold=self.qual_threshold) - grasps, scores = self.sort_grasps(grasps) - self.vis.grasps(self.base_frame, grasps, scores) + grasps, scores = self.sort_grasps(grasps, q) if len(grasps) > 0: + smin, smax = np.min(scores), np.max(scores) self.best_grasp = grasps[0] - self.vis.best_grasp(self.base_frame, grasps[0], scores[0]) + self.vis.grasps(self.base_frame, grasps, scores, smin, smax) + self.vis.best_grasp(self.base_frame, grasps[0], scores[0], smin, smax) def compute_error(x_d, x): diff --git a/active_grasp/visualization.py b/active_grasp/visualization.py index c549d7b..ce429d6 100644 --- a/active_grasp/visualization.py +++ b/active_grasp/visualization.py @@ -1,4 +1,4 @@ -from geometry_msgs.msg import PoseArray +from geometry_msgs.msg import PoseStamped import matplotlib.colors import numpy as np import rospy @@ -27,6 +27,7 @@ class Visualizer: latch=True, queue_size=1, ) + self.pose_pub = rospy.Publisher("pose", PoseStamped, queue_size=1) self.quality_pub = rospy.Publisher("quality", PointCloud2, queue_size=1) def clear(self): @@ -76,9 +77,6 @@ class Visualizer: self.draw(create_grasp_markers(frame, grasp, color, "best_grasp", radius=0.006)) def grasps(self, frame, grasps, scores, smin=0.9, smax=1.0, alpha=0.8): - if len(grasps) == 0: - return - markers = [] for i, (grasp, score) in enumerate(zip(grasps, scores)): color = cmap((score - smin) / (smax - smin)) @@ -137,6 +135,10 @@ class Visualizer: ) self.draw([marker]) + def pose(self, frame, pose): + msg = to_pose_stamped_msg(pose, frame) + self.pose_pub.publish(msg) + def quality(self, frame, voxel_size, quality, threshold=0.9): points, values = grid_to_map_cloud(voxel_size, quality, threshold) values = (values - 0.9) / (1.0 - 0.9) # to increase contrast