yolov5的flask部署python调用

yolov5 github：https://github.com/ultralytics/yolov5
跟踪：https://github.com/mikel-brostrom/Yolov5_DeepSort_Pytorch
TensorRT：https://github.com/TrojanXu/yolov5-tensorrt
NCNN：https://github.com/WZTENG/YOLOv5_NCNN

detect:

from torchvision import transforms
import torch
from PIL import Image,ImageDraw
from models import yolo
from utils.general import non_max_suppression
from models.experimental import attempt_load

# model = yolo.Model(r"D:\GoogleEarthProPortable\yolov5-master\models\yolov5s.yaml")
# model.load_state_dict(torch.load(r"D:\GoogleEarthProPortable\yolov5-master\weights\yolov5s.pt"))
model = attempt_load("weights/yolov5s.pt")  # load FP32 model
model.eval()

img = Image.open("inference/images/bus.jpg")

tf = transforms.Compose([
    transforms.Resize((512,640)),
    transforms.ToTensor()
])

print(img.size) # w,h
scale_w = img.size[0] /640
scale_h = img.size[1] /512
im = img.resize((640,512))

img_tensor = tf(img)

pred = model(img_tensor[None])[0]
pred = non_max_suppression(pred,0.3,0.5)

imgDraw = ImageDraw.Draw(img)
for box in pred[0]:
    b = box.cpu().detach().long().numpy()
    print(b)
    imgDraw.rectangle((b[0]*scale_w,b[1]*scale_h,b[2]*scale_w,b[3]*scale_h))
    # imgDraw.rectangle((b[0],b[1],b[2],b[3]))
img.show()

serving：

import io
import json

from torchvision import models
import torchvision.transforms as transforms
from PIL import Image,ImageDraw

from utils.general import non_max_suppression
from models.experimental import attempt_load

from flask import Flask, jsonify, request
app = Flask(__name__)

model = attempt_load("weights/yolov5s.pt")  # load FP32 model
model.eval()

names= ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
        'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
        'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
        'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
        'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
        'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
        'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
        'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
        'hair drier', 'toothbrush']

def transform_image(image_bytes):
    my_transforms = transforms.Compose([transforms.Resize((512,640)),
                                        transforms.ToTensor(),
                                        ])
    image = Image.open(io.BytesIO(image_bytes))
    return my_transforms(image)

def get_prediction(image_bytes):
    tensor = transform_image(image_bytes=image_bytes)
    outputs = model(tensor[None])[0]
    print(outputs)
    outputs = non_max_suppression(outputs,0.3,0.5)
    boxs = outputs[0]
    print(boxs[0])
    print(int(boxs[0][-1].item()))
    class_name = names[int(boxs[0][5].item())]
    print(boxs.shape)
    boxes = []
    for i in range(boxs.shape[0]):
        boxes.append([boxs[i][0].item(),boxs[i][1].item(),boxs[i][2].item(),boxs[i][3].item(),boxs[i][4].item(),boxs[i][5].item()])

    return boxes

@app.route('/predict', methods=['POST'])
def predict():
    if request.method == 'POST':
        file = request.files['file']
        img_bytes = file.read()
        boxes = get_prediction(image_bytes=img_bytes)
        return ({'boxes': boxes})


if __name__ == '__main__':
    app.run()

client：

import requests
import os

for i in os.listdir("inference/images"):
    image = open("inference/images/"+i,'rb')
    payload = {'file':image}
    r = requests.post(" http://localhost:5000/predict", files=payload).json()
    print(r)

git bash控制台：
启动flask服务器：FLASK_ENV=development FLASK_APP=app.py flask run
测试命令：curl -X POST -F file=@test_img/dog.jpg http://localhost:5000/predict