YOLOv11实现区域人流量统计系统

2,714 阅读7分钟

基于YOLOv11的目标检测模型,我们可以实现区域人流量统计系统。

界面预览

无标题视频——使用Clipchamp制作.gif

实现原理

目标检测是用来识别图片或视频中物体的位置和类别的。

简单来说,目标检测的结果就是一组框框(我们叫它“边界框”),这些框把图片里的物体圈起来,同时还会告诉你每个框里的物体是什么类别以及识别的可信度分数。当你需要找出画面里有哪些感兴趣的物体,但不需要精确知道它们的具体形状或具体位置时,目标检测就是一个很好的选择。

我们用的模型是 yolo11n.pt,在线推理速度非常快,可以检测出多种类型的目标,还可以结合track算法实现目标追踪。

代码实现

代码实现非常通俗易懂,重点要注意queue_region,因为该变量是用来表示检测区域的,设置的不对,检测位置就会出错。

还要注意 QueueManager中的classes,因为是预训练的模型,所以当 classes=[0] 时,默认检测出的是人。

# Init Queue Manager
queue = solutions.QueueManager(
    show=True,  # Display the output
    model="yolo11n.pt",  # Path to the YOLO11 model file
    region=queue_region,  # Pass queue region points
    classes=[0],  # If you want to count specific classes i.e person and car with COCO pretrained model.
    # line_width=2,  # Adjust the line width for bounding boxes and text display
)

完整源码如下:

import cv2

from ultralytics import solutions

cap = cv2.VideoCapture("monitor.mp4")

assert cap.isOpened(), "Error reading video file"
w, h, fps = (int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT, cv2.CAP_PROP_FPS))

# Video writer
video_writer = cv2.VideoWriter("queue_management.avi", cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h))

# Define queue region points
queue_region = [(20, 400), (1080, 400), (1080, 800), (20, 800)]  # Define queue region points

# Init Queue Manager
queue = solutions.QueueManager(
    show=True,  # Display the output
    model="yolo11n.pt",  # Path to the YOLO11 model file
    region=queue_region,  # Pass queue region points
    classes=[0],  # If you want to count specific classes i.e person and car with COCO pretrained model.
    # line_width=2,  # Adjust the line width for bounding boxes and text display
)

# Process video
while cap.isOpened():
    success, im0 = cap.read()

    if success:
        out = queue.process_queue(im0)
        video_writer.write(im0)
        if cv2.waitKey(1) & 0xFF == ord("q"):
            break
        continue

    print("Video frame is empty or video processing has been successfully completed.")
    break

cap.release()
cv2.destroyAllWindows()

源码解读

仔细分析 QueueManager 类源码,方便我们后期定制出自己的业务 QueueManager 类源码。

源码实现了一个基于实时视频流的队列管理器(QueueManager),可以在画面中检测物体、追踪它们的运动轨迹,并统计某个指定区域内的物体数量。


1. 整体功能概览

  • 检测物体:在视频流中识别目标,比如人或车辆。
  • 追踪物体:跟踪目标的运动轨迹。
  • 统计数量:判断目标是否进入某个指定区域,并统计区域内的目标数量。
  • 可视化:将检测结果(如目标框、轨迹、区域、数量)绘制在视频帧上。

2. 类结构和继承

class QueueManager(BaseSolution)

  • QueueManager 是一个类,继承自 BaseSolution
  • BaseSolution 是一个基础功能类,提供了一些通用的操作,比如视频帧的处理、目标检测等。
  • QueueManager 在此基础上扩展了一些功能,比如队列统计、区域绘制、轨迹记录等。

3. 构造函数:初始化 QueueManager

def __init__(self, **kwargs):
    super().__init__(**kwargs)
    self.initialize_region()
    self.counts = 0  # Queue counts Information
    self.rect_color = (255, 255, 255)  # Rectangle color
    self.region_length = len(self.region)  # Store region length for further usage

初始化主要做了以下几件事:

  1. 继承初始化:调用父类的初始化方法(super().__init__),加载父类里的功能参数,比如输入视频的路径等。
  2. 初始化区域信息:调用 initialize_region() 方法,定义统计的区域位置(比如某个矩形或多边形区域)。
  3. 计数变量:初始化 self.counts = 0,用来存储当前帧中区域内的目标数量。
  4. 颜色和区域长度
    • self.rect_color 定义统计区域的颜色(白色)。
    • self.region_length 保存统计区域的点数,比如三角形是 3 个点,矩形是 4 个点。

4. 核心方法:process_queue

  1. 清空计数器。
  2. 初始化绘制工具。
  3. 提取当前帧中的目标和轨迹信息。
  4. 绘制区域、目标框、轨迹和质心。
  5. 判断目标是否进入统计区域,更新计数。
  6. 在画面上显示统计结果。
  7. 返回处理后的帧。

我们重点看下检查目标是否在统计区域内的实现思路

track_history = self.track_history.get(track_id, [])
prev_position = None
if len(track_history) > 1:
    prev_position = track_history[-2]
if self.region_length >= 3 and prev_position and self.r_s.contains(self.Point(self.track_line[-1])):
    self.counts += 1
  1. 获取目标的历史轨迹:从 self.track_history 中取出当前目标的轨迹历史。

    • track_history 是一个字典,存储每个目标的运动路径(用 track_id 区分不同的目标)。
    • 如果目标的轨迹点数超过 1,取倒数第二个位置点 prev_position
  2. 判断条件

    • 确保统计区域是一个合法的多边形(点数 >=3)。
    • 确保目标有上一帧的轨迹点(prev_position 存在)。
    • 检查目标当前质心是否在统计区域内(self.r_s.contains(self.Point(self.track_line[-1])),这里的 self.Point 表示目标当前的位置点)。
  3. 更新计数器

    • 如果目标满足上述条件,就增加计数 self.counts += 1

第二个重点是显示队列数量实现思路,也就是数据可视化,文字、线框、颜色可以自由定制。

self.annotator.queue_counts_display(
    f"Queue Counts : {str(self.counts)}",
    points=self.region,
    region_color=self.rect_color,
    txt_color=(104, 31, 17),
)
  • 功能:在画面上显示当前统计区域中的目标数量。
  • 参数
    • Queue Counts:显示的文字内容,比如 "Queue Counts: 5"。
    • points=self.region:区域的位置(多边形的顶点)。
    • region_color=self.rect_color:区域的颜色(白色)。
    • txt_color=(104, 31, 17):文字颜色。

# Ultralytics YOLO 🚀, AGPL-3.0 license

from ultralytics.solutions.solutions import BaseSolution
from ultralytics.utils.plotting import Annotator, colors


class QueueManager(BaseSolution):
    """
    Manages queue counting in real-time video streams based on object tracks.

    This class extends BaseSolution to provide functionality for tracking and counting objects within a specified
    region in video frames.

    Attributes:
        counts (int): The current count of objects in the queue.
        rect_color (Tuple[int, int, int]): RGB color tuple for drawing the queue region rectangle.
        region_length (int): The number of points defining the queue region.
        annotator (Annotator): An instance of the Annotator class for drawing on frames.
        track_line (List[Tuple[int, int]]): List of track line coordinates.
        track_history (Dict[int, List[Tuple[int, int]]]): Dictionary storing tracking history for each object.

    Methods:
        initialize_region: Initializes the queue region.
        process_queue: Processes a single frame for queue management.
        extract_tracks: Extracts object tracks from the current frame.
        store_tracking_history: Stores the tracking history for an object.
        display_output: Displays the processed output.

    Examples:
        >>> queue_manager = QueueManager(source="video.mp4", region=[100, 100, 200, 200, 300, 300])
        >>> for frame in video_stream:
        ...     processed_frame = queue_manager.process_queue(frame)
        ...     cv2.imshow("Queue Management", processed_frame)
    """

    def __init__(self, **kwargs):
        """Initializes the QueueManager with parameters for tracking and counting objects in a video stream."""
        super().__init__(**kwargs)
        self.initialize_region()
        self.counts = 0  # Queue counts Information
        self.rect_color = (255, 255, 255)  # Rectangle color
        self.region_length = len(self.region)  # Store region length for further usage

    def process_queue(self, im0):
        """
        Processes the queue management for a single frame of video.

        Args:
            im0 (numpy.ndarray): Input image for processing, typically a frame from a video stream.

        Returns:
            (numpy.ndarray): Processed image with annotations, bounding boxes, and queue counts.

        This method performs the following steps:
        1. Resets the queue count for the current frame.
        2. Initializes an Annotator object for drawing on the image.
        3. Extracts tracks from the image.
        4. Draws the counting region on the image.
        5. For each detected object:
           - Draws bounding boxes and labels.
           - Stores tracking history.
           - Draws centroids and tracks.
           - Checks if the object is inside the counting region and updates the count.
        6. Displays the queue count on the image.
        7. Displays the processed output.

        Examples:
            >>> queue_manager = QueueManager()
            >>> frame = cv2.imread("frame.jpg")
            >>> processed_frame = queue_manager.process_queue(frame)
        """
        self.counts = 0  # Reset counts every frame
        self.annotator = Annotator(im0, line_width=self.line_width)  # Initialize annotator
        self.extract_tracks(im0)  # Extract tracks

        self.annotator.draw_region(
            reg_pts=self.region, color=self.rect_color, thickness=self.line_width * 2
        )  # Draw region

        for box, track_id, cls in zip(self.boxes, self.track_ids, self.clss):
            # Draw bounding box and counting region
            self.annotator.box_label(box, label=self.names[cls], color=colors(track_id, True))
            self.store_tracking_history(track_id, box)  # Store track history

            # Draw tracks of objects
            self.annotator.draw_centroid_and_tracks(
                self.track_line, color=colors(int(track_id), True), track_thickness=self.line_width
            )

            # Cache frequently accessed attributes
            track_history = self.track_history.get(track_id, [])

            # store previous position of track and check if the object is inside the counting region
            prev_position = None
            if len(track_history) > 1:
                prev_position = track_history[-2]
            if self.region_length >= 3 and prev_position and self.r_s.contains(self.Point(self.track_line[-1])):
                self.counts += 1

        # Display queue counts
        self.annotator.queue_counts_display(
            f"Queue Counts : {str(self.counts)}",
            points=self.region,
            region_color=self.rect_color,
            txt_color=(104, 31, 17),
        )
        self.display_output(im0)  # display output with base class function

        return im0  # return output image for more usage

功能拓展

本文将对接基于RTSP协议的监控摄像头,实现在线分析视频功能。 很简单,我们只需改动一行代码,如下所示。

cap = cv2.VideoCapture("rtsp://admin:admin@IP地址:端口号/live0")

总结

本文介绍了基于 YOLOv11 模型的目标检测与追踪技术,开发了一套区域人流量统计系统。通过将目标检测与轨迹追踪算法结合,系统可以实时识别视频流中的目标(如人或车辆),并统计指定区域内的目标数量。该系统功能强大,代码实现简洁高效,适用于公共场所人流量监控等场景。