class Collect_Demos():
def __init__(self):
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.start_recording = False
self.stop_recording = False
self.com = COM()
self.com.go_to_initial_state(self.whole_body, self.gripper)
#self.whole_body.move_to_joint_positions({'head_tilt_joint':-0.3})
self.cam = RGBD()
self.joystick = JoyStick()
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 = self.joystick.get_current_control()
state = self.com.format_data(img_rgb, img_depth)
cv2.imshow('debug', state[0])
cv2.waitKey(30)
#Save all data
data = {}
data['action'] = cur_action
data['color_img'] = state[0]
data['depth_img'] = state[1]
data['noisey_twist'] = self.joystick.get_current_twist()
data['gripper_torque'] = self.torque.read_data()
data['joint_positions'] = self.joints.read_data()
if (LA.norm(cur_action) > 1e-4):
self.trajectory.append(data)
def run(self):
cur_recording = self.joystick.get_record_actions()
if (cur_recording[0] < -0.1):
print "BEGIN DATA COLLECTION"
self.start_recording = True
count = 0
if (self.start_recording):
self.trajectory = []
while not self.stop_recording:
#while count < 20:
if (count % 6 == 0):
self.proess_state()
cur_recording = self.joystick.get_record_actions()
if (cur_recording[1] < -0.1):
print "END DATA COLLECTION"
self.stop_recording = True
count += 1
q = input('SAVE DATA: y or n: ')
if (q == 'y'):
self.com.save_recording(self.trajectory)
self.start_recording = False
self.stop_recording = 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)