调用自己训练的pt模型并通过接口形式去访问—web精简版

前言

本篇文章主要是给大家分享，如何调用自己训练号的pt模型文件，并通过接口的方式将结果可视化显示到web端，下面是完整的代码，代码中有详细注释，有问题可私信

完整代码

# 引入相关包
from fastapi import FastAPI
import uvicorn
from fastapi.responses import StreamingResponse
import torch
import cv2

# 实例化fastapi
app = FastAPI()


# 创建测试模型的函数
def model_test():
	# 加载本地训练好的pt模型文件
    model = torch.hub.load('./', 'custom', 'best.pt', source='local')
    # 设置置信度
    model.conf = 0.3

	# 通过VideoCapture(id)打开摄像头，这里的id号需要自行修改，一般内置摄像头是0，外置是1、2...
    cap = cv2.VideoCapture(0)
    if not cap.isOpened():
        raise RuntimeError("VideoCapture打开失败.")
	
    try:
    	# 当摄像头成功打开后，不停的捕捉视频流
        while cap.isOpened():
        	# 读取视频帧
            ret, frame = cap.read()
            if not ret:
                break
			
			# 将视频帧传入model中进行识别
            results = model(frame)
            # 读取模型识别的结果
            boxs = results.pandas().xyxy[0].values
			
			# 遍历模型中的结果，并可视化识别框和识别结果，其中box[0]到box[3]是目标物左上和右下的像素坐标，box[4]是识别的准确率，box[6]是识别的目标物类型
            for box in boxs:
                cv2.rectangle(frame, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])), (0, 255, 0), 2)
                cv2.putText(frame, str(box[6])+" "+str(box[4]), (int(box[0]), int(box[1])), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)
			
			# 将视频流返回到浏览器
            ret, buffer = cv2.imencode('.jpg', frame)
            images = buffer.tobytes()
            yield (b'--frame\r\n'
                   b'Content-Type: image/jpeg\r\n\r\n' + images + b'\r\n')
    finally:
        cap.release()


@app.get("/cj")
def cj():
    return StreamingResponse(model_test(), media_type='multipart/x-mixed-replace; boundary=frame')


if __name__ == "__main__":
	# 这里端口可自定义，只要不与现有的运行端口冲突就可以
    uvicorn.run(app=app, host="127.0.0.1", port=8087, workers=1)