def __init__(self,features,user_name = None):
self.com = COM()
self.robot = hsrb_interface.Robot()
self.noise = 0.1
self.features = features#self.com.binary_image
self.count = 0
if(not user_name == None):
self.com.Options.setup(self.com.Options.root_dir,user_name)
self.omni_base = self.robot.get('omni_base')
self.whole_body = self.robot.get('whole_body')
self.gripper = self.robot.get('gripper')
self.com.go_to_initial_state(self.whole_body,self.gripper)
self.joystick = JoyStick_X(self.com)
self.policy = Policy(self.com,self.features)
self.com.change_grip(self.gripper)
def __init__(self):
'''
Initialization class for a Policy
Parameters
----------
yumi : An instianted yumi robot
com : The common class for the robot
cam : An open bincam class
debug : bool
A bool to indicate whether or not to display a training set point for
debuging.
'''
self.robot = hsrb_interface.Robot()
self.omni_base = self.robot.get('omni_base')
self.whole_body = self.robot.get('whole_body')
self.cam = RGBD()
self.com = COM()
#self.com.go_to_initial_state(self.whole_body)
self.count = 425
self.joystick = JoyStick_X(self.com)
self.true_count = 0
def __init__(self, user_name=None, inject_noise=False, noise_scale=1.0):
self.robot = hsrb_interface.Robot()
self.noise = 0.1
self.omni_base = self.robot.get('omni_base')
self.whole_body = self.robot.get('whole_body')
self.gripper = self.robot.get('gripper')
self.tl = TransformListener()
self.cam = RGBD()
time.sleep(5)
self.b_d = Bottle_Detect(self.cam.read_info_data())
self.start_recording = False
self.stop_recording = False
self.com = COM()
if (not user_name == None):
self.com.Options.setup(self.com.Options.root_dir, user_name)
#self.com.go_to_initial_state(self.whole_body,self.gripper)
#self.whole_body.move_to_joint_positions({'head_tilt_joint':-0.3})
self.joystick = JoyStick_X(self.com,
inject_noise=inject_noise,
noise_scale=noise_scale)
self.torque = Gripper_Torque()
self.joints = Joint_Positions()
def __init__(self):
'''
Initialization class for a Policy
Parameters
----------
yumi : An instianted yumi robot
com : The common class for the robot
cam : An open bincam class
debug : bool
A bool to indicate whether or not to display a training set point for
debuging.
'''
self.robot = hsrb_interface.Robot()
self.rgbd_map = RGBD2Map()
self.omni_base = self.robot.get('omni_base')
self.whole_body = self.robot.get('whole_body')
#PARAMETERS TO CHANGE
self.side = 'TOP'
self.r_count = 0
self.grasp_count = 0
self.success_count = 0
self.true_count = 0
self.grasp = True
self.r_count = self.get_rollout_number()
self.cam = RGBD()
self.com = COM()
self.joystick = JoyStick_X(self.com)
if cfg.USE_WEB_INTERFACE:
self.wl = Web_Labeler()
else:
self.wl = Python_Labeler(cam=self.cam)
self.com.go_to_initial_state(self.whole_body)
self.tt = TableTop()
self.tt.find_table(self.robot)
self.position_head()
self.br = tf.TransformBroadcaster()
self.tl = TransformListener()
#self.test_current_point()
time.sleep(4)
#thread.start_new_thread(self.ql.run,())
print "after thread"
class CardPicker():
def __init__(self):
'''
Initialization class for a Policy
Parameters
----------
yumi : An instianted yumi robot
com : The common class for the robot
cam : An open bincam class
debug : bool
A bool to indicate whether or not to display a training set point for
debuging.
'''
self.robot = hsrb_interface.Robot()
self.omni_base = self.robot.get('omni_base')
self.whole_body = self.robot.get('whole_body')
self.cam = RGBD()
self.com = COM()
#self.com.go_to_initial_state(self.whole_body)
self.count = 425
self.joystick = JoyStick_X(self.com)
self.true_count = 0
def collect_data(self):
while True:
c_img = self.cam.read_color_data()
d_img = self.cam.read_depth_data()
c_img, d_img = self.com.format_data(c_img, d_img)
cv2.imshow('video_feed', c_img)
cv2.waitKey(30)
cur_recording = self.joystick.get_record_actions_passive()
if (cur_recording[0] < -0.1 and self.true_count % 5 == 0):
print "PHOTO SNAPPED " + str(self.count)
cv2.imwrite(
cfg.IMAGE_PATH + '/frame_' + str(self.count) + '.png',
c_img)
self.count += 1
self.true_count += 1
def __init__(self):
'''
Initialization class for a Policy
Parameters
----------
yumi : An instianted yumi robot
com : The common class for the robot
cam : An open bincam class
debug : bool
A bool to indicate whether or not to display a training set point for
debuging.
'''
self.robot = hsrb_interface.Robot()
self.omni_base = self.robot.get('omni_base')
self.whole_body = self.robot.get('whole_body')
self.cam = RGBD()
self.com = COM()
#self.com.go_to_initial_state(self.whole_body)
self.br = tf.TransformBroadcaster()
self.tl = TransformListener()
self.gp = GraspPlanner()
self.detector = Detector()
self.joystick = JoyStick_X(self.com)
self.suction = Suction(self.gp, self.cam, self.com.Options)
#self.suction.stop()
#thread.start_new_thread(self.ql.run,())
print "after thread"
class CardPicker():
def __init__(self):
'''
Initialization class for a Policy
Parameters
----------
yumi : An instianted yumi robot
com : The common class for the robot
cam : An open bincam class
debug : bool
A bool to indicate whether or not to display a training set point for
debuging.
'''
self.robot = hsrb_interface.Robot()
self.whole_body = self.robot.get('whole_body')
#self.cam = RGBD()
self.com = COM()
self.com.go_to_initial_state(self.whole_body)
self.br = tf.TransformBroadcaster()
self.tl = TransformListener()
self.gp = GraspPlanner()
#self.detector = Detector()
self.joystick = JoyStick_X(self.com)
#self.suction = Suction(self.gp,self.cam,self.com.Options)
#self.suction.stop()
#thread.start_new_thread(self.ql.run,())
print "after thread"
def find_mean_depth(self, d_img):
'''
Evaluates the current policy and then executes the motion
specified in the the common class
'''
indx = np.nonzero(d_img)
mean = np.mean(d_img[indx])
return
def card_pick(self):
while True:
cur_recording = self.joystick.get_record_actions_passive()
self.broadcast_transform()
if (cur_recording[0] < -0.1):
self.go_to_centroid(self.whole_body)
#self.com.go_to_initial_state(self.whole_body)
def broadcast_transform(self):
try:
self.br.sendTransform(
(0.0, 0.0, -0.02),
tf.transformations.quaternion_from_euler(ai=-0.785,
aj=0.0,
ak=0.0),
rospy.Time.now(), 'transform_ar_marker', 'ar_marker/11')
except:
rospy.logerr('ar marker not found')
def go_to_centroid(self, whole_body):
whole_body.end_effector_frame = 'hand_l_finger_vacuum_frame'
nothing = True
#self.whole_body.move_to_neutral()
whole_body.move_end_effector_pose(geometry.pose(z=-0.02, ei=-0.785),
'ar_marker/11')
#whole_body.move_end_effector_by_line((0,0,1),0.02)
#self.start()
#whole_body.move_to_joint_positions({'arm_lift_joint':0.23})
def check_card_found(self):
# try:
transforms = self.tl.getFrameStrings()
cards = []
for transform in transforms:
print transform
if 'card' in transform:
print 'got here'
f_p = self.tl.lookupTransform('head_rgbd_sensor_rgb_frame',
transform, rospy.Time(0))
cards.append(transform)
return True, cards
# except:
return False, []
class FindObject():
def __init__(self,features,user_name = None):
self.com = COM()
self.robot = hsrb_interface.Robot()
self.noise = 0.1
self.features = features#self.com.binary_image
self.count = 0
if(not user_name == None):
self.com.Options.setup(self.com.Options.root_dir,user_name)
self.omni_base = self.robot.get('omni_base')
self.whole_body = self.robot.get('whole_body')
self.gripper = self.robot.get('gripper')
self.com.go_to_initial_state(self.whole_body,self.gripper)
self.joystick = JoyStick_X(self.com)
self.policy = Policy(self.com,self.features)
self.com.change_grip(self.gripper)
def run(self):
if(self.policy.cam.is_updated):
self.com.load_net()
self.policy.rollout()
self.com.clean_up()
self.policy.cam.is_updated = False
self.count += 1
def check_initial_state(self):
go = False
while not go:
img_rgb = self.policy.cam.read_color_data()
img_depth = self.policy.cam.read_depth_data()
state = self.com.format_data(img_rgb,img_depth)
cv2.imshow('initial_state', state[0] )
cv2.waitKey(30)
self.joystick.apply_control()
cur_recording = self.joystick.get_record_actions()
if(cur_recording[1] < -0.1):
print "BEGIN ROLLOUT"
go = True
def go_to_initial_state(self):
self.com.go_to_initial_state(self.whole_body,self.gripper)
def check_success(self):
self.com.clean_up()
self.b_d = Bottle_Detect()
img_rgb = self.policy.cam.read_color_data()
success = self.b_d.detect_bottle(img_rgb)
self.b_d.clean_up()
print "BOTTLE FOUND ",success
return success
def save_data(self,img_rgb,img_detect):
state = self.com.format_data(img_rgb,None)
state['found_object'] = img_detect
state['object_poses']
def check_marker(self):
try:
A = self.tl.lookupTransform('head_l_stereo_camera_frame','ar_marker/12', rospy.Time(0))
except:
rospy.logerr('trash not found')
return False
return True
def execute_grasp(self):
self.com.grip_open(self.gripper)
self.whole_body.end_effector_frame = 'hand_palm_link'
nothing = True
try:
self.whole_body.move_end_effector_pose(geometry.pose(y=0.15,z=-0.11, ek=-1.57),'ar_marker/9')
nothing = False
except:
rospy.logerr('mustard bottle found')
self.com.grip_squeeze(self.gripper)
self.whole_body.move_end_effector_pose(geometry.pose(z=-0.9),'hand_palm_link')
self.com.grip_open(self.gripper)
self.com.grip_squeeze(self.gripper)
class FindObject():
def __init__(self,features,user_name = None):
self.com = COM()
self.robot = hsrb_interface.Robot()
self.noise = 0.1
self.features = features#self.com.binary_image
self.count = 0
if(not user_name == None):
self.com.Options.setup(self.com.Options.root_dir,user_name)
self.omni_base = self.robot.get('omni_base')
self.whole_body = self.robot.get('whole_body')
self.gripper = self.robot.get('gripper')
#self.com.go_to_initial_state(self.whole_body,self.gripper)
self.joystick = JoyStick_X(self.com)
self.tl = TransformListener()
self.policy = Policy(self.com,self.features)
def run(self):
if(self.policy.cam.is_updated):
self.com.load_net()
self.policy.rollout()
self.com.clean_up()
self.policy.cam.is_updated = False
self.count += 1
def check_initial_state(self):
go = False
sticky = True
while sticky:
self.joystick.apply_control()
cur_recording = self.joystick.get_record_actions()
if(cur_recording[1] > -0.1):
print "BEGIN ROLLOUT"
sticky = False
while not go:
img_rgb = self.policy.cam.read_color_data()
img_depth = self.policy.cam.read_depth_data()
state = self.com.format_data(img_rgb,img_depth)
cv2.imshow('initial_state', state[0] )
cv2.waitKey(30)
self.joystick.apply_control()
cur_recording = self.joystick.get_record_actions()
if(cur_recording[1] < -0.1):
print "BEGIN ROLLOUT"
go = True
def go_to_initial_state(self):
self.com.go_to_initial_state(self.whole_body,self.gripper)
def clear_data(self):
self.trajectory = []
def mark_success(self,q):
state = {}
if(q == 'y'):
state['success'] = 1
else:
state['success'] = 0
self.trajectory.append(state)
def is_goal_objet(self):
try:
full_pose = self.tl.lookupTransform('head_l_stereo_camera_frame','ar_marker/9', rospy.Time(0))
trans = full_pose[0]
transforms = self.tl.getFrameStrings()
poses = []
for transform in transforms:
if 'bottle' in transform:
f_p = self.tl.lookupTransform('head_rgbd_sensor_link',transform, rospy.Time(0))
poses.append(f_p[0])
print 'augmented pose ',f_p
print 'ar_marker ', trans
for pose in poses:
if(LA.norm(pose[2]-0.1-trans[2])< 0.03):
return True
except:
rospy.logerr('AR MARKER NOT THERE')
return False
def check_success(self):
self.com.clean_up()
self.b_d = Bottle_Detect(self.policy.cam.read_info_data())
img_rgb = self.policy.cam.read_color_data()
img_depth = self.policy.cam.read_depth_data()
s_obj,img_detect,poses = self.b_d.detect_bottle(img_rgb,img_depth)
success = self.is_goal_objet()
self.b_d.clean_up()
self.process_data(img_rgb,img_detect,poses,success)
print "BOTTLE FOUND ",success
return success
def process_data(self,img_rgb,img_detect,object_poses,success):
img_rgb_cr,img_d = self.com.format_data(img_rgb,None)
state = {}
state['color_img'] = img_rgb_cr
state['found_object'] = img_detect
state['object_poses'] = object_poses
state['was_success'] = success
self.trajectory.append(state)
def save_data(self):
self.com.save_evaluation(self.trajectory)
def check_marker(self):
try:
A = self.tl.lookupTransform('head_l_stereo_camera_frame','ar_marker/9', rospy.Time(0))
except:
rospy.logerr('trash not found')
return False
return True
def execute_grasp(self):
self.com.grip_open(self.gripper)
self.whole_body.end_effector_frame = 'hand_palm_link'
nothing = True
try:
self.whole_body.move_end_effector_pose(geometry.pose(y=0.15,z=-0.09, ek=-1.57),'ar_marker/9')
except:
rospy.logerr('mustard bottle found')
self.com.grip_squeeze(self.gripper)
self.whole_body.move_end_effector_pose(geometry.pose(z=-0.9),'hand_palm_link')
self.com.grip_open(self.gripper)
self.com.grip_squeeze(self.gripper)
class BedMaker():
def __init__(self):
'''
Initialization class for a Policy
Parameters
----------
yumi : An instianted yumi robot
com : The common class for the robot
cam : An open bincam class
debug : bool
A bool to indicate whether or not to display a training set point for
debuging.
'''
self.robot = hsrb_interface.Robot()
self.rgbd_map = RGBD2Map()
self.omni_base = self.robot.get('omni_base')
self.whole_body = self.robot.get('whole_body')
#PARAMETERS TO CHANGE
self.side = 'TOP'
self.r_count = 0
self.grasp_count = 0
self.success_count = 0
self.true_count = 0
self.grasp = True
self.r_count = self.get_rollout_number()
self.cam = RGBD()
self.com = COM()
self.joystick = JoyStick_X(self.com)
if cfg.USE_WEB_INTERFACE:
self.wl = Web_Labeler()
else:
self.wl = Python_Labeler(cam=self.cam)
self.com.go_to_initial_state(self.whole_body)
self.tt = TableTop()
self.tt.find_table(self.robot)
self.position_head()
self.br = tf.TransformBroadcaster()
self.tl = TransformListener()
#self.test_current_point()
time.sleep(4)
#thread.start_new_thread(self.ql.run,())
print "after thread"
def get_rollout_number(self):
if self.side == "BOTTOM":
rollouts = glob.glob(cfg.FAST_PATH + 'b_grasp/*.png')
else:
rollouts = glob.glob(cfg.FAST_PATH + 't_grasp/*.png')
r_nums = []
for r in rollouts:
a = r[56:]
i = a.find('_')
r_num = int(a[:i])
r_nums.append(r_num)
return max(r_nums) + 1
def position_head(self):
if self.side == "TOP":
self.tt.move_to_pose(self.omni_base, 'right_down')
self.tt.move_to_pose(self.omni_base, 'right_up')
self.tt.move_to_pose(self.omni_base, 'top_mid')
self.whole_body.move_to_joint_positions({'head_tilt_joint': -0.8})
elif self.side == "BOTTOM":
self.tt.move_to_pose(self.omni_base, 'lower_start')
self.whole_body.move_to_joint_positions({'head_tilt_joint': -0.8})
def collect_data_bed(self):
while True:
c_img = self.cam.read_color_data()
d_img = self.cam.read_depth_data()
cv2.imshow('video_feed', c_img)
cv2.waitKey(30)
cur_recording = self.joystick.get_record_actions_passive()
if (cur_recording[0] < -0.1 and self.true_count % 20 == 0):
print "PHOTO SNAPPED "
self.save_image(c_img)
if self.grasp:
self.grasp_count += 1
self.grasp = False
else:
self.success_count += 1
self.grasp = True
if (cur_recording[1] < -0.1 and self.true_count % 20 == 0):
print "ROLLOUT DONE "
self.r_count += 1
self.grasp_count = 0
self.success_count = 0
self.grasp = True
self.true_count += 1
def save_image(self, c_img):
if self.side == "BOTTOM":
if self.grasp:
cv2.imwrite(
cfg.FAST_PATH + 'b_grasp/frame_' + str(self.r_count) +
'_' + str(self.grasp_count) + '.png', c_img)
else:
cv2.imwrite(
cfg.FAST_PATH + 'b_success/frame_' + str(self.r_count) +
'_' + str(self.success_count) + '.png', c_img)
else:
if self.grasp:
cv2.imwrite(
cfg.FAST_PATH + 't_grasp/frame_' + str(self.r_count) +
'_' + str(self.grasp_count) + '.png', c_img)
else:
cv2.imwrite(
cfg.FAST_PATH + 't_success/frame_' + str(self.r_count) +
'_' + str(self.success_count) + '.png', c_img)
class CardPicker():
def __init__(self):
'''
Initialization class for a Policy
Parameters
----------
yumi : An instianted yumi robot
com : The common class for the robot
cam : An open bincam class
debug : bool
A bool to indicate whether or not to display a training set point for
debuging.
'''
self.robot = hsrb_interface.Robot()
self.omni_base = self.robot.get('omni_base')
self.whole_body = self.robot.get('whole_body')
self.cam = RGBD()
self.com = COM()
#self.com.go_to_initial_state(self.whole_body)
self.br = tf.TransformBroadcaster()
self.tl = TransformListener()
self.gp = GraspPlanner()
self.detector = Detector()
self.joystick = JoyStick_X(self.com)
self.suction = Suction(self.gp, self.cam, self.com.Options)
#self.suction.stop()
#thread.start_new_thread(self.ql.run,())
print "after thread"
def find_mean_depth(self, d_img):
'''
Evaluates the current policy and then executes the motion
specified in the the common class
'''
indx = np.nonzero(d_img)
mean = np.mean(d_img[indx])
return
def find_card(self):
self.whole_body.move_to_neutral()
self.whole_body.move_to_joint_positions({
'arm_flex_joint': -1.57,
'head_tilt_joint': -.785
})
print self.check_card_found()[0]
i = 0
while True:
print "panning ", i
self.whole_body.move_to_joint_positions({'head_pan_joint': .30})
i += 1
if self.check_card_found()[0]:
break
print "found card!"
self.card_pick()
def card_pick(self):
while True:
c_img = self.cam.read_color_data()
c_img_c = np.copy(c_img)
cv2.imshow('debug_true', c_img_c)
cv2.waitKey(30)
d_img = self.cam.read_depth_data()
if (not c_img == None and not d_img == None):
c_img_cropped, d_img = self.com.format_data(
np.copy(c_img), d_img)
data = self.detector.numpy_detector(np.copy(c_img_cropped))
cur_recording = self.joystick.get_record_actions_passive()
self.suction.find_pick_region_net(data, c_img, d_img, c_img_c)
if (cur_recording[0] < -0.1):
card_found, cards = self.check_card_found()
if (card_found):
self.suction.execute_grasp(cards, self.whole_body)
self.com.go_to_initial_state(self.whole_body)
def check_card_found(self):
# try:
transforms = self.tl.getFrameStrings()
cards = []
for transform in transforms:
#print transform
if 'card' in transform:
print 'got here'
f_p = self.tl.lookupTransform('head_rgbd_sensor_rgb_frame',
transform, rospy.Time(0))
cards.append(transform)
return True, cards
# except:
return False, []
from il_ros_hsr.core.grasp_planner import GraspPlanner
from il_ros_hsr.p_pi.bed_making.com import Bed_COM as COM
import sys
sys.path.append('/home/autolab/Workspaces/michael_working/fast_grasp_detect/')
from skvideo.io import vwrite
from il_ros_hsr.core.joystick_X import JoyStick_X
if __name__ == '__main__':
robot = hsrb_interface.Robot()
videos_color = []
cam = RGBD()
jy = JoyStick_X()
while True:
time.sleep(0.05)
img = cam.read_color_data()
if (not img == None):
cv2.imshow('debug2', img)
cv2.waitKey(30)
img = img[:, :, (2, 1, 0)]
videos_color.append(img)
cur_recording = jy.get_record_actions_passive()
if (cur_recording[0] < -0.1):
print "SAVING VIDEO"
def __init__(self):
"""
For faster data collection where we manually simulate it.
We move with our hands. This will give us the large datasets we need.
Supports both grasping and success net data collection. If doing the
grasping, DON'T MAKE IT A SUCCESS CASE where the blanket is all the way
over the corner. That way we can use the images for both grasping and
as failure cases for the success net.
For the success net data collection, collect data at roughly a 5:1 ratio
of successes:failures, and make failures the borderline cases. Then we
borrow data from the grasping network to make it 5:5 or 1:1 for the actual
success net training process (use another script for forming the data).
We use the keys on the joystick to indicate the success/failure class.
"""
makedirs()
self.robot = robot = hsrb_interface.Robot()
self.rgbd_map = RGBD2Map()
self.omni_base = self.robot.get('omni_base')
self.whole_body = self.robot.get('whole_body')
self.cam = RGBD()
self.com = COM()
self.wl = Python_Labeler(cam=self.cam)
# ----------------------------------------------------------------------
# PARAMETERS TO CHANGE (well, really the 'side' and 'grasp' only).
# We choose a fixed side and collect data from there, no switching.
# Automatically saves based on `r_count` and counting the saved files.
# `self.grasp` remains FIXED in the code, so we're either only
# collecting grasp or only collecting success images.
# ----------------------------------------------------------------------
self.side = 'BOTTOM' # CHANGE AS NEEDED
self.grasp = False # CHANGE AS NEEDED
self.grasp_count = 0
self.success_count = 0
self.true_count = 0
self.r_count = self.get_rollout_number()
self.joystick = JoyStick_X(self.com)
print("NOTE: grasp={} (success={}), side: {}, rollout num: {}".format(
self.grasp, not self.grasp, self.side, self.r_count))
print("Press X for any SUCCESS (class 0), Y for FAILURES (class 1).")
# Set up initial state, table, etc.
self.com.go_to_initial_state(self.whole_body)
self.tt = TableTop()
# For now, a workaround. Ugly but it should do the job ...
#self.tt.find_table(robot)
self.tt.make_fake_ar()
self.tt.find_table_workaround(robot)
self.br = tf.TransformBroadcaster()
self.tl = TransformListener()
time.sleep(4)
# When we start, spin this so we can check the frames. Then un-comment,
# etc. It's the current hack we have to get around crummy AR marker detection.
#rospy.spin()
# THEN we position the head since that involves moving the _base_.
self.position_head()
class BedMaker():
def __init__(self):
"""
For faster data collection where we manually simulate it.
We move with our hands. This will give us the large datasets we need.
Supports both grasping and success net data collection. If doing the
grasping, DON'T MAKE IT A SUCCESS CASE where the blanket is all the way
over the corner. That way we can use the images for both grasping and
as failure cases for the success net.
For the success net data collection, collect data at roughly a 5:1 ratio
of successes:failures, and make failures the borderline cases. Then we
borrow data from the grasping network to make it 5:5 or 1:1 for the actual
success net training process (use another script for forming the data).
We use the keys on the joystick to indicate the success/failure class.
"""
makedirs()
self.robot = robot = hsrb_interface.Robot()
self.rgbd_map = RGBD2Map()
self.omni_base = self.robot.get('omni_base')
self.whole_body = self.robot.get('whole_body')
self.cam = RGBD()
self.com = COM()
self.wl = Python_Labeler(cam=self.cam)
# ----------------------------------------------------------------------
# PARAMETERS TO CHANGE (well, really the 'side' and 'grasp' only).
# We choose a fixed side and collect data from there, no switching.
# Automatically saves based on `r_count` and counting the saved files.
# `self.grasp` remains FIXED in the code, so we're either only
# collecting grasp or only collecting success images.
# ----------------------------------------------------------------------
self.side = 'BOTTOM' # CHANGE AS NEEDED
self.grasp = False # CHANGE AS NEEDED
self.grasp_count = 0
self.success_count = 0
self.true_count = 0
self.r_count = self.get_rollout_number()
self.joystick = JoyStick_X(self.com)
print("NOTE: grasp={} (success={}), side: {}, rollout num: {}".format(
self.grasp, not self.grasp, self.side, self.r_count))
print("Press X for any SUCCESS (class 0), Y for FAILURES (class 1).")
# Set up initial state, table, etc.
self.com.go_to_initial_state(self.whole_body)
self.tt = TableTop()
# For now, a workaround. Ugly but it should do the job ...
#self.tt.find_table(robot)
self.tt.make_fake_ar()
self.tt.find_table_workaround(robot)
self.br = tf.TransformBroadcaster()
self.tl = TransformListener()
time.sleep(4)
# When we start, spin this so we can check the frames. Then un-comment,
# etc. It's the current hack we have to get around crummy AR marker detection.
#rospy.spin()
# THEN we position the head since that involves moving the _base_.
self.position_head()
def get_rollout_number(self):
"""Had to modify this a bit from Michael's code. Test+see if it works.
For now, let's save based on how many `data.pkl` files we have in the
appropriate directory.
"""
if self.side == "BOTTOM":
nextdir = 'b_grasp'
if not self.grasp:
nextdir = 'b_success'
rollouts = sorted([
x for x in os.listdir(join(FAST_PATH, nextdir))
if 'data' in x and 'pkl' in x
])
else:
nextdir = 't_grasp'
if not self.grasp:
nextdir = 't_success'
rollouts = sorted([
x for x in os.listdir(join(FAST_PATH, nextdir))
if 'data' in x and 'pkl' in x
])
return len(rollouts)
def position_head(self):
"""Ah, I see, we can go straight to the top. Whew.
It's new code reflecting the different poses and HSR joints:
But, see important note below about commenting out four lines ...
"""
self.whole_body.move_to_go()
if self.side == "BOTTOM":
self.tt.move_to_pose(self.omni_base, 'lower_start_tmp')
else:
self.tt.move_to_pose(self.omni_base, 'right_down')
self.tt.move_to_pose(self.omni_base, 'right_up')
self.tt.move_to_pose(self.omni_base, 'top_mid_tmp')
# NOTE: If robot is already at the top, I don't want to adjust position.
# Thus, comment out the three lines and the preceding `else` statement.
self.whole_body.move_to_joint_positions(
{'arm_flex_joint': -np.pi / 16.0})
self.whole_body.move_to_joint_positions(
{'head_pan_joint': np.pi / 2.0})
self.whole_body.move_to_joint_positions({'arm_lift_joint': 0.120})
self.whole_body.move_to_joint_positions(
{'head_tilt_joint': -np.pi / 4.0}) # -np.pi/36.0})
def collect_data_grasp_only(self):
"""Collect data for the grasping network only, like H's method.
Actually, some of these images should likely be part of the success
network training, where the 'success=False' because I don't think I
collected data here that was considered a 'success=True' ...
"""
data = []
assert self.grasp
rc = str(self.r_count).zfill(3)
while True:
c_img = self.cam.read_color_data()
d_img = self.cam.read_depth_data()
# Continually show the camera image on screen and wait for user.
# Doing `k=cv2.waitKey(33)` and `print(k)` results in -1 except for
# when we press the joystick in a certain configuration.
cv2.imshow('video_feed', c_img)
cv2.waitKey(30)
# Here's what they mean. Y is top, and going counterclockwise:
# Y: [ 1, 0]
# X: [-1, 0] # this is what we want for image collection
# A: [ 0, 1]
# B: [ 0, -1] # use to terminate the rollout
#
# There is also a bit of a delay embedded, i.e., repeated clicking
# of `X` won't save images until some time has passed. Good! It is
# also necessary to press for a few milliseconds (and not just tap).
cur_recording = self.joystick.get_record_actions_passive()
if (cur_recording[0] < -0.1 and self.true_count % 20 == 0):
print(
"PHOTO SNAPPED (cur_recording: {})".format(cur_recording))
self.save_image(c_img, d_img)
self.grasp_count += 1
# Add to dictionary info we want, including target pose.
# Also add 'type' key since data augmentation code uses it.
pose = red_contour(c_img)
info = {
'c_img': c_img,
'd_img': d_img,
'pose': pose,
'type': grasp
}
data.append(info)
print(" image {}, pose: {}".format(len(data), pose))
# --------------------------------------------------------------
# We better save each time since we might get a failure to
# detect, thus losing some data. We overwrite existing saved
# files, which is fine since it's the current rollout `r_count`.
# Since we detect pose before this, if the pose isn't detected,
# we don't save. Good.
# --------------------------------------------------------------
if self.side == 'BOTTOM':
save_path = join(FAST_PATH, 'b_grasp',
'data_{}.pkl'.format(rc))
else:
save_path = join(FAST_PATH, 't_grasp',
'data_{}.pkl'.format(rc))
with open(save_path, 'w') as f:
pickle.dump(data, f)
# Kill the script and re-position HSR to get diversity in camera views.
if (cur_recording[1] < -0.1 and self.true_count % 20 == 0):
print("ROLLOUT DONE (cur_recording: {})".format(cur_recording))
print("Length is {}. See our saved pickle files.".format(
len(data)))
sys.exit()
# Necessary, otherwise we'd save 3-4 times per click.
self.true_count += 1
def collect_data_success_only(self):
"""Collect data for the success network.
Should be more emphasis on the success cases (not failures) because the
grasing network data can supplement the failures. Focus on _borderline_
failures in this method.
Recall that 0 = successful grasp, 1 = failed grasp.
SAVE AND EXIT FREQUENTLY, perhaps after every 15-20 images. It's easy to
make a mistake with the class label, so better to exit early often.
"""
data = []
assert not self.grasp
rc = str(self.r_count).zfill(3)
while True:
c_img = self.cam.read_color_data()
d_img = self.cam.read_depth_data()
cv2.imshow('video_feed', c_img)
cv2.waitKey(30)
cur_recording = self.joystick.get_record_actions_passive()
# Joystick controllers. Y is top, and going counterclockwise:
# Y: [ 1, 0] # FAILURE images (class index 1)
# X: [-1, 0] # SUCCESS images (class index 0)
# A: [ 0, 1]
# B: [ 0, -1] # terminate data collection
# ------------------------------------------------------------------
if (cur_recording[0] < -0.1
or cur_recording[0] > 0.1) and self.true_count % 20 == 0:
print(
"PHOTO SNAPPED (cur_recording: {})".format(cur_recording))
if cur_recording[0] < -0.1:
s_class = 0
elif cur_recording[0] > 0.1:
s_class = 1
else:
raise ValueError(cur_recording)
self.save_image(c_img, d_img, success_class=s_class)
self.success_count += 1
# Add to dictionary info we want, including the class.
info = {
'c_img': c_img,
'd_img': d_img,
'class': s_class,
'type': 'success'
}
data.append(info)
print(" image {}, class: {}".format(len(data), s_class))
if self.side == 'BOTTOM':
save_path = join(FAST_PATH, 'b_success',
'data_{}.pkl'.format(rc))
else:
save_path = join(FAST_PATH, 't_success',
'data_{}.pkl'.format(rc))
with open(save_path, 'w') as f:
pickle.dump(data, f)
# Kill the script and re-position HSR to get diversity in camera views.
if (cur_recording[1] < -0.1 and self.true_count % 20 == 0):
print("ROLLOUT DONE (cur_recording: {})".format(cur_recording))
print("Length is {}. See our saved pickle files.".format(
len(data)))
sys.exit()
# Necessary, otherwise we'd save 3-4 times per click.
self.true_count += 1
def save_image(self, c_img, d_img, success_class=None):
"""Save images. Don't forget to process depth images.
For now I'm using a tuned cutoff like 1400, at least to _visualize_.
NOTE: since the cutoff for turning depth images into black may change,
it would be better to save the original d_img in a dictionary. Don't use
cv2.imwrite() as I know from experience that it won't work as desired.
"""
rc = str(self.r_count).zfill(3)
f_rc_grasp = 'frame_{}_{}.png'.format(rc,
str(self.grasp_count).zfill(2))
f_rc_success = 'frame_{}_{}_class_{}.png'.format(
rc,
str(self.success_count).zfill(2), success_class)
if np.isnan(np.sum(d_img)):
cv2.patchNaNs(d_img, 0.0)
d_img = depth_to_net_dim(d_img, cutoff=1400) # for visualization only
if self.side == "BOTTOM":
if self.grasp:
pth1 = join(FAST_PATH, 'b_grasp', 'rgb_' + f_rc_grasp)
pth2 = join(FAST_PATH, 'b_grasp', 'depth_' + f_rc_grasp)
else:
pth1 = join(FAST_PATH, 'b_success', 'rgb_' + f_rc_success)
pth2 = join(FAST_PATH, 'b_success', 'depth_' + f_rc_success)
else:
if self.grasp:
pth1 = join(FAST_PATH, 't_grasp', 'rgb_' + f_rc_grasp)
pth2 = join(FAST_PATH, 't_grasp', 'depth_' + f_rc_grasp)
else:
pth1 = join(FAST_PATH, 't_success', 'rgb_' + f_rc_success)
pth2 = join(FAST_PATH, 't_success', 'depth_' + f_rc_success)
cv2.imwrite(pth1, c_img)
cv2.imwrite(pth2, d_img)
class Collect_Demos():
def __init__(self, user_name=None, inject_noise=False, noise_scale=1.0):
self.robot = hsrb_interface.Robot()
self.noise = 0.1
self.omni_base = self.robot.get('omni_base')
self.whole_body = self.robot.get('whole_body')
self.gripper = self.robot.get('gripper')
self.tl = TransformListener()
self.cam = RGBD()
time.sleep(5)
self.b_d = Bottle_Detect(self.cam.read_info_data())
self.start_recording = False
self.stop_recording = False
self.com = COM()
if (not user_name == None):
self.com.Options.setup(self.com.Options.root_dir, user_name)
#self.com.go_to_initial_state(self.whole_body,self.gripper)
#self.whole_body.move_to_joint_positions({'head_tilt_joint':-0.3})
self.joystick = JoyStick_X(self.com,
inject_noise=inject_noise,
noise_scale=noise_scale)
self.torque = Gripper_Torque()
self.joints = Joint_Positions()
def proess_state(self):
img_rgb = self.cam.read_color_data()
img_depth = self.cam.read_depth_data()
cur_action, noise_action, time_pub = self.joystick.apply_control()
state = self.com.format_data(img_rgb, img_depth)
#cv2.imwrite('frame_'+str(self.count)+'.png',state[0])
#Save all data
data = {}
data['action'] = cur_action
data['color_img'] = state[0]
data['depth_img'] = state[1]
data['noise_action'] = noise_action
pose = self.whole_body.get_end_effector_pose().pos
pose = np.array([pose.x, pose.y, pose.z])
data['robot_pose'] = pose
# data['action_time'] = time_pub
# data['image_time'] = self.cam.color_time_stamped
print "ACTION AT COUNT ", self.count
print cur_action
self.count += 1
self.trajectory.append(data)
# if(LA.norm(cur_action) > 1e-3):
# print "GOT ACCEPTED"
# self.trajectory.append(data)
def check_success(self):
img_rgb = self.cam.read_color_data()
img_depth = self.cam.read_depth_data()
success = self.b_d.detect_bottle(img_rgb, img_depth)
print "BOTTLE FOUND ", success
return success
def run(self):
self.joystick.apply_control()
cur_recording = self.joystick.get_record_actions_passive()
if (cur_recording[0] < -0.1):
print "BEGIN DATA COLLECTION"
self.start_recording = True
count = 0
if (self.start_recording):
self.count = 0
self.trajectory = []
while not self.stop_recording:
#while count < 20:
if (self.cam.is_updated):
self.proess_state()
self.cam.is_updated = False
cur_recording = self.joystick.get_record_actions()
if (cur_recording[1] < -0.1):
print "END DATA COLLECTION"
self.stop_recording = True
count += 1
self.check_success()
q = input('SAVE DATA: y or n: ')
if (q == 'y'):
self.com.save_recording(self.trajectory)
self.start_recording = False
self.stop_recording = False