yolov5 6.1 flask返回检测图片

server.py

from flask import Flask, request
import torch
import io
from io import BytesIO
import base64
import json
from PIL import Image
#加载模型,'ultralytics/yolov5'为yolov5源码路径，下载存在本地。'custom'：通过这个设置可以加载自己的模型。path = 'yolov5s.pt':权重地址，source='local'：从本地加载yolov5源码，如果不设置默认为github，force_reload=True是否强制更新代码
model = torch.hub.load('ultralytics/yolov5','custom',path = 'yolov5s.pt',force_reload=True,source='local')

# 设置允许上传的文件格式
ALLOW_EXTENSIONS = ['png', 'jpg', 'jpeg', 'gif']

# 判断文件后缀是否存在
def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[-1] in ALLOW_EXTENSIONS
# 定义路由
app = Flask(__name__)
@app.route("/photo", methods=['POST','GET'])
def uploads():
    if not request.method == "POST":
        return
    if request.files.get("image"):
        image_file = request.files["image"]
        file_name = image_file.filename
        if allowed_file(file_name)==True:
            image_bytes = image_file.read()
            img = Image.open(io.BytesIO(image_bytes))
            if img.mode != 'RGB':
                img = img.convert('RGB')
            results = model(img, size=640)
            str = results.pandas().xyxy[0].to_json(orient="records")
            im = request.values.to_dict()
            #判断是否返回图片，这里的True和False都是字符串，只有True返回，其他字符串都不返回图片
            value = im['download_image']
            if value == 'True':
                # results.imgs  # array of original images (as np array) passed to model for inference
                results.render()  # updates results.imgs with boxes and labels
                image_str = []
                for img in results.ims:
                    buffered = BytesIO()
                    img_base64 = Image.fromarray(img)
                    img_base64.save(buffered, format="JPEG")
                    image = base64.b64encode(buffered.getvalue()).decode('utf-8')  # base64 encoded image with results
                    image_str.append(image)
                s = {'result': str, 'image': image_str[0],"name": file_name}
            else:
                s = {'result': str,"name": file_name}
        else:
            s = {file_name: "格式错误，无法正常打开",'error': 'error',}
            s = json.dumps(s)
        return s
if __name__ == "__main__":
    app.run(host='127.0.0.1',port=5000)

client.py

import requests
import os
import time
import argparse
from PIL import Image
from io import BytesIO
import base64

#加载图片地址
#图片文件夹，可以修改为自己的路径
path = r'C:\Users\Administrator\Desktop\test'
#连接地址，服务部署的地址，默认本机ip，端口5000
url = "http://127.0.0.1:5000/photo"
file_path = os.listdir(path)
s = len(file_path)
t1 = time.time()
for i in file_path:
    img_path = os.path.join(path+'/'+i)
    path1 = img_path.replace("\\","/")
    #路径拼接
    file_name = img_path.split('/')[-1]
    #图片名分割
# 二进制打开图片
    file = open(img_path, 'rb')
# # 拼接参数，False为不返回图片，True为返回图片
    data = {'download_image': False}
    files = {'image': file}
# 发送post请求到服务器端
    r = requests.post(url, files=files,data=data)
    file.close()
# # 获取服务器返回的图片，字节流返回
    result = r.json()
    try:
        #对图片进行解码并打开
        img = Image.open(BytesIO(base64.b64decode(result['image'])))
        #结果删除图片编码进行打印
        result.pop('image')
        print(result)
        #储存图片，在本项目目录，需要建立一个名字为1的文件夹
        img.save('1/'+'result'+file_name)
        # img.show()#展示图片
        img.close()
    except KeyError:
        result['path'] = path1
        print(result)
t2 = time.time()
t3 = t2 -t1
print('耗费时间',t3)
print('图片总数',s)
print('平均时间',t3/s)