179 lines
6.3 KiB
Python
179 lines
6.3 KiB
Python
import numpy as np
|
|
from sensor_msgs.msg import CameraInfo
|
|
from pathlib import Path
|
|
import rospy
|
|
|
|
from .timer import Timer
|
|
from .visualization import Visualizer
|
|
from robot_helpers.model import KDLModel
|
|
from robot_helpers.ros import tf
|
|
from robot_helpers.ros.conversions import *
|
|
from vgn.detection import *
|
|
from vgn.perception import UniformTSDFVolume
|
|
|
|
|
|
class Policy:
|
|
def __init__(self):
|
|
self.load_parameters()
|
|
self.init_robot_model()
|
|
self.init_visualizer()
|
|
|
|
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")
|
|
msg = rospy.wait_for_message(info_topic, CameraInfo, rospy.Duration(2.0))
|
|
self.intrinsic = from_camera_info_msg(msg)
|
|
self.qual_threshold = rospy.get_param("vgn/qual_threshold")
|
|
|
|
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()
|
|
|
|
def is_feasible(self, view, q_init=None):
|
|
q_init = q_init if q_init else [0.0, -0.79, 0.0, -2.356, 0.0, 1.57, 0.79]
|
|
return self.model.ik(q_init, view) is not None
|
|
|
|
def activate(self, bbox, view_sphere):
|
|
self.vis.clear()
|
|
|
|
self.bbox = bbox
|
|
self.view_sphere = view_sphere
|
|
self.vis.bbox(self.base_frame, self.bbox)
|
|
|
|
self.calibrate_task_frame()
|
|
|
|
self.tsdf = UniformTSDFVolume(0.3, 40)
|
|
self.vgn = VGN(Path(rospy.get_param("vgn/model")))
|
|
|
|
self.views = []
|
|
self.best_grasp = None
|
|
self.x_d = None
|
|
self.done = False
|
|
self.info = {}
|
|
|
|
def calibrate_task_frame(self):
|
|
self.T_base_task = Transform.translation(self.bbox.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(0.5) # Wait for tf tree to be updated
|
|
self.vis.workspace(self.task_frame, 0.3)
|
|
|
|
def update(self, img, x, q):
|
|
raise NotImplementedError
|
|
|
|
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 = [], []
|
|
|
|
for grasp in in_grasps:
|
|
pose = self.T_base_task * grasp.pose
|
|
R, t = pose.rotation, pose.translation
|
|
tip = pose.rotation.apply([0, 0, 0.05]) + pose.translation
|
|
|
|
# Filter out artifacts close to the support
|
|
if t[2] < self.bbox.min[2] + 0.05 or tip[2] < self.bbox.min[2] + 0.02:
|
|
continue
|
|
|
|
if self.bbox.is_inside(tip):
|
|
grasp.pose = pose
|
|
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, q, q_grasp):
|
|
return -np.linalg.norm(q - q_grasp)
|
|
|
|
|
|
class SingleViewPolicy(Policy):
|
|
def update(self, img, x, q):
|
|
linear, _ = compute_error(self.x_d, x)
|
|
if np.linalg.norm(linear) < 0.02:
|
|
self.views.append(x)
|
|
self.tsdf.integrate(img, self.intrinsic, x.inv() * self.T_base_task)
|
|
tsdf_grid, voxel_size = self.tsdf.get_grid(), self.tsdf.voxel_size
|
|
self.vis.scene_cloud(self.task_frame, self.tsdf.get_scene_cloud())
|
|
self.vis.map_cloud(self.task_frame, self.tsdf.get_map_cloud())
|
|
|
|
out = self.vgn.predict(tsdf_grid)
|
|
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, 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.done = True
|
|
|
|
|
|
class MultiViewPolicy(Policy):
|
|
def activate(self, bbox, view_sphere):
|
|
super().activate(bbox, view_sphere)
|
|
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, q):
|
|
self.views.append(x)
|
|
self.vis.path(self.base_frame, self.views)
|
|
|
|
with Timer("tsdf_integration"):
|
|
self.tsdf.integrate(img, self.intrinsic, x.inv() * self.T_base_task)
|
|
self.vis.scene_cloud(self.task_frame, self.tsdf.get_scene_cloud())
|
|
self.vis.map_cloud(self.task_frame, self.tsdf.get_map_cloud())
|
|
|
|
with Timer("grasp_prediction"):
|
|
tsdf_grid, voxel_size = self.tsdf.get_grid(), self.tsdf.voxel_size
|
|
out = self.vgn.predict(tsdf_grid)
|
|
self.vis.quality(self.task_frame, self.tsdf.voxel_size, out.qual, 0.5)
|
|
|
|
t = (len(self.views) - 1) % self.T
|
|
self.qual_hist[t, ...] = out.qual
|
|
|
|
with Timer("grasp_selection"):
|
|
grasps = select_grid(voxel_size, out, threshold=self.qual_threshold)
|
|
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)
|
|
else:
|
|
self.best_grasp = None
|
|
# TODO clear grasp markers
|
|
|
|
|
|
def compute_error(x_d, x):
|
|
linear = x_d.translation - x.translation
|
|
angular = (x_d.rotation * x.rotation.inv()).as_rotvec()
|
|
return linear, angular
|
|
|
|
|
|
registry = {}
|
|
|
|
|
|
def register(id, cls):
|
|
global registry
|
|
registry[id] = cls
|
|
|
|
|
|
def make(id, *args, **kwargs):
|
|
if id in registry:
|
|
return registry[id](*args, **kwargs)
|
|
else:
|
|
raise ValueError("{} policy does not exist.".format(id))
|