Commit 218f90ad authored by fouad5's avatar fouad5
Browse files

features

parent 3df47a8d
#Learning imports
from features_detection.recognizer import getFeatures
from features_detection.linear_classifier import evaluate
from decision_api.decision_tree_gr import DecisionTreeGR
from planning.baxter import Baxter
import cv2
......@@ -26,12 +27,12 @@ class BaxterAgent(object):
'torso_right': baxter_interface.Navigator('torso_right')}
if self.is_train:
self.tree = DecisionTreeGR()
self.tree = DecisionTreeGR(features=['color', 'shape'])
else :
self.tree = self.load_decision_tree()
def extract_features(self, img):
f = getFeatures(img)
f = evaluate(img)
return f
def _train_sort(self, arm_name, msg):
......@@ -40,6 +41,7 @@ class BaxterAgent(object):
arm.move(arm.top)
rospy.sleep(2)
img = arm.take_snapshot()
cv2.imwrite('bluesquare.1.jpg', img)
try:
f = self.extract_features(img)
f['direction'] = arm_name.upper()
......
......@@ -78,7 +78,7 @@ def predict(learner, x):
features[k] = learner[k].predict([x])
return features
def add_samples(x, y, baxter, features, n=10, dmin=0., dmax=0.05):
def add_samples(x, y, baxter, features, n=5, dmin=0., dmax=0.05):
# f = 1 if features['color'] == 'BLUE' else 0
shifts = np.random.uniform(dmin, dmax, (n, 3))
# start pose
......@@ -89,9 +89,10 @@ def add_samples(x, y, baxter, features, n=10, dmin=0., dmax=0.05):
for shift in tqdm(shifts):
p = np.copy(pose)
p[:3] += shift
#baxter.left.move(baxter.left.top)
baxter.left.move(list(p))
img = baxter.left.take_snapshot()
# cv2.imwrite('{}.jpg'.format(uuid.uuid4()), img)
cv2.imwrite('{}.jpg'.format(uuid.uuid4()), img)
x.append(img)
y.append(features)
baxter.left.move(baxter.left.neutral)
......@@ -115,7 +116,8 @@ def train(robot, scene=None):
if scene is not None:
scene.set_item(features['color'].lower())
x, y = add_samples(x, y, robot.baxter, features)
if inp != 'q':
x, y = add_samples(x, y, robot.baxter, features)
x, y = preprocess_data(x, y)
learner = create_model(x, y)
......
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