def lego_demo(self):
self.rollout_data = []
self.get_new_grasp = True
if not DEBUG:
self.gm.position_head()
while True:
time.sleep(1) #making sure the robot is finished moving
c_img = self.cam.read_color_data()
d_img = self.cam.read_depth_data()
cv2.imwrite("debug_imgs/c_img.png", c_img)
print "\n new iteration"
if (not c_img == None and not d_img == None):
#label image
data = self.wl.label_image(c_img)
c_img = self.cam.read_color_data()
grasp_boxes = []
suction_boxes = []
singulate_boxes = []
for i in range(data['num_labels']):
bbox = data['objects'][i]['box']
classnum = data['objects'][i]['class']
classname = ['grasp', 'singulate', 'suction',
'quit'][classnum]
if classname == "grasp":
grasp_boxes.append(bbox)
elif classname == "suction":
suction_boxes.append(bbox)
elif classname == "singulate":
singulate_boxes.append(bbox)
elif classname == "quit":
self.ds.save_rollout()
self.gm.move_to_home()
self.ds.append_data_to_rollout(c_img, data)
main_mask = crop_img(c_img)
col_img = ColorImage(c_img)
workspace_img = col_img.mask_binary(main_mask)
grasps = []
viz_info = []
for i in range(len(grasp_boxes)):
bbox = grasp_boxes[i]
center_mass, direction = bbox_to_grasp(bbox, c_img, d_img)
viz_info.append([center_mass, direction])
pose, rot = self.gm.compute_grasp(center_mass, direction,
d_img)
grasps.append(
self.gripper.get_grasp_pose(pose[0],
pose[1],
pose[2],
rot,
c_img=workspace_img.data))
suctions = []
for i in range(len(suction_boxes)):
suctions.append("compute_suction?")
if len(grasps) > 0 or len(suctions) > 0:
cv2.imwrite(
"grasps.png",
visualize(workspace_img, [v[0] for v in viz_info],
[v[1] for v in viz_info]))
print "running grasps"
for grasp in grasps:
print "grasping", grasp
self.gm.execute_grasp(grasp)
print "running suctions"
for suction in suctions:
print "suctioning", suction
#execute suction
elif len(singulate_boxes) > 0:
print("singulating")
bbox = singulate_boxes[0]
obj_mask = bbox_to_mask(bbox, c_img)
start, end = find_singulation(col_img, main_mask, obj_mask)
display_singulation(start, end, workspace_img)
self.gm.singulate(start, end, c_img, d_img)
else:
print("cleared workspace")
self.whole_body.move_to_go()
self.gm.position_head()
for i in range(data['num_labels']):
bbox = data['objects'][i]['box']
classnum = data['objects'][i]['class']
classname = ['grasp', 'singulate', 'suction', 'quit'][classnum]
if classname == "grasp":
grasp_boxes.append(bbox)
elif classname == "suction":
suction_boxes.append(bbox)
elif classname == "singulate":
singulate_boxes.append(bbox)
elif classname == "quit":
break
main_mask = crop_img(c_img)
col_img = ColorImage(c_img)
workspace_img = col_img.mask_binary(main_mask)
grasps = []
viz_info = []
for i in range(len(grasp_boxes)):
bbox = grasp_boxes[i]
center_mass, direction = bbox_to_grasp(bbox, c_img, d_img)
viz_info.append([center_mass, direction])
suctions = []
for i in range(len(suction_boxes)):
suctions.append("compute_suction?")
if len(grasps) > 0 or len(suctions) > 0:
print("grasping")
