def _add_cube_into_image(self, obs, location, color=[20,200,200]):
cube = Cube(origin=np.array(location), scale=self.cube_size)
self.cube_image = final = copy(obs["rgb_filled"])
roll, pitch, yaw = self.robot.get_rpy()
x, y, z = self.robot.eyes.get_position()
size = self.config["resolution"] // 2
fov = self.config["fov"]
world_to_image_mat = generate_projection_matrix(x, y, z, yaw, pitch, roll, fov, fov, size, size)
if self.use_texture:
masks, xx_faces, yy_faces = get_cube_depth_and_faces(cube, world_to_image_mat, size*2, size*2)
for face_mask, xx, yy in zip(masks, xx_faces, yy_faces):
transform = skimage.transform.ProjectiveTransform()
dest = np.array([yy,xx]).T
try:
transform.estimate(self.texture_points, dest)
self.new_img = skimage.transform.warp(self.texture_image, transform.inverse)
except:
continue
img_mask = face_mask < self.depth
self.cube_image[img_mask] = np.array(self.new_img[img_mask] * 255, dtype=np.uint8)
self.depth[img_mask] = face_mask[img_mask]
else:
cube_idx = draw_cube(cube, world_to_image_mat, size*2, size*2, fast_depth=True)
self.cube_image[cube_idx < self.depth] = np.array(color)
return self.cube_image
def _add_obstacle_into_image(self, obs, location, color=[0, 159, 107]):
cube = Cube(origin=np.array(location), scale=self.cube_size)
roll, pitch, yaw = self.robot.get_rpy()
x, y, z = self.robot.eyes.get_position()
size = self.config["resolution"] // 2
fov = self.config["fov"]
world_to_image_mat = generate_projection_matrix(x, y, z, yaw, pitch, roll, fov, fov, size, size)
cube_idx = draw_cube(cube, world_to_image_mat, size*2, size*2, fast_depth=True)
self.cube_image[np.logical_and(cube_idx < self.depth, cube_idx < self.target_depth)] = np.array(color)
return self.cube_image