147 lines
3.9 KiB
Python
147 lines
3.9 KiB
Python
# python -m pip install flask
|
|
# export FLASK_APP=main.py
|
|
# flask run --without-threads
|
|
|
|
from flask import Flask, request
|
|
from tensorflow import keras
|
|
from image_detector.yolo3.yolo import YOLO
|
|
from image_detector.robot_detector import detect_robot
|
|
from LED_color_detector.led_detector import get_led_color
|
|
from image_detector.rat_detector import detect_rat
|
|
from config import img_base_path
|
|
|
|
app = Flask(__name__)
|
|
|
|
led_model = keras.models.load_model('./model_VGG16_LED_color_bigger_scene2/')
|
|
|
|
""" Automatic call while FLASK init """
|
|
yolo_model = YOLO()
|
|
|
|
|
|
# def deinit_yolo(yolo):
|
|
# yolo.close_session()
|
|
|
|
"""API_address/detectRobot1?img="""
|
|
@app.get("/detectRobot1")
|
|
def detectRobot1():
|
|
robot_led_colors = []
|
|
response_model = {}
|
|
|
|
# build path
|
|
image_path = img_base_path + "robot1/" + request.args['img']
|
|
detected_objects = detect_robot(model=yolo_model, img_path=image_path)
|
|
|
|
if not detected_objects:
|
|
return {
|
|
0: ["None"],
|
|
}, 200
|
|
|
|
for robot in detected_objects:
|
|
color = get_led_color(robot_image=robot[0], model=led_model)
|
|
robot_led_colors.append(color)
|
|
|
|
""" for emphatic robot """
|
|
### led color and robot positions will be sended to empathy model
|
|
|
|
""" for egoistic robot """
|
|
robot_id = 0
|
|
for robot in range(len(detected_objects)):
|
|
response_model[int(robot_id)] = [
|
|
robot_led_colors[robot],
|
|
int(detected_objects[robot][1][0]),
|
|
int(detected_objects[robot][1][1])
|
|
]
|
|
robot_id += 1
|
|
return response_model, 200
|
|
|
|
|
|
|
|
"""API_address/detectRobot2?img="""
|
|
@app.get("/detectRobot2")
|
|
def detectRobot2():
|
|
robot_led_colors = []
|
|
response_model = {}
|
|
|
|
# build path
|
|
image_path = img_base_path + "robot2/" + request.args['img']
|
|
detected_objects = detect_robot(model=yolo_model, img_path=image_path)
|
|
|
|
if not detected_objects:
|
|
return {
|
|
0: ["None"],
|
|
}, 200
|
|
|
|
for robot in detected_objects:
|
|
color = get_led_color(robot_image=robot[0], model=led_model)
|
|
robot_led_colors.append(color)
|
|
|
|
""" for emphatic robot """
|
|
### led color and robot positions will be sended to empathy model
|
|
|
|
""" for egoistic robot """
|
|
robot_id = 0
|
|
for robot in range(len(detected_objects)):
|
|
response_model[int(robot_id)] = [
|
|
robot_led_colors[robot],
|
|
int(detected_objects[robot][1][0]),
|
|
int(detected_objects[robot][1][1])
|
|
]
|
|
robot_id += 1
|
|
return response_model, 200
|
|
|
|
|
|
"""API_address/detectRat1?img="""
|
|
@app.get("/detectRat1")
|
|
def detectRat1():
|
|
response_model = {}
|
|
|
|
# build path
|
|
image_path = img_base_path + "robot1/" + request.args['img']
|
|
detected_objects = detect_rat(model=yolo_model, img_path=image_path)
|
|
|
|
if not detected_objects:
|
|
return {
|
|
0: ["None"],
|
|
}, 200
|
|
|
|
""" for emphatic robot """
|
|
### led color and robot positions will be sended to empathy model
|
|
|
|
""" for egoistic robot """
|
|
rat_id = 0
|
|
for rat in range(len(detected_objects)):
|
|
response_model[int(rat_id)] = [
|
|
int(detected_objects[rat][1][0]),
|
|
int(detected_objects[rat][1][1])
|
|
]
|
|
rat_id += 1
|
|
return response_model, 200
|
|
|
|
|
|
"""API_address/detectRat2?img="""
|
|
@app.get("/detectRat2")
|
|
def detectRat2():
|
|
response_model = {}
|
|
|
|
# build path
|
|
image_path = img_base_path + "robot2/" + request.args['img']
|
|
detected_objects = detect_rat(model=yolo_model, img_path=image_path)
|
|
|
|
if not detected_objects:
|
|
return {
|
|
0: ["None"],
|
|
}, 200
|
|
|
|
""" for emphatic robot """
|
|
### led color and robot positions will be sended to empathy model
|
|
|
|
""" for egoistic robot """
|
|
rat_id = 0
|
|
for rat in range(len(detected_objects)):
|
|
response_model[int(rat_id)] = [
|
|
int(detected_objects[rat][1][0]),
|
|
int(detected_objects[rat][1][1])
|
|
]
|
|
rat_id += 1
|
|
return response_model, 200
|