Projekt-PBR_Sztuczna_Empatia/main.py

# 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