fastapi封装yolov5自训练模型

简介：本次用到的模型为基于yolov5自训练的口罩检测识别模型，测试工具是postman，简单在次做一个知识分享，也是给自己留一份记忆，比较适合刚入门人工智能行业的小白，如有问题请在下方留言，多多探讨，如有错误，请指正 抱拳抱拳抱拳

1、必备资源

        github下载yolov5源码，下载yolov5s.pt模型文件

        使用yolov5s.pt模型或者自己基于yolov5源码训练的模型

        下载postman工具（fastapi进行模型测试时使用）

2、安装必备依赖

        其中unvicron是基于 uvloop 和 httptools 构建的非常快速的 ASGI 服务器

        pip install fastapi

        pip install unvicron

3、找到yolov5源码中的detect.py文件,然后修改其中的如图所示：

在run方法最后新增代码，返回推理完成后图片的路径

4、新建webinfer.py文件，来搭建fastapi框架进行模型的推理，代码如下：

import io
import cv2
from starlette.responses import StreamingResponse
from fastapi import FastAPI

from detect import run

app = FastAPI()


@app.post('/predict/')
def infer(source: str):
    image_path = run(weights='best.pt', source=source)
    # 读取保存的图片
    saved_image = cv2.imread(image_path)
    _, img_encoded = cv2.imencode('.jpeg', saved_image)

    return StreamingResponse(io.BytesIO(img_encoded.tobytes()), media_type="image/jpeg")


if __name__ == '__main__':
    import uvicorn

    uvicorn.run(app, host="127.0.0.1", port=8000)

5、执行webinfer.py文件，得到如图所示即为正常：

 6、打开postman工具，配置参数，根据代码中的信息，填写对应参数（source对应的是测试图片的路径）

 点击发送得到最后结果

 后端将会显示200 ok表示响应成功