def load_datasets(self):
"""
Loads the datasets.
"""
for root, dirs, files in os.walk(self.circle_templates_dir):
self.circle_templates = [im2tensor(cv2.imread(os.path.join(self.circle_templates_dir, x), 1), channels=3) for x in files if x.endswith('.jpg')]
break
for root, dirs, files in os.walk(self.gripper_templates_dir):
self.gripper_templates = [im2tensor(cv2.imread(os.path.join(self.gripper_templates_dir, x), 1), channels=3) for x in files if x.endswith('.jpg')]
break
def output(self, sess, im,channels):
"""
accepts batch of 3d images, converts to tensor
and returns four element list of controls
"""
im = inputdata.im2tensor(im,channels)
shape = np.shape(im)
im = np.reshape(im, (-1, shape[0], shape[1], shape[2]))
with sess.as_default():
return sess.run(self.y_out, feed_dict={self.x:im}) [0]
def deploy(self, path, im):
"""
accepts 3-channel image with pixel values from
0-255 and returns controls in four element list
"""
sess = self.load(var_path=path)
im = inputdata.im2tensor(im)
shape = np.shape(im)
im = np.reshape(im, (-1, shape[0], shape[1], shape[2]))
with sess.as_default():
return sess.run(self.y_out, feed_dict={self.x:im})
def deploy(self, path, im):
"""
accepts 3-channel image with pixel values from
0-255 and returns controls in four element list
"""
sess = self.load(var_path=path)
im = inputdata.im2tensor(im)
shape = np.shape(im)
im = np.reshape(im, (-1, shape[0], shape[1], shape[2]))
with sess.as_default():
return sess.run(self.y_out, feed_dict={self.x:im})
def class_dist(self,sess,im,channels=3, mask=True):
"""
accepts batch of 3d images, converts to tensor
and returns four element list of controls
"""
if mask:
im = inputdata.im2tensor(im,channels)
shape = np.shape(im)
# print shape
im = np.reshape(im, (-1, shape[0], shape[1], shape[2]))
with sess.as_default():
dists = sess.run(self.y_out, feed_dict={self.x:im}) [0]
return np.reshape(dists, [2,5])
def reward_function(self, im, state, dist=True, success=False):
"""
Reward function for the current state, given an input
image and state.
Parameters:
im: image
Input image of the state.
Assumes image is already filtered.
state: numpy array
Internal state of the gripper.
dist: boolean
Whether to reward for distance.
"""
im = im2tensor(im, channels=3)
gc_pos = self.get_pos(self.circle_templates, im, gripper=False)
gripper_pos = self.get_pos(self.gripper_templates, im, gripper=True)
bounding_box = self.compute_bounding_box(-state[0], gripper_pos, w=125.0, h=50.0)
point_label = self.compute_point_label(gc_pos, bounding_box)
grip_reward = int(point_label) * self.success_reward
# Debugging
print "gc_pos", gc_pos
print "gripper_pos", gripper_pos
print "bounding_box", bounding_box
print "point_label", point_label
print "success", success
print ""
cv2.circle(im, tuple(gc_pos.astype(int)), radius=1, color=0, thickness=3)
cv2.circle(im, tuple(gripper_pos.astype(int)), radius=1, color=0, thickness=3)
for i in range(len(bounding_box)-1):
cv2.line(im, tuple(bounding_box[i].astype(int)), tuple(bounding_box[i+1].astype(int)), 0)
cv2.line(im, tuple(bounding_box[0].astype(int)), tuple(bounding_box[3].astype(int)), 0)
cv2.imshow('figure', im)
cv2.waitKey(100)
if dist:
dist_reward = -la.norm(gc_pos - gripper_pos)
return dist_reward + grip_reward
else:
return grip_reward
def copy_processed_image(src, dst): im = plt.imread(src) processed = im2tensor(im, channels=3) plt.imsave(dst, processed)
