def select_index(qual_vol, rot_vol, width_vol, index):
i, j, k = index
score = qual_vol[i, j, k]
ori = Rotation.from_quat(rot_vol[:, i, j, k])
pos = np.array([i, j, k], dtype=np.float64)
width = width_vol[i, j, k]
return Grasp(Transform(ori, pos), width), score
def read_grasp(df, i):
scene_id = df.loc[i, "scene_id"]
orientation = Rotation.from_quat(df.loc[i, "qx":"qw"].to_numpy(np.double))
position = df.loc[i, "x":"z"].to_numpy(np.double)
width = df.loc[i, "width"]
label = df.loc[i, "label"]
grasp = Grasp(Transform(orientation, position), width)
return scene_id, grasp, label
def main(args):
rospy.init_node("vgn_vis", anonymous=True)
dataset = Dataset(args.dataset, augment=args.augment)
i = np.random.randint(len(dataset))
voxel_grid, (label, rotations, width), index = dataset[i]
grasp = Grasp(Transform(Rotation.from_quat(rotations[0]), index), width)
vis.clear()
vis.draw_workspace(40)
vis.draw_tsdf(voxel_grid, 1.0)
vis.draw_grasp(grasp, float(label), 40.0 / 6.0)
rospy.sleep(1.0)
def __getitem__(self, i):
scene_id = self.df.loc[i, "scene_id"]
ori = Rotation.from_quat(self.df.loc[i, "qx":"qw"].to_numpy(np.single))
pos = self.df.loc[i, "i":"k"].to_numpy(np.single)
width = self.df.loc[i, "width"].astype(np.single)
label = self.df.loc[i, "label"].astype(np.long)
voxel_grid = read_voxel_grid(self.root, scene_id)
if self.augment:
voxel_grid, ori, pos = apply_transform(voxel_grid, ori, pos)
index = np.round(pos).astype(np.long)
rotations = np.empty((2, 4), dtype=np.single)
R = Rotation.from_rotvec(np.pi * np.r_[0.0, 0.0, 1.0])
rotations[0] = ori.as_quat()
rotations[1] = (ori * R).as_quat()
x, y, index = voxel_grid, (label, rotations, width), index
return x, y, index
def __init__(self, model_path):
# define frames
self.task_frame_id = "task"
self.cam_frame_id = "camera_depth_optical_frame"
self.T_cam_task = Transform(
Rotation.from_quat([-0.679, 0.726, -0.074, -0.081]),
[0.166, 0.101, 0.515])
# broadcast the tf tree (for visualization)
self.tf_tree = ros_utils.TransformTree()
self.tf_tree.broadcast_static(self.T_cam_task, self.cam_frame_id,
self.task_frame_id)
# define camera parameters
self.cam_topic_name = "/camera/depth/image_rect_raw"
self.intrinsic = CameraIntrinsic(640, 480, 383.265, 383.26, 319.39,
242.43)
# setup a CV bridge
self.cv_bridge = cv_bridge.CvBridge()
# construct the grasp planner object
self.device = torch.device(
"cuda" if torch.cuda.is_available() else "cpu")
self.net = load_network(model_path, self.device)
# initialize the visualization
vis.clear()
vis.draw_workspace(0.3)
# subscribe to the camera
self.img = None
rospy.Subscriber(self.cam_topic_name, sensor_msgs.msg.Image,
self.sensor_cb)
# setup cb to detect grasps
rospy.Timer(rospy.Duration(0.1), self.detect_grasps)
def from_quat_msg(msg):
"""Convert a Quaternion message to a Rotation object."""
return Rotation.from_quat([msg.x, msg.y, msg.z, msg.w])
def get_pose(self):
link_state = self.p.getLinkState(self.body_uid, self.link_index)
pos, ori = link_state[0], link_state[1]
return Transform(Rotation.from_quat(ori), pos)
def get_pose(self):
pos, ori = self.p.getBasePositionAndOrientation(self.uid)
return Transform(Rotation.from_quat(ori), np.asarray(pos))