def draw_around_box(image: OrthographicImage,
box: Dict[str, List[float]],
draw_lines=False) -> None:
bin_border = get_rect(size=box['size'], center=box['center'])
image_border = get_rect(size=[10.0, 10.0, box['size'][2]])
image_pose = Affine(image.pose.x, image.pose.y, -image.pose.z,
image.pose.a, image.pose.b, image.pose.c)
bin_border_projection = [image.project(image_pose * p) for p in bin_border]
color = [
max(
image.value_from_depth((image_pose * border.inverse()).z)
for border in bin_border)
] * 3
if draw_lines:
color_direction = (0, 255 * 255, 0) # Green
for i in range(len(bin_border)):
cv2.line(image.mat, tuple(bin_border_projection[i]),
tuple(bin_border_projection[(i + 1) % len(bin_border)]),
color_direction, 1)
else:
image_border_projection = [
image.project(image_pose * p) for p in image_border
]
cv2.fillPoly(
image.mat,
np.array([image_border_projection, bin_border_projection]), color)
def draw_around_box(image: OrthographicImage, box: Dict[str, List[float]], draw_lines=False) -> None:
bin_border = get_rect(size=box['size'], center=box['center'])
image_border = get_rect(size=[10.0, 10.0, box['size'][2]])
image_pose = Affine(image.pose)
color_divisor = 255 * 255 if image.mat.dtype == np.float32 else 1
bin_border_projection = [image.project((image_pose * p).to_array()) for p in bin_border]
if draw_lines:
color_direction = np.array([0, 255 * 255, 0]) / color_divisor # Green
for i in range(len(bin_border)):
cv2.line(
image.mat, tuple(bin_border_projection[i]),
tuple(bin_border_projection[(i + 1) % len(bin_border)]),
color_direction, 1
)
else:
color = np.array([
max(image.value_from_depth((image_pose * border.inverse()).z) for border in bin_border)
] * 3) / color_divisor
image_border_projection = [image.project((image_pose * p).to_array()) for p in image_border]
cv2.fillPoly(image.mat, np.array([image_border_projection, bin_border_projection]), color)
def draw_pose(image: OrthographicImage,
action_pose: RobotPose,
convert_to_rgb=False,
reference_pose=None) -> None:
if convert_to_rgb and image.mat.ndim == 2:
image.mat = cv2.cvtColor(image.mat, cv2.COLOR_GRAY2RGB)
color_rect = (255 * 255, 0, 0) # Blue
color_lines = (0, 0, 255 * 255) # Red
color_direction = (0, 255 * 255, 0) # Green
rect = get_rect([200.0 / image.pixel_size, 200.0 / image.pixel_size, 0.0])
for i, r in enumerate(rect):
draw_line(image, action_pose, r, rect[(i + 1) % len(rect)], color_rect,
2, reference_pose)
draw_line(image, action_pose, Affine(90 / image.pixel_size, 0),
Affine(100 / image.pixel_size, 0), color_rect, 2, reference_pose)
draw_line(image, action_pose, Affine(0.012, action_pose.d / 2),
Affine(-0.012, action_pose.d / 2), color_lines, 1,
reference_pose)
draw_line(image, action_pose, Affine(0.012, -action_pose.d / 2),
Affine(-0.012, -action_pose.d / 2), color_lines, 1,
reference_pose)
draw_line(image, action_pose, Affine(0, action_pose.d / 2),
Affine(0, -action_pose.d / 2), color_lines, 1, reference_pose)
draw_line(image, action_pose, Affine(0.006, 0), Affine(-0.006, 0),
color_lines, 1, reference_pose)
if not isinstance(action_pose.z, str) and np.isfinite(action_pose.z):
draw_line(image, action_pose, Affine(z=0.14), Affine(),
color_direction, 1, reference_pose)
def test_drawing(self):
image = OrthographicImage(self.read_image(self.file_path / self.image_name), 2000.0, 0.22, 0.4, '', Config.default_image_pose)
self.assertEqual(image.mat.dtype, np.uint16)
draw_around_box(image, box=Config.box)
imageio.imwrite(str(self.file_path / 'gen-draw-around-bin-unit.png'), image.mat)
def test_shift_conversion(self):
conv = Converter(shift_z_offset=0.0, box=self.box)
image = OrthographicImage(
self.read_image(self.file_path / self.image_name), 2000.0, 0.22,
0.4, '', Config.default_image_pose)
action = Action()
action.type = 'shift'
action.pose = RobotPose()
action.pose.x = -0.02
action.pose.y = 0.105
action.pose.a = 1.52
action.pose.d = 0.078
action.index = 1
action.shift_motion = [0.03, 0.0]
draw_pose(image, action.pose, convert_to_rgb=True)
draw_around_box(image, box=self.box, draw_lines=True)
imageio.imwrite(str(self.file_path / 'gen-shift-draw-pose.png'),
image.mat)
self.assertTrue(conv.shift_check_safety(action, [image]))
conv.calculate_pose(action, [image])
self.assertLess(action.pose.z, 0.0)
def infer_grasp(self, req: InferGraspRequest) -> InferGraspResponse:
image_res = self.get_images()
cv_image = self.bridge.imgmsg_to_cv2(image_res.image, 'mono16')
image = OrthographicImage(cv_image, image_res.pixel_size,
image_res.min_depth, image_res.max_depth)
images = [image]
action = self.inference.infer(images, req.method)
draw_image = clone(image)
draw_image.mat = cv2.cvtColor(draw_image.mat, cv2.COLOR_GRAY2RGB)
if self.draw_on_image:
draw_around_box(draw_image, box=self.box, draw_lines=True)
draw_pose(draw_image, action)
if self.save_tmp_image:
cv2.imwrite('/tmp/graspro/current-image.png', draw_image.mat)
draw_image_msg = self.bridge.cv2_to_imgmsg(draw_image.mat, 'rgb8')
self.image_publisher.publish(draw_image_msg)
if self.publish_heatmap:
heatmap = self.heatmapper.render(image, alpha=0.5)
heatmap_msg = self.bridge.cv2_to_imgmsg(heatmap, 'rgb8')
self.heatmap_publisher.publish(heatmap_msg)
return InferGraspResponse(action=action)
def add_camera(service, suffixes: List[str],
images: Dict[str, OrthographicImage]) -> None:
result = service(suffixes)
for i, img in enumerate(result.images):
mat = self.bridge.imgmsg_to_cv2(img.image, img.image.encoding)
images[suffixes[i]] = OrthographicImage(
mat, img.pixel_size, img.min_depth, img.max_depth,
img.camera)
def clone(image: OrthographicImage) -> OrthographicImage:
return OrthographicImage(
image.mat.copy(),
image.pixel_size,
image.min_depth,
image.max_depth,
image.camera,
image.pose,
)
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 get_image(cls,
collection: str,
episode_id: str,
action_id: int,
suffix: str,
images=None,
image_data=None,
as_float=False) -> OrthographicImage:
image = cv2.imread(
str(cls.get_image_path(collection, episode_id, action_id, suffix)),
cv2.IMREAD_UNCHANGED)
if image is None:
raise FileNotFoundError(
f'Image {collection} {episode_id} {action_id} {suffix} not found.'
)
if not as_float and image.dtype == np.uint8:
image = image.astype(np.uint16)
image *= 255
elif as_float:
mult = 255 if image.dtype == np.uint8 else 255 * 255
image = image.astype(np.float32)
image /= mult
if image_data:
image_data_result = {
'info': image_data['info'],
'pose': RobotPose(data=image_data['pose']),
}
elif images:
image_data_result = images[suffix]
else:
episode = cls.database[collection].find_one({'id': episode_id},
{'actions.images': 1})
if not episode or not episode['actions'] or not (
0 <= action_id < len(episode['actions'])):
raise Exception(
f'Internal mismatch of image {collection} {episode_id} {action_id} {suffix} not found.'
)
image_data_result = Action(
data=episode['actions'][action_id]).images[suffix]
return OrthographicImage(
image,
image_data_result['info']['pixel_size'],
image_data_result['info']['min_depth'],
image_data_result['info']['max_depth'],
suffix.split('-')[0],
image_data_result['pose'].to_array(),
# list(image_data_result['pose'].values()),
# Affine(image_data_result['pose']).to_array(),
)
def pose_from_index(self, index, index_shape, example_image: OrthographicImage) -> RobotPose:
vec = self.rotate_vector((
example_image.position_from_index(index[1], index_shape[1]),
example_image.position_from_index(index[2], index_shape[2])
), self.a_space[index[0]])
pose = RobotPose()
pose.x = -vec[0] * self.resolution_factor * self.scale_factors[0] # [m]
pose.y = -vec[1] * self.resolution_factor * self.scale_factors[1] # [m]
pose.a = -self.a_space[index[0]] # [rad]
return pose
def draw_line(
image: OrthographicImage,
action_pose: Affine,
pt1: Affine,
pt2: Affine,
color,
thickness=1,
reference_pose=None,
) -> None:
action_pose = Frames.get_pose_in_image(action_pose=action_pose, image_pose=Affine(image.pose), reference_pose=reference_pose)
pt1_projection = tuple(image.project((action_pose * pt1).to_array()))
pt2_projection = tuple(image.project((action_pose * pt2).to_array()))
cv2.line(image.mat, pt1_projection, pt2_projection, color, thickness, lineType=cv2.LINE_AA)
def image_difference(image: OrthographicImage, image_comp: OrthographicImage) -> OrthographicImage:
kernel = np.ones((5, 5), np.uint8)
diff = np.zeros(image.mat.shape, np.uint8)
diff[(image.mat > image_comp.mat + 5) & (image.mat > 0)] = 255
diff = cv2.morphologyEx(diff, cv2.MORPH_OPEN, kernel, iterations=2)
return OrthographicImage(
diff,
image.pixel_size,
image.min_depth,
image.max_depth,
image.camera,
image.pose,
)
def get_area_of_interest_new(
image: OrthographicImage,
pose: Affine,
size_cropped: Tuple[float, float] = None,
size_result: Tuple[float, float] = None,
border_color=None,
project_3d=True,
flags=cv2.INTER_LINEAR,
) -> OrthographicImage:
size_input = (image.mat.shape[1], image.mat.shape[0])
center_image = (size_input[0] / 2, size_input[1] / 2)
action_pose = Frames.get_pose_in_image(action_pose=pose, image_pose=Affine(image.pose)) if project_3d else pose
angle_a = action_pose.a
if abs(action_pose.b) > np.pi - 0.1 and abs(action_pose.c) > np.pi - 0.1:
angle_a = action_pose.a - np.pi
if size_result and size_cropped:
scale = size_result[0] / size_cropped[0]
assert scale == (size_result[1] / size_cropped[1])
elif size_result:
scale = size_result[0] / size_input[0]
assert scale == (size_result[1] / size_input[1])
else:
scale = 1.0
size_final = size_result or size_cropped or size_input
trans = get_transformation(
image.pixel_size * action_pose.y,
image.pixel_size * action_pose.x,
-angle_a,
center_image,
scale=scale,
cropped=size_cropped,
)
mat_result = cv2.warpAffine(image.mat, trans, size_final, flags=flags, borderValue=border_color) # INTERPOLATION_METHOD
image_pose = Affine(x=action_pose.x, y=action_pose.y, a=-action_pose.a) * Affine(image.pose)
return OrthographicImage(
mat_result,
scale * image.pixel_size,
image.min_depth,
image.max_depth,
image.camera,
image_pose.to_array(),
)
def render(self, image: OrthographicImage, alpha=0.5, draw_directions=False):
prob = self.inf.model.predict(self.inf.get_images(image))
prob = np.maximum(prob, 0)
heat_values = self.calculate_heat(prob)
heatmap = cv2.applyColorMap(heat_values.astype(np.uint8), cv2.COLORMAP_JET)
result = cv2.cvtColor(image.mat, cv2.COLOR_GRAY2RGB) / 255 + alpha * cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB)
result = OrthographicImage(result, image.pixel_size, image.min_depth, image.max_depth)
if draw_directions:
for _ in range(10):
self.draw_arrow(result, prob, np.unravel_index(prob.argmax(), prob.shape))
prob[np.unravel_index(prob.argmax(), prob.shape)] = 0
return result.mat
def render(
self,
image: OrthographicImage,
goal_image: OrthographicImage = None,
alpha=0.5,
save_path=None,
reward_index=None,
draw_directions=False,
indices=None,
):
base = image.mat
inputs = [self.inf.get_images(image)]
if goal_image:
base = goal_image.mat
inputs += [self.inf.get_images(goal_image)]
reward = self.model.predict(inputs)
if reward_index is not None:
reward = reward[reward_index]
# reward = np.maximum(reward, 0)
reward_mean = np.mean(reward, axis=3)
# reward_mean = reward[:, :, :, 0]
heat_values = self.calculate_heat(reward_mean)
heatmap = cv2.applyColorMap(heat_values.astype(np.uint8),
cv2.COLORMAP_JET)
base_heatmap = cv2.cvtColor(base,
cv2.COLOR_GRAY2RGB) / 255 + alpha * heatmap
result = OrthographicImage(base_heatmap, image.pixel_size,
image.min_depth, image.max_depth)
if indices is not None:
self.draw_indices(result, reward, indices)
if draw_directions:
for _ in range(10):
self.draw_arrow(
result, reward,
np.unravel_index(reward.argmax(), reward.shape))
reward[np.unravel_index(reward.argmax(), reward.shape)] = 0
if save_path:
cv2.imwrite(str(save_path), result.mat)
return result.mat
def patch_image_at(
image: OrthographicImage,
patch: np.ndarray,
pose: Affine,
size_cropped=None,
operation='replace',
) -> OrthographicImage:
if size_cropped:
patch = cv2.resize(patch, size_cropped)
size_input = (image.mat.shape[1], image.mat.shape[0])
center_image = (size_input[0] / 2, size_input[1] / 2)
center_cropped_image = (patch.shape[1] / 2, patch.shape[0] / 2)
dx = center_image[0] - center_cropped_image[0] - image.pixel_size * pose.y
dy = center_image[1] - center_cropped_image[1] - image.pixel_size * pose.x
trans = get_transformation(
dx * np.cos(pose.a) - dy * np.sin(pose.a),
dy * np.cos(pose.a) + dx * np.sin(pose.a),
pose.a,
center_cropped_image
)
result = cv2.warpAffine(patch, trans, size_input, borderValue=np.iinfo(np.uint16).max)
mask = np.array(result < np.iinfo(np.uint16).max, dtype=np.uint16)
mask = cv2.erode(mask, np.ones((5, 5), np.uint16), iterations=1)
if operation == 'replace':
mat_patched_image = np.where(mask, result, image.mat)
elif operation == 'add':
mat_patched_image = image.mat + np.where(mask, result, np.zeros(image.mat.shape, dtype=np.uint16))
else:
raise Exception(f'Operation {operation} not implemented in patch image!')
return OrthographicImage(
mat_patched_image,
image.pixel_size,
image.min_depth,
image.max_depth,
image.camera,
image.pose,
)
def get_area_of_interest(
image: OrthographicImage,
pose: Affine,
size_cropped: Tuple[float, float] = None,
size_result: Tuple[float, float] = None,
border_color=None,
project_3d=True,
) -> OrthographicImage:
size_input = (image.mat.shape[1], image.mat.shape[0])
center_image = (size_input[0] / 2, size_input[1] / 2)
action_pose = Frames.get_pose_in_image(
action_pose=pose, image_pose=image.pose) if project_3d else pose
trans = get_transformation(image.pixel_size * action_pose.y,
image.pixel_size * action_pose.x,
-action_pose.a, center_image)
mat_result = cv2.warpAffine(image.mat,
trans,
size_input,
borderValue=border_color)
new_pixel_size = image.pixel_size
if size_cropped:
mat_result = crop(mat_result, size_output=size_cropped)
if size_result:
mat_result = cv2.resize(mat_result, size_result)
new_pixel_size *= size_result[0] / size_cropped[
0] if size_cropped else size_result[0] / size_input[0]
return OrthographicImage(
mat_result,
new_pixel_size,
image.min_depth,
image.max_depth,
image.camera,
Affine(x=action_pose.x, y=action_pose.y, a=-action_pose.a) *
image.pose,
)
def test_grasp_conversion(self):
conv = Converter(grasp_z_offset=0.0, box=self.box)
image = OrthographicImage(
self.read_image(self.file_path / self.image_name), 2000.0, 0.22,
0.4, '', Config.default_image_pose)
action = Action()
action.type = 'grasp'
action.pose.x = -0.06
action.pose.y = 0.099
action.pose.a = 0.523
action.pose.d = 0.078
action.index = 1
draw_pose(image, action.pose, convert_to_rgb=True)
draw_around_box(image, box=self.box, draw_lines=True)
imageio.imwrite(str(self.file_path / 'gen-grasp-draw-pose.png'),
image.mat)
self.assertTrue(conv.grasp_check_safety(action, [image]))
conv.calculate_pose(action, [image])
self.assertLess(action.pose.z, 0.0)
def get_images(self, image: OrthographicImage):
root = tk.Tk()
image_array = Image.fromarray(cv2.merge(cv2.split(image.mat)))
image_gtk = ImageTk.PhotoImage(image=image_array)
load_new_image = False
ann = [-1000, -1000]
def mouse_clicked(event):
x, y = event.x, event.y
ann[0] = x - self.size_input[0] / 2
ann[1] = y - self.size_input[1] / 2
image_show = Image.fromarray(cv2.merge(cv2.split(image.mat)))
cv2.circle(image_show, (x, y), 2, (255, 255, 255), -1)
image_array = Image.fromarray(cv2.merge(cv2.split(image_show)))
image_gtk = ImageTk.PhotoImage(image=image_array)
panel.configure(image=image_gtk)
panel.image = image_gtk
def button_submit():
root.destroy()
def new_image():
load_new_image = True
root.destroy()
panel = tk.Label(root, image=image_gtk)
button = tk.Button(root, text='Grasp Object', command=button_submit)
button2 = tk.Button(root, text='New Image', command=new_image)
panel.pack(side='top')
button.pack(side='bottom')
button2.pack(side='bottom')
panel.bind('<Button 1>', mouse_clicked)
root.mainloop()
if load_new_image:
return Action()
draw_around_box(image, box=self.box)
background_color = image.value_from_depth(-(image.pose * Affine(0, 0, self.box['size'][2])).z)
image_resized = cv2.resize(image, self.size_resized)
images = []
anns = []
for a in self.a_space:
rot_mat = cv2.getRotationMatrix2D(
(self.size_resized[0] / 2, self.size_resized[1] / 2),
a * 180.0 / np.pi,
1.0,
)
rot_mat[:, 2] += [
(self.size_rotated[0] - self.size_resized[0]) / 2,
(self.size_rotated[1] - self.size_resized[1]) / 2,
]
dst_depth = cv2.warpAffine(image_resized, rot_mat, self.size_rotated, borderValue=background_color)
images.append(crop(dst_depth, self.size_cropped))
ann_scaled = (float(ann[0]) / self.scale_factors[0], float(ann[1]) / self.scale_factors[1])
ann_scaled_rot = self.rotate_vector(ann_scaled, a)
ann_scaled_rot_pos = (
int(round(ann_scaled_rot[0] + self.size_cropped[0] / 2)),
int(round(ann_scaled_rot[1] + self.size_cropped[1] / 2)),
)
x_vec = np.zeros(self.size_cropped[0])
y_vec = np.zeros(self.size_cropped[1])
if 0 <= ann_scaled_rot_pos[0] < self.size_cropped[0]:
np.put(x_vec, ann_scaled_rot_pos[0], 1)
if 0 <= ann_scaled_rot_pos[1] < self.size_cropped[1]:
np.put(y_vec, ann_scaled_rot_pos[1], 1)
ann_matrix = 255 * np.outer(y_vec, x_vec)
anns.append(ann_matrix)
return images, anns
def predict_actions(
self,
images: List[OrthographicImage],
method: SelectionMethod,
N=1,
verbose=True,
) -> List[Action]:
actions: List[Action] = []
uncertainty_images = [
OrthographicImage(np.zeros(i.mat.shape,
dtype=np.uint16), i.pixel_size,
i.min_depth, i.max_depth, i.camera, i.pose)
for i in images
]
for i in range(N):
for image in images:
draw_around_box(image, box=Config.box)
grasp = self.grasp_inference.infer(
images, method, uncertainty_images=uncertainty_images)
self.grasp_indexer.to_action(grasp)
# Shift actions
if Config.shift_objects and grasp.estimated_reward < Config.grasp_shift_threshold:
shift = self.shift_inference.infer(images, method)
self.shift_indexer.to_action(shift)
bin_empty = shift.estimated_reward < Config.shift_empty_threshold
if bin_empty:
actions.append(Action('bin_empty', safe=1))
else:
self.converter.calculate_pose(shift, images)
actions.append(shift)
# Grasp actions
else:
estimated_reward_lower_than_threshold = grasp.estimated_reward < Config.bin_empty_at_max_probability
bin_empty = estimated_reward_lower_than_threshold and Epoch.selection_method_should_be_high(
method)
new_image = False
if bin_empty:
actions.append(Action('bin_empty', safe=1))
elif grasp.estimated_reward_std > 0.9: # default=0.25
actions.append(Action('new_image', safe=1))
else:
self.converter.calculate_pose(grasp, images)
actions.append(grasp)
actions[-1].step = i
action = actions[-1]
logger.info(f'{i}: {action}')
if verbose:
image_copy = clone(images[0])
uncertainty_image_copy = clone(uncertainty_images[0])
draw_pose(image_copy, action.pose, convert_to_rgb=True)
draw_pose(uncertainty_image_copy,
action.pose,
convert_to_rgb=True)
cv2.imwrite(str(self.output_path / f'result-{i}.png'),
image_copy.mat)
cv2.imwrite(str(self.output_path / f'uncertainty-{i}.png'),
uncertainty_image_copy.mat)
if action.type == 'bin_empty' or action.type == 'new_image':
break
# Predict next image
reward = action.estimated_reward > Config.bin_empty_at_max_probability if action.type == 'grasp' else action.estimated_reward
action_type = self.grasp_shift_indexer.from_action(action)
images = self.predict_images_after_action(
images,
action,
reward=reward,
action_type=action_type,
uncertainty_images=uncertainty_images,
)
if isinstance(images, tuple):
images, uncertainty_images = images
else:
uncertainty_images = None
if verbose:
cv2.imwrite(str(self.output_path / f'result-{i+1}.png'),
images[0].mat)
cv2.imwrite(str(self.output_path / f'uncertainty-{i+1}.png'),
uncertainty_images[0].mat)
return actions
def plan_actions(
self,
images: List[OrthographicImage],
method: SelectionMethod,
depth=1,
leaves=1,
verbose=False,
) -> List[Action]:
uncertainty_images = [
OrthographicImage(np.zeros(i.mat.shape,
dtype=np.uint16), i.pixel_size,
i.min_depth, i.max_depth, i.camera, i.pose)
for i in images
]
tree = PlanningTree(images, uncertainty_images)
for node, i in tree.fill_nodes(leaves=leaves, depth=depth):
# Visited actions: node.actions
for image in node.images:
draw_around_box(image, box=Config.box)
grasp = self.grasp_inference.infer(
node.images,
method,
uncertainty_images=node.uncertainty_images)
self.grasp_indexer.to_action(grasp)
# Shift actions
if Config.shift_objects and grasp.estimated_reward < Config.grasp_shift_threshold:
shift = self.shift_inference.infer(node.images, method)
self.shift_indexer.to_action(shift)
bin_empty = shift.estimated_reward < Config.shift_empty_threshold
if bin_empty:
action = Action('bin_empty', safe=1)
else:
self.converter.calculate_pose(shift, node.images)
action = shift
# Grasp actions
else:
estimated_reward_lower_than_threshold = grasp.estimated_reward < Config.bin_empty_at_max_probability
bin_empty = estimated_reward_lower_than_threshold and Epoch.selection_method_should_be_high(
method)
new_image = False
if bin_empty:
action = Action('bin_empty', safe=1)
elif grasp.estimated_reward_std > 0.9: # default=0.25
action = Action('new_image', safe=1)
else:
self.converter.calculate_pose(grasp, node.images)
action = grasp
logger.info(f'{i}: {action}')
if verbose:
image_copy = clone(images[0])
uncertainty_image_copy = clone(uncertainty_images[0])
draw_pose(image_copy, action.pose, convert_to_rgb=True)
draw_pose(uncertainty_image_copy,
action.pose,
convert_to_rgb=True)
cv2.imwrite(str(self.output_path / f'result-{i}.png'),
image_copy.mat)
cv2.imwrite(str(self.output_path / f'uncertainty-{i}.png'),
uncertainty_image_copy.mat)
if action.type == 'bin_empty' or action.type == 'new_image':
break
# Predict next image
reward = action.estimated_reward > Config.bin_empty_at_max_probability if action.type == 'grasp' else action.estimated_reward
action_type = self.grasp_shift_indexer.from_action(action)
images = self.predict_images_after_action(
node.images,
action,
reward=reward,
action_type=action_type,
uncertainty_images=node.uncertainty_images,
)
if isinstance(images, tuple):
images, uncertainty_images = images
else:
uncertainty_images = None
node.add_action(action, images, uncertainty_images)
if verbose:
cv2.imwrite(str(self.output_path / f'result-{i+1}.png'),
node.images[0].mat)
cv2.imwrite(str(self.output_path / f'uncertainty-{i+1}.png'),
node.uncertainty_images[0].mat)
actions, max_reward, mean_reward = tree.get_actions_maximize_reward(
max_depth=depth)
print(
f'Max reward: {max_reward:0.3f}, Mean reward: {mean_reward:0.3f}, Length: {len(actions)}'
)
# actions, max_steps, mean_steps = tree.get_actions_minimize_steps()
return actions
