from ultralytics import YOLO
import numpy as np
import cv2
#设置参数
#1.输入源
source=0
#2.目标检查模型文件路径
model_path='models/best1.pt'
#3.需要检测在电子元器件序号列表##
classes=[0,1,2,3,4]
#4.颜色池
colors=[(255,255,0),
(0,255,0),
(0,0,255),(255,50,200),(0,0,60)]
#导入模型
model=YOLO('models/best1.pt')
#功能1实时计数功能函数
def realtime_count(classes=[0,1,2,3,4]):
# 初始化结果##
n_result={'C':0,'R':0,'LED':0,'A':0,'B':0}
#模型推理
results=model.predict(source=source,classes=classes,stream=True,show=False)
for result in results:
#类别字典
cls_names=result.names
#输入图片
img=result.orig_img
#boxes对象
boxes=result.boxes
#目标的坐标、置信度、类别序号等信息
data=boxes.data
#GPU张量-->Nnmpy矩阵
data=data.cpu().numpy()
# 展示检测结果图像
k=show_result(data, img,cls_names)
# 如果用户输入字母q,则退出本功能
if k==ord('q'):
break
else:
#获取各类电子元器件的数量
for i in classes:
#计算数量
n=len(data[data[:,-1]==i])
# 更新数量结果
n_result[cls_names[i]]=n
print('计数结果',n_result)
# 销毁窗口
cv2.destroyAllWindows()
#功能2:累计计数功能函数#
def accu_count(classes=[0,1,2,3,4]):
# 初始化结果##
n_result = {'C': 0, 'R': 0, 'LED': 0, 'A': 0, 'B': 0}
# 模型推理
results = model.predict(source=source, classes=classes, stream=True, show=False)
for result in results:
# 类别字典
cls_names = result.names
# 输入图片
img = result.orig_img
# boxes对象
boxes = result.boxes
# 目标的坐标、置信度、类别序号等信息
data = boxes.data
# GPU张量-->Nnmpy矩阵
data = data.cpu().numpy()
# 展示检测结果图像
k = show_result(data, img, cls_names)
# 如果用户输入字母q,则退出本功能
if k == ord('q'):
break
elif k==ord('a'):
# 获取各类电子元器件的数量
for i in classes:
# 计算数量
n = len(data[data[:, -1] == i])
# 更新数量结果
n_result[cls_names[i]] += n
print('计数结果', n_result)
# 销毁窗口
cv2.destroyAllWindows()
def show_result(data,img,cls_names):
for i in range(len(data)):
x1, y1 = data[i][0],data[i][1]
x2, y2= data[i][2], data[i][3]
#cls_index = data[:, -1]
cls=cls_names[data[i][-1]]
#name = cls_names[cls_index]
conf = str(data[i][-2])[:4]
#画框
cv2.rectangle(img,(int(x1),int(y1)),(int(x2),int(y2)),
color=colors[int(data[i][-1])],thickness=2)
#写文字
cv2.putText(img,cls+''+str(conf)[:4],(int(x1),int(y1)),2,1,
color=colors[int(data[i][-1])],thickness=2)
#展示图片
cv2.imshow('',img)
k=cv2.waitKey(1)
return k
# realtime_count(classes=[0,2])
# accu_count(classes=[1,3])
12-16
2373
![](https://csdnimg.cn/release/blogv2/dist/pc/img/readCountWhite.png)
03-05
2705
![](https://csdnimg.cn/release/blogv2/dist/pc/img/readCountWhite.png)
05-22
4933
![](https://csdnimg.cn/release/blogv2/dist/pc/img/readCountWhite.png)
12-10
5110
![](https://csdnimg.cn/release/blogv2/dist/pc/img/readCountWhite.png)