基于YOLOv8模型实现电子元器件实时计数和累计计数

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])