collect training images

src/features.py

+#!/usr/bin/python
+import moveit_commander
+import rospy
+import sys
+import cv2
+import numpy as np
+import uuid
+import pickle
+import os.path
+from tqdm import tqdm
+
+from baxter_agent import BaxterAgent
+from planning.scene import Scene
+from features_detection.background_remove import segment
+
+from sklearn import tree, svm
+
+config = {
+	'is_train': True,
+	'tree_dir': './saved_trees/',
+	'scene': 'scene description', #Should contain a description of the scene to be loaded
+	'learner': 'learner.pkl'
+}
+
+# def features_to_class(features):
+# 	col = features['color'] == 'BLUE'
+# 	sha = features['shape'] == 'CUBE'
+# 	siz = features['size'] == 'SMALL'
+# 	f = sum(v<<i for i, v in enumerate([col, sha, siz]))
+# 	return f
+
+def create_model(x, y):
+	learner = {}
+	for k in y[0].keys():
+		learner[k] = svm.LinearSVC()
+	return learner
+	# net = tflearn.input_data(shape=[None].extend(x.shape[1:]))
+	# net = tflearn.conv2d(net, 32, (3, 3), activation='relu')
+	# net = tflearn.max_pool_2d(net, (2,2))
+	# net = tflearn.conv2d(net, 32, (3, 3), activation='relu')
+	# net = tflearn.max_pool_2d(net, (2,2))
+	# net = tflearn.conv2d(net, 64, (3, 3), activation='relu')
+	# net = tflearn.max_pool_2d(net, (2,2))
+	# net = tflearn.flatten(net)
+	# net = tflearn.fully_connected(net, 64, activation='relu')
+	# net = tflearn.dropout(net, 0.5)
+	# net = tflearn.fully_connected(net, 1, activation='sigmoid')
+	# net = tflearn.regression(net, optimizer='adam', loss='categorical_crossentropy')
+
+	# return tflearn.DNN(net)
+
+def preprocess_sample(img):
+	img = cv2.resize(img, (0,0), fx=0.5, fy=0.5)
+	return segment(img)
+
+def preprocess_data(x, y):
+	x = np.array([preprocess_sample(img) for img in x])
+	y = np.array(y)
+	shuffled = range(0, x.shape[0])
+	np.random.shuffle(shuffled)
+
+	x = x[shuffled, :, :, :]
+	y = y[shuffled]
+	x = np.reshape(x, (x.shape[0], -1))
+	return x, y
+
+def train_model(learner, x, y):
+	for k in learner.keys():
+		y_k = map(lambda fs: fs[k], y)
+		learner[k].fit(x, y_k)
+
+def predict(learner, x):
+	x = preprocess_sample(x)
+	cv2.imwrite('{}.jpg'.format(uuid.uuid4()), x)
+	x = np.array(x).flatten()
+	features = {}
+	for k in learner.keys():
+		features[k] = learner[k].predict([x])
+	return features
+
+def add_samples(x, y, baxter, features, n=10, dmin=0., dmax=0.05):
+	# f = 1 if features['color'] == 'BLUE' else 0
+	shifts = np.random.uniform(dmin, dmax, (n, 3))
+	# start pose
+	baxter.left.move(baxter.left.top)
+	rospy.sleep(1)
+	pose = baxter.left.current_pose()
+	baxter.left.take_snapshot()
+	for shift in tqdm(shifts):
+		p = np.copy(pose)
+		p[:3] += shift
+		baxter.left.move(list(p))
+		img = baxter.left.take_snapshot()
+		# cv2.imwrite('{}.jpg'.format(uuid.uuid4()), img)
+		x.append(img)
+		y.append(features)
+	baxter.left.move(baxter.left.neutral)
+	return x, y
+
+
+def train(robot, scene=None):
+	x = []
+	y = []
+	# features = {'shape': 'CUBE', 'size': 'SMALL'}
+	inp = None
+	while inp != 'q':
+		features = {}
+		inp = None
+		while inp != 'q' and inp != 'd':
+			print('name:feature, d: done, q: quit')
+			inp = raw_input()
+			feature = inp.split(':')
+			if len(feature) == 2:
+				features[feature[0]] = feature[1].upper()
+
+		if scene is not None:
+			scene.set_item(features['color'].lower())
+		x, y = add_samples(x, y, robot.baxter, features)
+
+	x, y = preprocess_data(x, y)
+	learner = create_model(x, y)
+	train_model(learner, x, y)
+
+	print('done fitting')
+	with open('learner.pkl', 'wb') as f:
+		pickle.dump(learner, f, protocol=pickle.HIGHEST_PROTOCOL)
+	print('saved learner')
+	return learner
+
+def run(robot, learner, scene=None):
+	baxter = robot.baxter
+	rospy.sleep(1)
+	inp = None
+	while inp != 'q':
+		print('p: predict, q: quit')
+		if scene is not None:
+			print('[opt] write a color:')
+		inp = raw_input()
+		if scene is not None and inp != 'p' and inp != 'q':
+			scene.set_item(inp.lower())
+
+		baxter.left.move(baxter.left.top)
+		rospy.sleep(1)
+		img = baxter.left.take_snapshot()
+		p = predict(learner, img)
+		print(p)
+
+
+if __name__ == '__main__':
+	moveit_commander.roscpp_initialize(sys.argv)
+	rospy.init_node('garbage_bot', anonymous=True)
+
+	b = BaxterAgent(config)
+
+	scene = None
+	if len(sys.argv) > 1:
+		scene = Scene(b.baxter.get_planning_frame())
+		rospy.sleep(2)
+		scene.clear_scene()
+		scene.set_table()
+		scene.set_item()
+		rospy.sleep(2)
+
+	b.baxter.both.move(b.baxter.both.neutral)
+	rospy.sleep(1)
+	if os.path.exists(config['learner']):
+		with open(config['learner'], 'rb') as f:
+			learner = pickle.load(f)
+	else:
+		learner = train(b, scene)
+	run(b, learner, scene)