YOLOv5构建安全帽检测和识别系统使用详解

git clone https://github.com/ultralytics/yolov5.git
cd yolov5
pip install -r requirements.txt

train: ../dataset/train/
val: ../dataset/val/
nc: 2
names: ['no_helmet', 'helmet']

python train.py --img 640 --batch 16 --epochs 100 --data data/helmet_detection.yaml --cfg models/yolov5s.yaml --weights yolov5s.pt

python detect.py --source ../dataset/val/ --weights runs/train/exp/weights/best.pt --conf 0.5

import cv2
from yolov5 import YOLOv5
def main():
    # 加载模型
    model = YOLOv5('runs/train/exp/weights/best.pt')
    # 打开摄像头
    cap = cv2.VideoCapture(0)
    while True:
        # 读取摄像头画面
        ret, frame = cap.read()
        # 检测安全帽
        results = model.detect(frame)
        # 在画面上绘制检测框
        for result in results:
            x, y, w, h, cls, conf = result
            cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2)
            cv2.putText(frame, f'{cls}: {conf:.2f}', (x, y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
        # 显示画面
        cv2.imshow('Helmet Detection', frame)
        # 按'q'键退出
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    cap.release()
    cv2.destroyAllWindows()
if __name__ == '__main__':
    main()