app.py

from cmath import inf
from flask import request, Response
from flask import Flask
import numpy as np
from gevent import pywsgi
from yolov5.main import *
from flask import render_template
import base64
import json
app = Flask(__name__)
weights="./yolov5/best.pt"
device="cpu"
img_name = "./static/pre-img.png"
@app.route('/api',methods=['POST','GET'])
def api():
    if request.method == 'POST':
        print("api post ok")
        data = request.get_json()
        img_base64 = data['img_base64']
        #print(img_base64)
        img_file = base64.b64decode(img_base64)
        img = cv2.imdecode(np.fromstring(img_file, np.uint8), cv2.IMREAD_GRAYSCALE)
        #print(img)
        cv2.imwrite("./static/pre-img.png",img)
        infer = detect(weights,img_name,device)
        print("infer ok")
        print(infer.shape)
        d={}
        for i in range(infer.shape[0]):
            if int(infer[i][5])==0:
                infer_value = "oilcan"
            elif int(infer[i][5])==1:
                infer_value = "vessel"
            elif int(infer[i][5])==2:
                infer_value = "bridge"
            elif int(infer[i][5])==3:
                infer_value = "plane" 
            else:
                infer_value = "none"
            data = {
            "x1":int(infer[i][0]),
            "y1":int(infer[i][1]),
            "x2":int(infer[i][2]),
            "y2":int(infer[i][3]),
            "confidence":float(infer[i][4]),
            "infer":infer_value
            }
            d[i]=data
        ret = json.dumps(d)
        return ret
    else:
        data = {
            "x1":0,
            "y1":0,
            "x2":0,
            "y2":0,
            "confidence":0.0,
            "infer":"none",
        }
        ret = json.dumps(data)
        return ret


@app.route('/',methods=['POST','GET'])
def hello():
    pre_img="pre-img.jpg"
    infer_img="infer-img.jpg"
    post = False
    data = detect(weights,img_name,device)
    data1 = data.numpy()
    data2 = []
    for i in range(data1.shape[0]):
        if int(data1[i][5])==0:
            data2.append("oilcan")
        elif int(data1[i][5])==1:
            data2.append("vessel")
        elif int(data1[i][5])==2:
            data2.append("bridge")
        elif int(data1[i][5])==3:
            data2.append("plane") 
        else:
            data2.append ("none")

    if request.method == 'POST':
        post = True
        pre_img = request.files['upload']
        print("post ok")
        try:
            img = cv2.imread(pre_img)
            cv2.imwrite("./static/pre-img.png",img)
            infer = detect(weights,img_name,device)
            infer_img = "./static/infer-img.png"
            data1 = infer.numpy()
            data2 = []
            for i in range(data1.shape[0]):
                if int(data1[i][5])==0:
                    data2.append("oilcan")
                elif int(data1[i][5])==1:
                    data2.append("vessel")
                elif int(data1[i][5])==2:
                    data2.append("bridge")
                elif int(data1[i][5])==3:
                    data2.append("plane") 
                else:
                    data2.append("none")
            print(infer)
            print("infer ok")
        except:
            print("read error")
    elif request.method == 'GET':
        print("get ok")
        infer=False
    else:
        print("request error")
        infer=False

    return render_template('index.html',infer=post,data=data1,data2=data2)

server = pywsgi.WSGIServer(('0.0.0.0', 5000), app)
server.serve_forever()