Commit 612625e0 authored by sleepless-se's avatar sleepless-se
Browse files

adjust format to ImageAI API format

parent 87986d46
Pipeline #68034105 passed with stage
in 12 minutes and 16 seconds
......@@ -48,23 +48,27 @@ def predict_image(image):
image = Image.open(image)
image = image.convert("RGB")
with graph.as_default():
result = yolo.detect_image(image)
results = yolo.detect_image(image)
import cv2
cv2.imwrite("out.jpg", np.asarray(result['image'])[..., ::-1])
return result
images = []
for result in results:
images = result['image']
cv2.imwrite("out.jpg", np.asarray(images)[..., ::-1])
return results
def return_format(data):
new_data = {}
boxs = []
names = []
for box in data['box_points']:
boxs.append([box[1],box[0],box[3],box[2]])
for name in data['name']:
names.append(classes[name])
new_data['box_points'] = boxs
new_data['percentage_probability'] = data['percentage_probability'] * 100
new_data['name'] = names
return new_data
results = []
for result in data:
box = result['box_points']
new_result = {}
new_result['box_points'] = [box[1],box[0],box[3],box[2]]
new_result['percentage_probability'] = result['percentage_probability'] * 100
new_result['name'] = result['name']
print(result)
results.append(new_result)
return new_result
app = Flask(__name__)
......
......@@ -124,10 +124,7 @@ class YOLO(object):
self.input_image_shape: [image.size[1], image.size[0]],
K.learning_phase(): 0
})
result = {}
result['name'] = out_classes
result['percentage_probability'] = out_scores
result['box_points'] = out_boxes
print('Found {} boxes for {}'.format(len(out_boxes), 'img'))
......@@ -136,11 +133,17 @@ class YOLO(object):
size=np.floor(3e-2 * image.size[1] + 0.5).astype('int32'))
thickness = (image.size[0] + image.size[1]) // 300
results = []
for i, c in reversed(list(enumerate(out_classes))):
predicted_class = self.class_names[c]
box = out_boxes[i]
score = out_scores[i]
result = {}
result['name'] = predicted_class
result['percentage_probability'] = score
result['box_points'] = box
label = '{} {:.2f}'.format(predicted_class, score)
draw = ImageDraw.Draw(image)
label_size = draw.textsize(label, font)
......@@ -166,13 +169,14 @@ class YOLO(object):
[tuple(text_origin), tuple(text_origin + label_size)],
fill=self.colors[c])
draw.text(text_origin, label, fill=(0, 0, 0), font=font)
result['image'] = image
results.append(result)
del draw
end = timer()
print(end - start)
result['image'] = image
print(result)
return result
return results
def close_session(self):
self.sess.close()
......
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