def grasp_convert(self, action: Action,
images: List[OrthographicImage]) -> None:
image = images[0]
mat_area_image = get_area_of_interest(image,
action.pose,
border_color=np.nan,
project_3d=False).mat
# cv2.imwrite('/tmp/test_grasp.png', mat_area_image)
mat_area_image = mat_area_image.astype(np.float32)
mat_area_image[mat_area_image ==
0] = np.nan # Make every not found pixel NaN
# Get distance at gripper for possible collisions
gripper_one_side_size = 0.5 * image.pixel_size * (action.pose.d + 0.002
) # [px]
area_center = crop(mat_area_image,
(image.pixel_size * 0.01, image.pixel_size * 0.03))
side_gripper_image_size = (image.pixel_size * 0.036,
image.pixel_size * 0.036)
area_left = crop(mat_area_image, side_gripper_image_size,
(-gripper_one_side_size, 0))
area_right = crop(mat_area_image, side_gripper_image_size,
(gripper_one_side_size, 0))
# z_raw = image.depth_from_value(np.nanmax(area_center))
z_raw = image.depth_from_value(np.nanmedian(area_center))
z_raw_left = image.depth_from_value(np.nanmin(area_left))
z_raw_right = image.depth_from_value(np.nanmin(area_right))
z_raw += self.grasp_z_offset
z_raw_collision = max(z_raw_left, z_raw_right) - 0.008
action.pose.z = -min(
z_raw, z_raw_collision) # Get the maximum [m] for impedance mode
def test_image_transformation(self):
image = OrthographicImage(self.read_image(self.file_path / self.image_name), 2000.0, 0.22, 0.4, '', Config.default_image_pose)
pose = RobotPose(affine=Affine(0.02, 0.0, 0.0))
area_image = get_area_of_interest(image, pose, border_color=(0))
imageio.imwrite(str(self.file_path / 'gen-x20-b.png'), area_image.mat)
pose = RobotPose(affine=Affine(0.03, 0.03, 0.0))
area_image = get_area_of_interest(image, pose, border_color=(0))
imageio.imwrite(str(self.file_path / 'gen-x30-y30-b.png'), area_image.mat)
pose = RobotPose(affine=Affine(0.01, 0.04, 0.0, 0.4))
area_image = get_area_of_interest(image, pose, border_color=(0))
imageio.imwrite(str(self.file_path / 'gen-x10-y40-a04-b.png'), area_image.mat)
image = image.translate([0.0, 0.0, 0.05])
image = image.rotate_x(-0.3, [0.0, 0.25])
imageio.imwrite(str(self.file_path / 'gen-rot0_3.png'), image.mat)
def print_before_after_image(episode_id: str):
action, before_image, after_image = Loader.get_action(
'shifting', episode_id, ['ed-v', 'ed-after'])
area_before_image = get_area_of_interest(before_image,
action.pose,
size_cropped=(300, 300),
size_result=(150, 150))
area_after_image = get_area_of_interest(after_image,
action.pose,
size_cropped=(300, 300))
white = [255 * 255] * 3
draw_line(area_before_image,
action.pose,
Affine(0, 0),
Affine(0.036, 0),
color=white,
thickness=2)
draw_line(area_before_image,
action.pose,
Affine(0.036, 0.0),
Affine(0.026, -0.008),
color=white,
thickness=2)
draw_line(area_before_image,
action.pose,
Affine(0.036, 0.0),
Affine(0.026, 0.008),
color=white,
thickness=2)
cv2.imwrite(
str(Path.home() / 'Desktop' / f'example-{episode_id}-before.png'),
area_before_image.mat)
cv2.imwrite(
str(Path.home() / 'Desktop' / f'example-{episode_id}-after.png'),
area_after_image.mat)
print('---')
print(episode_id)
def save_episode(predictor, database, episode_id, reward=1, action_type=1):
action, image, image_after = Loader.get_action(database, episode_id,
['ed-v', 'ed-after'])
draw_around_box(image, box=Config.box)
draw_around_box(image_after, box=Config.box)
# background_color = image.value_from_depth(get_distance_to_box(image, Config.box))
area = get_area_of_interest(image,
action.pose,
size_cropped=(256, 256),
size_result=predictor.size)
area_after = get_area_of_interest(image_after,
action.pose,
size_cropped=(256, 256),
size_result=predictor.size)
result = predictor.predict(area,
reward=reward,
action_type=action_type,
sampling=True,
number=20)
save_dir = Path.home() / 'Desktop' / 'predict-examples' / episode_id
save_dir.mkdir(exist_ok=True)
cv2.imwrite(str(save_dir / f'{predictor.name}_s_bef.png'), area.mat)
cv2.imwrite(str(save_dir / f'{predictor.name}_s_aft.png'), area_after.mat)
cv2.imwrite(str(save_dir / f'{predictor.name}_result.png'),
result[0] * 255)
if predictor.uncertainty:
result[result < 0.1] = np.nan
uncertainty = np.nanvar(result, axis=0)
uncertainty /= np.nanmax(uncertainty) * 0.25
uncertainty = np.clip(uncertainty * 255, 0, 255).astype(np.uint8)
uncertainty = cv2.applyColorMap(uncertainty, cv2.COLORMAP_JET)
cv2.imwrite(str(save_dir / f'{predictor.name}_unc.png'), uncertainty)
def get_area_image(self,
image,
pose,
size_cropped,
size_output,
border_color=None):
box_distance = get_distance_to_box(image, self.box)
border_color = border_color if border_color else image.value_from_depth(
box_distance)
return get_area_of_interest(image,
pose,
size_cropped=size_cropped,
size_result=size_output,
border_color=border_color)
def place_convert(self, action: Action,
images: List[OrthographicImage]) -> None:
image = images[0]
mat_area_image = get_area_of_interest(image,
action.pose,
border_color=np.nan).mat
mat_area_image = mat_area_image.astype(np.float32)
mat_area_image[mat_area_image ==
0] = np.nan # Make every not found pixel NaN
# Get distance at gripper for possible collisions
area_center = crop(mat_area_image,
(image.pixel_size * 0.01, image.pixel_size * 0.03))
z_raw = image.depth_from_value(np.nanmedian(area_center))
z_raw += self.place_z_offset
action.pose.z = -z_raw # [m] Move slightly up for gripper center point
def infer_at_pose(self, images: List[OrthographicImage], pose: Affine):
input_images = []
for image in images:
draw_around_box(image, box=self.box)
area_mat = get_area_of_interest(
image,
pose,
size_cropped=(200, 200),
size_result=(32, 32),
border_color=image.value_from_depth(get_distance_to_box(image, self.box)),
).mat
area_mat = np.array(area_mat, dtype=np.float32) / np.iinfo(area_mat.dtype).max # 2 * x - 1.0
area_mat_exp = np.expand_dims(area_mat, axis=-1)
input_images.append(area_mat_exp)
if self.monte_carlo:
input_images_mc = [np.array([image for i in range(self.monte_carlo)]) for image in input_images]
estimated_rewards_sampling = self.model.predict(input_images_mc)
estimated_reward = np.mean(estimated_rewards_sampling, axis=0)
estimated_reward_std = self.mutual_information(estimated_rewards_sampling)
return estimated_reward, estimated_reward_std
return self.model.predict([input_images])[0]
def predict_images_after_action(
self,
images: List[OrthographicImage],
action: Action,
reward: float,
action_type: int,
uncertainty_images=None,
) -> List[OrthographicImage]:
image = images[0]
uncertainty_image = uncertainty_images[0]
start = time.time()
draw_around_box(image, box=Config.box)
area = get_area_of_interest(image,
action.pose,
size_cropped=(256, 256),
size_result=(64, 64))
area_input = np.expand_dims(area.mat, axis=2).astype(
np.float32) / np.iinfo(np.uint16).max * 2 - 1
reward = np.expand_dims(np.expand_dims(np.expand_dims(reward, axis=1),
axis=1),
axis=1).astype(np.float32)
action_type = np.expand_dims(np.expand_dims(action_type, axis=1),
axis=1)
use_monte_carlo = self.monte_carlo and self.monte_carlo > 1
if not use_monte_carlo:
area_result = self.prediction_model.predict([[area_input],
[reward],
[action_type],
np.zeros(
(1, 1, 1, 8))])[0]
area_result = np.array(np.iinfo(np.uint16).max *
(area_result + 1) / 2,
dtype=np.uint16)
else:
latent = np.random.normal(scale=0.05,
size=(self.monte_carlo, 1, 1, 8))
if self.monte_carlo > 3:
latent[0, :, :, :] = 0.0
predictions = self.prediction_model.predict([
[area_input for _ in range(self.monte_carlo)],
[reward for _ in range(self.monte_carlo)],
[action_type for _ in range(self.monte_carlo)],
latent,
])
predictions = (predictions + 1) / 2
predictions = np.array(predictions, dtype=np.float32)
area_result = predictions[0]
area_result = np.array(np.iinfo(np.uint16).max * area_result,
dtype=np.uint16)
predictions[predictions < 0.1] = np.nan
area_uncertainty = np.nanvar(predictions, axis=0)
area_uncertainty *= 7e3
area_uncertainty[area_uncertainty > 1] = 1
area_uncertainty = np.array(np.iinfo(np.uint16).max *
area_uncertainty,
dtype=np.uint16)
uncertainty_image = patch_image_at(
uncertainty_image,
area_uncertainty,
action.pose,
size_cropped=(256, 256),
operation='add',
)
result = patch_image_at(image,
area_result,
action.pose,
size_cropped=(256, 256))
logger.info(f'Predicted image [s]: {time.time() - start:0.3f}')
if use_monte_carlo:
return [result], [uncertainty_image]
return [result]