Commit 3df47a8d authored by Mathieu Reymond's avatar Mathieu Reymond
Browse files

collect training images

parent 95eed8ef
#!/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)
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment