Commit b28a1fd1 authored by fouad5's avatar fouad5
Browse files

final setup (works)

parent dced312d
......@@ -33,9 +33,7 @@ class BaxterAgent(object):
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)
return dm.predict_features(self.ir, img)
def _train_sort(self, arm_name, msg):
arm = getattr(self.baxter, arm_name)
......
......@@ -10,9 +10,9 @@ config = {
'is_train': True,
'tree_dir': './saved_trees/',
'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_path': '/home/ubuntu/Desktop/garbagebot/ros/src/GarbageBot/src/detecting/checkpoint/model.m',
'ir_x':[None, 100, 110, 3],
'ir_y': 9
'ir_y': 3
}
def add_item(scene, color, msg):
......
......@@ -9,6 +9,44 @@ from imgaug import augmenters as iaa
from uuid import uuid4
from math import sqrt
def isolate_shape(img, rect=None, only_shape=True):
img = cv2.GaussianBlur(img, (5, 5), 0)
# plt.imshow(img)
# plt.show(block=False)
height, width = img.shape[:2]
# Create a mask holder
mask = np.zeros(img.shape[:2], np.uint8)
# Grab Cut the object
bgdModel = np.zeros((1, 65), np.float64)
fgdModel = np.zeros((1, 65), np.float64)
# "Hard Coding the Rect The object must lie within this rect."
if not rect:
rect = (300, 0, 100, height / 2)
cv2.grabCut(img, mask, rect, bgdModel, fgdModel, 2, cv2.GC_INIT_WITH_RECT)
# grabCut
mask = np.where((mask == 2) | (mask == 0), 0, 1).astype('uint8')
img1 = img * mask[:, :, np.newaxis]
# Get the background
background = img - img1
# Change all pixels in the background that are not black to white
background[np.where((background > [0, 0, 0]).all(axis=2))] = [
255, 255, 255
]
# Add the background and the image
final = background + img1
if only_shape:
return background
return final
def name_to_features(name):
color = {'RED': 0, 'BLUE': 1, 'GREEN': 2}
shape = {'CUBE': 0, 'SPHERE': 1, 'TRIANGLE': 2}
......@@ -30,10 +68,11 @@ def out_to_features(out):
shape = ['CUBE', 'SPHERE', 'TRIANGLE']
# outputs ranked according to probabilities,
# take the most probable one
print(out)
out = out[0]
c_i = out/3
s_i = out%3
# out = out[0]
# c_i = out/3
# s_i = out%3
c_i = np.argmax(out[3:])
s_i = np.argmax(out[:3])
return {'color': color[c_i], 'shape': shape[s_i]}
def load(path):
......@@ -71,6 +110,7 @@ def augment(x, y, nb=20):
return np.append(x, np.array(x_aug), axis=0), np.append(y, np.array(y_aug), axis=0)
def preprocess_sample(img, crop_x=(210,430), crop_y=(0, 200)):
img = isolate_shape(img, only_shape=False)
img = img[crop_y[0]:crop_y[1], crop_x[0]:crop_x[1], :]
img = cv2.resize(img, (0,0), fx=0.5, fy=0.5)
# img = cv2.GaussianBlur(img,(5,5),0)
......@@ -117,6 +157,13 @@ def train_model(model, x, y):
outs = [y[:,3:], y[:,:3]]
train_net(model, x, outs)
def predict_features(model, img):
# cv2.imwrite('tosort.jpg', img)
img = preprocess_sample(img)
pred = model.predict([img])[0]
print(pred)
return out_to_features(pred)
def load_model(path, x_shape, out):
model = create_model(x_shape, out)
model.load(path)
......@@ -145,7 +192,4 @@ if __name__ == '__main__':
else:
model = load_model(model_path, [None, 100, 110, 3], 3)
img = cv2.imread(pred_path)
img = preprocess_sample(img)
print(model.predict([img]))
pred = model.predict_label([img])[0]
print(out_to_features(pred))
print(predict_features(model, img))
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