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

发布于:2024-11-29 ⋅ 阅读:(42) ⋅ 点赞:(0)

在计算机视觉领域,实时人类检测是一个重要的应用场景。本文将介绍如何使用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进行人类检测。