hofee/nbv_sim
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Visualize quality and grasps

Browse Source
This commit is contained in:
Michel Breyer
2021-08-10 18:52:03 +02:00
parent 5a33561abb
commit 1e6d933e53
3 changed files with 127 additions and 64 deletions
Download Patch File Download Diff File Expand all files Collapse all files

79
active_grasp/policy.py
View File

@@ -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(tsdf_grid)
out = self.vgn.predict(self.tsdf.get_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
mask = self.qual_hist[:T, ...] > 0.0
# The next line prints a warning since some voxels have no grasp
# 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))
qual = self.compute_mean_quality()
index_list = select_local_maxima(qual, 0.9, 3)
grasps = [g for i in index_list if (g := self.select_mean_at(i))]
else:
tsdf_grid = self.tsdf.get_grid()
out = self.vgn.predict(tsdf_grid)
qual = threshold_quality(out.qual, 0.9)
index_list = select_local_maxima(qual, 3)
out = self.vgn.predict(self.tsdf.get_grid())
qual = out.qual
index_list = select_local_maxima(qual, 0.9, 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.select_grasps_on_target_object(grasps)
grasps = self.transform_and_reject(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):
for grasp in grasps:
grasp.pose = self.T_base_task * grasp.pose
return grasps
def compute_mean_quality(self):
qual = np.mean(self.qual_hist, axis=0, where=self.qual_hist > 0.0)
return np.nan_to_num(qual, copy=False) # mean of empty slices returns nan
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

31
active_grasp/visualization.py
View File

@@ -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])
self.draw([Marker(action=Marker.DELETEALL)])
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)