使用YOLOv5进行实时人类检测——摄像头版本

在计算机视觉领域，实时人类检测是一个重要的应用场景。本文将介绍如何使用YOLOv5模型和OpenCV库来实现一个简单的实时人类检测程序。

环境准备

在开始之前，请确保您已经安装了以下库：

• OpenCV

• PyTorch

• NumPy

您可以使用以下命令安装这些库：

pip install opencv-python torch torchvision numpy

代码实现

以下是实现实时人类检测的完整代码：

import cv2
import torch
import numpy as np
import time

def detect_people():
    # 加载YOLOv5模型
    model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
    # 只检测人类（person类别）
    model.classes = [0]  # COCO数据集中人类的类别索引为0
    
    # 打开摄像头
    cap = cv2.VideoCapture(0)  # 使用外接摄像头，如果不行请改为0
    
    if not cap.isOpened():
        print("Error: 无法打开摄像头")
        return

    # 设置摄像头分辨率
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
    
    # 用于控制检测频率
    frame_count = 0
    detection_interval = 2  # 每隔2帧进行一次检测
    last_results = None
    
    print("按 'q' 退出程序")
    print("按 'esc' 退出程序")
    
    while True:
        start_time = time.time()
        ret, frame = cap.read()
        if not ret:
            print("Error: 无法读取摄像头画面")
            break
            
        # 每隔几帧进行一次检测
        if frame_count % detection_interval == 0:
            # 使用YOLOv5进行检测
            results = model(frame)
            last_results = results
        else:
            results = last_results
            
        # 获取检测结果
        if results is not None:
            # 获取检测框
            boxes = results.xyxy[0].cpu().numpy()
            
            # 在图像上绘制边界框和置信度
            for box in boxes:
                x1, y1, x2, y2, conf, cls = box
                if conf > 0.3:  # 置信度阈值
                    # 绘制边界框
                    cv2.rectangle(frame, 
                                (int(x1), int(y1)), 
                                (int(x2), int(y2)), 
                                (0, 255, 0), 2)
                    # 显示置信度
                    label = f'{conf:.2f}'
                    cv2.putText(frame, label, 
                              (int(x1), int(y1)-10), 
                              cv2.FONT_HERSHEY_SIMPLEX, 
                              0.5, (0, 255, 0), 2)
            
            # 显示人数
            people_count = len([box for box in boxes if box[4] > 0.3])
            cv2.putText(frame, f'People Count: {people_count}', 
                        (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, 
                        (0, 255, 0), 2)
        
        # 计算和显示FPS
        fps = 1.0 / (time.time() - start_time)
        cv2.putText(frame, f'FPS: {int(fps)}', 
                    (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 1, 
                    (0, 255, 0), 2)
        
        # 显示操作提示
        cv2.putText(frame, "Press 'q' or 'esc' to quit", 
                    (10, 90), cv2.FONT_HERSHEY_SIMPLEX, 0.7, 
                    (0, 255, 0), 2)
        
        # 显示结果
        cv2.imshow('People Detection', frame)
        
        frame_count += 1
        
        # 检测按键
        key = cv2.waitKey(1) & 0xFF
        # 按'q'或'esc'退出
        if key == ord('q') or key == 27:  # 27是esc键的ASCII码
            print("程序已退出")
            break
    
    # 释放资源
    cap.release()
    cv2.destroyAllWindows()

if __name__ == '__main__':
    try:
        detect_people()
    except KeyboardInterrupt:
        print("\n程序被用户中断")
    except Exception as e:
        print(f"发生错误: {str(e)}")
    finally:
        # 确保程序退出时释放摄像头资源
        cv2.destroyAllWindows()

代码解析

• 模型加载：使用torch.hub.load加载预训练的YOLOv5模型，并设置只检测人类。

• 摄像头设置：通过cv2.VideoCapture打开摄像头，并设置分辨率。

• 检测逻辑：每隔两帧进行一次检测，使用YOLOv5模型对当前帧进行处理，获取检测结果。

• 结果展示：在图像上绘制检测框和置信度，并显示检测到的人数和FPS。

• 退出机制：通过按键检测实现程序的退出。

总结

通过以上代码，您可以实现一个简单的实时人类检测程序。YOLOv5的高效性和准确性使其成为计算机视觉任务中的热门选择。希望这篇博客能帮助您更好地理解如何使用YOLOv5进行人类检测。