YOLOv5不用detect.py进行自定义模型预测

一般我们使用yolov5的时候，是通过train.py进行训练，通过detect.py进行预测。但是我们将yolov5部署到web应用时，detect.py只能返回标注好的图片，返回具体的物品名是比较麻烦的。能不能直接输入一张图片，返回物体的坐标、种类呢？其实是有方法的。

import torch

model = torch.hub.load('D:\yolov5-master', 'custom', 'best.pt', source='local')

这行代码可以使用本地训练的模型进行预测
其中第一个参数是yolov5的项目文件夹，第三个参数是训练好的模型

import torch
import pandas

model = torch.hub.load('D:\yolov5-master', 'custom', 'best.pt', source='local')

#读取摄像头
cap = cv2.VideoCapture()

# 逐帧读取视频  
while True:  
    ret, frame = cap.read()  
	
	
	results = model(frame)
	df = results.pandas().xyxy[0]

这里我们逐帧读取摄像头，然后将图片传入到模型，利用pandas对数据进行处理，可以看到模型返回的数据是这样的。

      xmin    ymin    xmax   ymax  confidence  class    name
 0  749.50   43.50  1148.0  704.5    0.874023      0  person
 1  433.50  433.50   517.5  714.5    0.687988     27     tie
 2  114.75  195.75  1095.0  708.0    0.624512      0  person
 3  986.00  304.00  1028.0  420.0    0.286865     27     tie

其中前四个元素是坐标，confidence是置信度，name是种类名。

这样我们就可以提取每个种类的坐标和名字了。

接下来我们尝试不借助detect.py，自己做一个标注后的图片

import torch
import pandas
import cv2

 
# Model
model = torch.hub.load('D:\yolov5-master', 'custom', 'best.pt', source='local')


#读取网络摄像头视频
video_path = "http://192.168.0.147:4747/video"
cap = cv2.VideoCapture(video_path)

# 逐帧读取视频  
while True:  
    ret, frame = cap.read()  
  
    # 如果正确读取帧，ret为True  
    if not ret:  
        print("Can't receive frame (stream end?). Exiting ...")  
        break  
  

	#将读取的图片传入模型
    results = model(frame)
    df = results.pandas().xyxy[0]
    
    #遍历所有种类元素的坐标和种类名
    for index,row in df.iterrows():
        xmin = int(row['xmin'])
        ymin = int(row['ymin'])
        xmax = int(row['xmax'])
        ymax = int(row['ymax'])
        name = row['name']
		#打上方框
        cv2.rectangle(frame,(xmin,ymin),(xmax,ymax),(255,0,0),3)
        
		#显示种类文字
        font = cv2.FONT_HERSHEY_SIMPLEX  
        font_scale = 1  # 字体大小
        color = (255, 255, 255)  
        cv2.putText(frame,name,(xmin,ymin),font,font_scale,color,2)
    #展示图片
    cv2.imshow('Frame', frame)
       

因为台式机没有摄像头，所以使用了网络摄像头。可以使用DroidCam软件进行实现。

然后逐帧导入模型，生成pandas数据。

遍历pandas数据，提取每个元素的坐标和名字，渲染到原图上。

可以简单看一下效果。

可见效果是很好的

使用自定义模型预测，我们可以简单将预测结构封装成一个json数据，可以很好的实现前后端分离

例如

results = model(img)
type = results.pandas().xyxy[0]['name']

res = {}


for i in type:
    print(i)
    if (i in res) == False:
        res[i] = 1
    else:
        res[i]+=1

print(res)

#返回值
#{'qingcai': 1, 'qiezi': 1, 'doufu': 1, 'baicai': 1, 'rou': 2, 'fan': 1}