Commit f9cc9799 authored by Mathieu Reymond's avatar Mathieu Reymond
Browse files

new model, new dataset

parent 3cc2e179

Too many changes to show.

To preserve performance only 1000 of 1000+ files are displayed.
#Learning imports
from features_detection.recognizer import getFeatures
from decision_api.decision_tree_gr import DecisionTreeGR
from planning.baxter import Baxter
import detecting.model as dm
import cv2
#Baxter and ROS imports
import baxter_interface
......@@ -9,118 +9,122 @@ import rospy
import time
class BaxterAgent(object):
"""
Agent that manages the baxter robot learning process and deployment
"""
def __init__(self, config):
self.config = config
self.baxter = Baxter();
self.baxter.both.move(self.baxter.both.neutral)
self.is_train = config['is_train']
self.navs = {'left': baxter_interface.Navigator('left'),
'right': baxter_interface.Navigator('right'),
'torso_left': baxter_interface.Navigator('torso_left'),
'torso_right': baxter_interface.Navigator('torso_right')}
if self.is_train:
self.tree = DecisionTreeGR()
else :
self.tree = self.load_decision_tree()
def extract_features(self, img):
f = getFeatures(img)
return f
def _train_sort(self, arm_name, msg):
arm = getattr(self.baxter, arm_name)
if msg:
arm.move(arm.top)
rospy.sleep(2)
img = arm.take_snapshot()
try:
f = self.extract_features(img)
f['direction'] = arm_name.upper()
print(f)
self.tree.add_training_sample(f)
self.baxter.sort(arm_name)
except:
print('no features found')
arm.move(arm.neutral)
def _train_end(self, msg):
if msg:
self.is_train = False
try:
self.tree.train()
self.save_decision_tree()
except:
print('could not train features')
def train(self):
self.drawTrainingUI()
train_left = lambda m: self._train_sort('left', m)
train_right = lambda m: self._train_sort('right', m)
self.navs['left'].button0_changed.connect(train_left)
self.navs['right'].button0_changed.connect(train_right)
self.navs['left'].button1_changed.connect(self._train_end)
self.navs['right'].button1_changed.connect(self._train_end)
# blocking call until end button is clicked
while self.is_train:
pass
print('done training')
def _sort(self, msg):
if msg:
self.baxter.left.move(self.baxter.left.top)
img = self.baxter.left.take_snapshot()
try:
features = self.extract_features(img)
cv2.imwrite('tosort.jpg', img)
print(features)
direction = self.tree.classify(features).lower()
print(direction)
if direction != 'left':
self.baxter.left.move(self.baxter.left.neutral)
self.baxter.sort(direction)
except Exception as e:
print(e)
self.baxter.both.move(self.baxter.both.neutral)
def _run_end(self, m):
if m:
self.is_run = False
def run(self):
self.drawProductionUI()
self.is_run = True
self.navs['torso_left'].button0_changed.connect(self._sort)
self.navs['torso_right'].button0_changed.connect(self._sort)
self.navs['torso_left'].button1_changed.connect(self._run_end)
self.navs['torso_right'].button1_changed.connect(self._run_end)
# blocking call until end button is clicked
while self.is_run:
pass
print('done')
def load_decision_tree(self):
return self.tree.load(self.config['tree_dir'])
def save_decision_tree(self):
self.tree.save(self.config['tree_dir'])
def init_move(self):
self.baxter.both.move(self.baxter.both.neutral)
self.baxter.right.gripper.close()
rospy.sleep(3)
def drawTrainingUI(self):
print 'Garbagebot now training. Please choose a side to sort after putting down an object'
#xdisplay_image.send_image('path/to/initial/image')
def drawProductionUI(self):
#TODO#
pass
"""
Agent that manages the baxter robot learning process and deployment
"""
def __init__(self, config):
self.config = config
self.baxter = Baxter();
self.baxter.both.move(self.baxter.both.neutral)
self.is_train = config['is_train']
self.navs = {'left': baxter_interface.Navigator('left'),
'right': baxter_interface.Navigator('right'),
'torso_left': baxter_interface.Navigator('torso_left'),
'torso_right': baxter_interface.Navigator('torso_right')}
self.ir = dm.load_model(config['ir_path'], config['ir_x'], config['ir_y'])
if self.is_train:
self.tree = DecisionTreeGR(features=['color', 'shape'])
else :
self.tree = self.load_decision_tree()
def extract_features(self, img):
img = dm.preprocess_sample(img)
pred = self.ir.predict_label([img])[0]
return dm.out_to_features(pred)
def _train_sort(self, arm_name, msg):
arm = getattr(self.baxter, arm_name)
if msg:
arm.move(arm.top)
rospy.sleep(2)
img = arm.take_snapshot()
try:
f = self.extract_features(img)
f['direction'] = arm_name.upper()
print(f)
self.tree.add_training_sample(f)
self.baxter.sort(arm_name)
except:
print('no features found')
arm.move(arm.neutral)
def _train_end(self, msg):
if msg:
self.is_train = False
try:
self.tree.train()
self.save_decision_tree()
except Exception as e:
print(e)
print('could not train features')
def train(self):
self.drawTrainingUI()
train_left = lambda m: self._train_sort('left', m)
train_right = lambda m: self._train_sort('right', m)
self.navs['left'].button0_changed.connect(train_left)
self.navs['right'].button0_changed.connect(train_right)
self.navs['left'].button1_changed.connect(self._train_end)
self.navs['right'].button1_changed.connect(self._train_end)
# blocking call until end button is clicked
while self.is_train:
pass
print('done training')
def _sort(self, msg):
if msg:
self.baxter.left.move(self.baxter.left.top)
img = self.baxter.left.take_snapshot()
try:
features = self.extract_features(img)
cv2.imwrite('tosort.jpg', img)
print(features)
direction = self.tree.classify(features).lower()
print(direction)
if direction != 'left':
self.baxter.left.move(self.baxter.left.neutral)
self.baxter.sort(direction)
except Exception as e:
print(e)
self.baxter.both.move(self.baxter.both.neutral)
def _run_end(self, m):
if m:
self.is_run = False
def run(self):
self.drawProductionUI()
self.is_run = True
self.navs['torso_left'].button0_changed.connect(self._sort)
self.navs['torso_right'].button0_changed.connect(self._sort)
self.navs['torso_left'].button1_changed.connect(self._run_end)
self.navs['torso_right'].button1_changed.connect(self._run_end)
# blocking call until end button is clicked
while self.is_run:
pass
print('done')
def load_decision_tree(self):
return self.tree.load(self.config['tree_dir'])
def save_decision_tree(self):
self.tree.save(self.config['tree_dir'])
def init_move(self):
self.baxter.both.move(self.baxter.both.neutral)
self.baxter.right.gripper.close()
rospy.sleep(3)
def drawTrainingUI(self):
print 'Garbagebot now training. Please choose a side to sort after putting down an object'
#xdisplay_image.send_image('path/to/initial/image')
def drawProductionUI(self):
#TODO#
pass
......@@ -9,7 +9,10 @@ from planning.scene import Scene
config = {
'is_train': True,
'tree_dir': './saved_trees/',
'scene': 'scene description' #Should contain a description of the scene to be loaded
'scene': 'scene description', #Should contain a description of the scene to be loaded
'ir_path': '/home/mreymond/dev/ros/src/GarbageBot/src/detecting/checkpoint/model.m',
'ir_x':[None, 100, 110, 3],
'ir_y': 9
}
def add_item(scene, color, msg):
......
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