使用YOLOv8和ByteTracker进行实时人员跟踪和计数

新机器视觉 · 公众号 · · 2024-10-22 17:22

正文

在计算机视觉领域，实时跟踪和统计人数对于各种应用至关重要，从监控到事件管理。在这篇博客文章中，我们将探讨如何利用YOLOv8和ByteTracker实现准确的人数统计。

引言

YOLOv8（You Only Look Once，第八版）是一种以其速度和准确性而闻名的最新对象检测模型。
ByteTracker是一种先进的跟踪算法，旨在维持对象在帧之间的身份，使其成为执行人数统计等任务的理想选择。

这种组合不仅允许我们在帧中检测到人，而且还能够跟踪他们在帧之间的移动，为实时人数统计提供了强大的解决方案。

先决条件

在深入实现之前，请确保您具备以下条件：

Python 3.10
Ultralytics

设置环境

首先，创建并激活虚拟环境以管理依赖项：

conda create -n person-tracker python==3.10conda activate person-tracker

安装必要的库：

pip install ultralytics

实现人数统计

1. 导入库 从导入所需的库开始：

from ultralytics import YOLOfrom datetime import datetimeimport os

2. 定义PersonTracker类

创建一个PersonTracker类，该类集成了用于检测的YOLOv8和用于跟踪的ByteTracker：

class PersonTracker:    def __init__(self, model_path, result_dir='results/', tracker_config="bytetrack.yaml", conf=0.5, device='cuda:0',                 iou=0.5, img_size=(720, 1080)):        self.model = YOLO(model_path)        self.result_dir = result_dir        self.tracker_config = tracker_config        self.conf = conf        self.device = device        self.iou = iou        self.img_size = img_size
    def create_result_file(self):        folder_name = datetime.now().strftime("%Y-%m-%d-%H-%M-%S")        result_file_path = os.path.join(self.result_dir, folder_name + ".txt")        os.makedirs(self.result_dir, exist_ok=True)        with open(result_file_path, 'w') as file:            file.write(folder_name + "\n")        return result_file_path
    def detect_and_track(self, source, show=True, logger=None):        result_file = self.create_result_file()        person_count = 0        previous_person_count = 0
        results = self.model.track(            source, show=show, stream=True, tracker=self.tracker_config, conf=self.conf,            device=self.device, iou=self.iou, stream_buffer=True, classes=[0], imgsz=self.img_size        )
        for i, result in enumerate(results):            boxes = result.boxes            try:                id_count = boxes.id.int().tolist()                max_id = max(id_count)
                if max_id > person_count:                    person_count = max_id
                if person_count != previous_person_count:                    previous_person_count = person_count                    with open(result_file, 'a') as filewrite:                        filewrite.write(f"Person count: {person_count}\n")
                    if logger:                        logger.info(f"Person count: {person_count}")
            except Exception as e:                pass

3. 运行人员跟踪器

视频文件跟踪

if __name__ == '__main__':    source = "path/to/your/video.mp4"    tracker = PersonTracker(model_path='path/to/yolov8_model.pt')    tracker.detect_and_track(source=source)

网络摄像头跟踪

if __name__ == '__main__':    source = 0  # Use 0 for the default webcam    tracker = PersonTracker(model_path='path/to/yolov8_model.pt')    tracker.detect_and_track(source=source)

RTSP流跟踪

if __name__ == '__main__':    source = 'rtsp://username:password@ip_address:port/path'    tracker = PersonTracker(model_path='path/to/yolov8_model.pt')    tracker.detect_and_track(source=source)