def apply_transform(voxel_grid, orientation, position):
angle = np.pi / 2.0 * np.random.choice(4)
R_augment = Rotation.from_rotvec(np.r_[0.0, 0.0, angle])
z_offset = np.random.uniform(6, 34) - position[2]
t_augment = np.r_[0.0, 0.0, z_offset]
T_augment = Transform(R_augment, t_augment)
T_center = Transform(Rotation.identity(), np.r_[20.0, 20.0, 20.0])
T = T_center * T_augment * T_center.inverse()
# transform voxel grid
T_inv = T.inverse()
matrix, offset = T_inv.rotation.as_matrix(), T_inv.translation
voxel_grid[0] = ndimage.affine_transform(voxel_grid[0],
matrix,
offset,
order=0)
# transform grasp pose
position = T.transform_point(position)
orientation = T.rotation * orientation
return voxel_grid, orientation, position
def generate_packed_scene(self, object_count, table_height):
attempts = 0
max_attempts = 12
while self.num_objects < object_count and attempts < max_attempts:
self.save_state()
urdf = self.rng.choice(self.object_urdfs)
x = self.rng.uniform(0.08, 0.22)
y = self.rng.uniform(0.08, 0.22)
z = 1.0
angle = self.rng.uniform(0.0, 2.0 * np.pi)
rotation = Rotation.from_rotvec(angle * np.r_[0.0, 0.0, 1.0])
pose = Transform(rotation, np.r_[x, y, z])
scale = self.rng.uniform(0.7, 0.9)
body = self.world.load_urdf(urdf,
pose,
scale=self.global_scaling * scale)
lower, upper = self.world.p.getAABB(body.uid)
z = table_height + 0.5 * (upper[2] - lower[2]) + 0.002
body.set_pose(pose=Transform(rotation, np.r_[x, y, z]))
self.world.step()
if self.world.get_contacts(body):
self.world.remove_body(body)
self.restore_state()
else:
self.remove_and_wait()
attempts += 1
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 draw_grasp(grasp, score, finger_depth):
radius = 0.1 * finger_depth
w, d = grasp.width, finger_depth
color = cmap(float(score))
markers = []
# left finger
pose = grasp.pose * Transform(Rotation.identity(), [0.0, -w / 2, d / 2])
scale = [radius, radius, d]
msg = _create_marker_msg(Marker.CYLINDER, "task", pose, scale, color)
msg.id = 0
markers.append(msg)
# right finger
pose = grasp.pose * Transform(Rotation.identity(), [0.0, w / 2, d / 2])
scale = [radius, radius, d]
msg = _create_marker_msg(Marker.CYLINDER, "task", pose, scale, color)
msg.id = 1
markers.append(msg)
# wrist
pose = grasp.pose * Transform(Rotation.identity(), [0.0, 0.0, -d / 4])
scale = [radius, radius, d / 2]
msg = _create_marker_msg(Marker.CYLINDER, "task", pose, scale, color)
msg.id = 2
markers.append(msg)
# palm
pose = grasp.pose * Transform(
Rotation.from_rotvec(np.pi / 2 * np.r_[1.0, 0.0, 0.0]),
[0.0, 0.0, 0.0])
scale = [radius, radius, w]
msg = _create_marker_msg(Marker.CYLINDER, "task", pose, scale, color)
msg.id = 3
markers.append(msg)
pubs["grasp"].publish(MarkerArray(markers=markers))
