📚 课程目录第1章基于YOLO的工业零件计数系统

1. 基于YOLO的工业零件计数系统

任务目标

构建一个基于YOLO的工业零件计数系统，对传送带上的零件（以螺栓为例）进行实时检测与计数。系统读取视频流，逐帧执行YOLO推理，维护一个跨帧的零件ID集合，最终输出检测到的唯一零件总数。

环境与依赖

Python 3.8+
OpenCV（pip install opencv-python）
Ultralytics YOLO（pip install ultralytics）
预训练模型：yolov8n.pt（首次运行自动下载）

核心实现

系统采用YOLOv8检测模型，对每帧图像中的零件目标（类别ID 0，即person类，此处作为演示占位，实际工业场景需替换为自定义训练模型）进行检测。通过ByteTrack跟踪算法为每个零件分配唯一ID，使用集合去重后统计总数。

import cv2
from ultralytics import YOLO
from collections import defaultdict

class PartCounter:
    def __init__(self, model_path='yolov8n.pt', conf_thres=0.5):
        self.model = YOLO(model_path)
        self.conf_thres = conf_thres
        self.seen_ids = set()  # 记录所有出现过的跟踪ID
        self.track_history = defaultdict(list)  # 可选：用于可视化轨迹

    def process_frame(self, frame):
        """处理单帧图像，返回标注后的帧和当前帧新增的零件数"""
        results = self.model.track(frame, persist=True, conf=self.conf_thres, tracker='bytetrack.yaml')
        new_count = 0

        if results[0].boxes.id is not None:
            boxes = results[0].boxes.xyxy.cpu().numpy()
            track_ids = results[0].boxes.id.cpu().numpy().astype(int)
            confs = results[0].boxes.conf.cpu().numpy()

            for box, track_id, conf in zip(boxes, track_ids, confs):
                # 过滤：仅统计类别为0（演示用），实际可改为自定义零件类别
                if int(results[0].boxes.cls[track_ids.tolist().index(track_id)]) == 0:
                    if track_id not in self.seen_ids:
                        self.seen_ids.add(track_id)
                        new_count += 1
                    # 绘制检测框和ID
                    x1, y1, x2, y2 = map(int, box)
                    cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
                    cv2.putText(frame, f'ID:{track_id}', (x1, y1-10),
                                cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)

        # 在左上角显示累计总数
        cv2.putText(frame, f'Total: {len(self.seen_ids)}', (20, 40),
                    cv2.FONT_HERSHEY_SIMPLEX, 1.2, (0, 0, 255), 3)
        return frame, new_count

    def get_total_count(self):
        """返回累计检测到的唯一零件总数"""
        return len(self.seen_ids)

def main(video_source=0):
    """主函数：支持摄像头(0)或视频文件路径"""
    cap = cv2.VideoCapture(video_source)
    if not cap.isOpened():
        print(f"错误：无法打开视频源 {video_source}")
        return

    counter = PartCounter()
    print("开始计数，按 'q' 退出...")

    while True:
        ret, frame = cap.read()
        if not ret:
            break

        annotated_frame, _ = counter.process_frame(frame)
        cv2.imshow('YOLO Part Counter', annotated_frame)

        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    cap.release()
    cv2.destroyAllWindows()
    print(f"最终零件总数: {counter.get_total_count()}")

if __name__ == '__main__':
    # 使用摄像头：main(0)
    # 使用视频文件：main('conveyor_parts.mp4')
    main(0)

运行方式

将上述代码保存为 part_counter.py
准备视频源：连接USB摄像头（默认0），或准备一段传送带零件视频文件
运行命令：python part_counter.py

验证方法

功能验证：启动后，画面中每个被检测的零件会显示绿色框和唯一ID；左上角红色数字实时更新累计总数
计数正确性：让同一零件从画面中消失再重新出现，其ID保持不变，总数不会重复累加；不同零件进入画面时总数递增
通过标准：当画面中有N个不同零件依次通过时，最终输出的总数等于N；按q退出后控制台打印的总数与画面显示一致

注意事项

当前使用yolov8n.pt预训练模型，其类别0为person，仅用于演示跟踪计数逻辑。实际工业场景需使用自定义数据集训练专用零件检测模型，并修改类别过滤条件
若视频帧率较低或零件运动过快，可适当降低conf_thres阈值（如0.3）以提高召回率
ByteTrack跟踪器在零件被完全遮挡或离开画面较久后可能分配新ID，导致计数偏高。对高精度要求场景，可结合ROI区域和运动方向做进一步过滤