AI项目十九:YOLOV8实现目标追踪

若该文为原创文章,转载请注明原文出处。

主要是学习一下实现目标追踪的原理,并测试一下效果。

目的是通过YOLOV8实现人员检测,并实现人员追踪,没个人员给分配一个ID,实现追踪的效果。

也可以统计人数。在小区办公楼的出入场所,这类很常见。

一、简介

追踪任务是指识别和跟踪特定目标在视频序列中的运动和位置,一般用唯一ID或固定颜色检测框表示),如下图:

AI项目十九:YOLOV8实现目标追踪_第1张图片

目标检测和目标跟踪的区别:
目标检测:目标检测任务要求同时完成对象的定位(即确定对象的边界框位置)和分类(即确定对象的类别)。这意味着目标检测算法必须不仅能够确定对象是否存在,还要知道它是什么。
目标检测通常用于识别和定位图像或视频帧中的对象,通常需要明确的目标类别信息。
目标跟踪:目标跟踪任务更关注对象在帧与帧之间的连续性,通常更注重对象的运动特征,而不要求进行目标的分类。
目标跟踪可以不涉及目标的类别,它的主要目标是维护对象的位置和轨迹,以实现在视频序列中的跟踪。

这里就有个问题,视频中不同时刻的同一个人,位置发生了变化,那么是如何关联上的呢?答案就是匈牙利算法和卡尔曼滤波。

  • 匈牙利算法可以告诉我们当前帧的某个目标,是否与前一帧的某个目标相同。
  • 卡尔曼滤波可以基于目标前一时刻的位置,来预测当前时刻的位置,并且可以比传感器(在目标跟踪中即目标检测器,比如Yolo等)更准确的估计目标的位置。

最经典的是DeepSORT,本篇记录的是如何使用IOU,所以了解下DeepSORT。

DeepSORT对每一帧的处理流程如下:

检测器得到bbox → 生成detections → 卡尔曼滤波预测→ 使用匈牙利算法将预测后的tracks和当前帧中的detecions进行匹配(级联匹配和IOU匹配) → 卡尔曼滤波更新

Frame 0:检测器检测到了3个detections,当前没有任何tracks,将这3个detections初始化为tracks
Frame 1:检测器又检测到了3个detections,对于Frame 0中的tracks,先进行预测得到新的tracks,然后使用匈牙利算法将新的tracks与detections进行匹配,得到(track, detection)匹配对,最后用每对中的detection更新对应的track

二、方法介绍

目前主流的目标跟踪算法都是基于Tracking-by-Detecton策略,即基于目标检测的结果来进行目标跟踪。

实现目标跟踪的方法:

1、IOU

比较前后两帧检测框IOU是否大于指定阈值,是则是同一个物体,不是则分配新ID,此方法对于运动慢的可以,效果差。

AI项目十九:YOLOV8实现目标追踪_第2张图片

2、卡尔曼滤波

卡尔曼滤波是一种用于估计系统状态的优秀算法。它结合了传感器测量和系统模型,通过递归地计算加权平均值,实时更新状态估计。卡尔曼滤波在众多领域,如导航、机器人技术和信号处理中广泛应用,以提高系统的准确性和鲁棒性。

可以用的库:GitHub - adipandas/multi-object-tracker: Multi-object trackers in Python

3、botsort&bytetrack

BoT-SORT是今年非常游戏的跟踪器模型。就所有主要 MOT 指标MOTA、IDF1 和 HOTA而言,BoT-SORT 和 BoT-SORT-ReID 在 MOT17 和 MOT20 测试集的 MOTChallenge数据集中排名第一。对于 MOT17:实现了 80.5 MOTA、80.2 IDF1 和 65.0 HOTA,在跟踪器的排行榜上暂居第一。
论文翻译:https://blog.csdn.net/hhhhhhhhhhwwwwwwwwww/article/details/126890651
论文链接:https://arxiv.org/pdf/2206.14651.pdf
代码:https://github.com/NirAharon/BOT-SORT
 

ByteTrack是基于tracking-by-detection范式的跟踪方法。大多数多目标跟踪方法通过关联分数高于阈值的检测框来获取目标ID。对于检测分数较低的目标,例如遮挡目标,会被简单的丢弃,这带来了不可忽略的问题,包括大量的漏检和碎片化轨迹。为了解决该问题,作者提出了一种简单、高效且通用的数据关联方法BYTE,通过关联每个检测框而不仅仅是高分检测框来进行跟踪。对于低分检测框,利用它们与轨迹的相似性来恢复真实目标并过滤掉背景检测。

BoT-SORT:https://github.com/NirAharon/BoT-SORT

ByteTrack :https://github.com/ifzhang/ByteTrack

三、IOU实现目标追踪

1、环境安装

本人的电脑使用的是CPU(无GPU)版本,所以直接安装,GPU需要安装CUDA等,自行安装。

# 使用Conda为本项目单独创建一个虚拟环境(python 3.8版本)
conda create -n yolov8_env python=3.8
# 激活进入环境
conda activate yolov8_env

# YOLOv8安装方式
pip install ultralytics

2、验证

# 图片cli验证
yolo predict model=yolov8n.pt source=./bus.jpg

# 视频cli验证
yolo predict model=yolov8n.pt source=./test.mp4

3、使用python验证

使用python语言验证主要是熟悉YOLOV8的API,知道如何调用,并测试。

from ultralytics import YOLO
import cv2
import numpy as np
import time

# 加载模型
model = YOLO("./yolov8n.pt")  # load a pretrained model (recommended for training)
objs_labels = model.names  # get class labels
print(objs_labels)


# 打开摄像头
cap = cv2.VideoCapture(0)

while True:
    # 读取一帧
    start_time = time.time()
    ret, frame = cap.read()
    if ret:
        # 检测
        result = list(model(frame, stream=True))[0]  # inference,如果stream=False,返回的是一个列表,如果stream=True,返回的是一个生成器

        boxes = result.boxes  # Boxes object for bbox outputs
        boxes = boxes.cpu().numpy()  # convert to numpy array
        
        # 遍历每个框
        for box in boxes.data:
            l,t,r,b = box[:4].astype(np.int32) # left, top, right, bottom
            conf, id = box[4:] # confidence, class
            # 绘制框
            cv2.rectangle(frame, (l,t), (r,b), (0,0,255), 2)
            # 绘制类别+置信度(格式:98.1%)
            cv2.putText(frame, f"{objs_labels[id]} {conf*100:.1f}%", (l, t-10), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
            
        end_time = time.time()
        fps = 1 / (end_time - start_time)
        # 绘制FPS
        cv2.putText(frame, f"FPS: {fps:.2f}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
                
        # 显示
        cv2.imshow("frame", frame)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    else:
        break

在终端里执行python demo.py

AI项目十九:YOLOV8实现目标追踪_第3张图片

4、IOU实现追踪

'''
iou追踪示例
'''
from ultralytics import YOLO
import cv2
import numpy as np
import time
import random
import os
from shapely.geometry import Polygon, LineString
import json

class IouTracker:
    def __init__(self):

        # 加载检测模型
        self.detection_model = YOLO("./yolov8n.pt")  
        # 获取类别 
        self.objs_labels = self.detection_model.names 
        # 打印类别
        print(self.objs_labels)
        # 只处理person
        self.track_classes = {0: 'person'}

        # 追踪的IOU阈值
        self.sigma_iou = 0.5
        # detection threshold
        self.conf_thresh = 0.3
   
    def iou(sel,bbox1, bbox2):
        """
        计算两个bounding box的IOU
        """

        (x0_1, y0_1, x1_1, y1_1) = bbox1
        (x0_2, y0_2, x1_2, y1_2) = bbox2

        # 计算重叠的矩形的坐标
        overlap_x0 = max(x0_1, x0_2)
        overlap_y0 = max(y0_1, y0_2)
        overlap_x1 = min(x1_1, x1_2)
        overlap_y1 = min(y1_1, y1_2)

        # 检查是否有重叠
        if overlap_x1 - overlap_x0 <= 0 or overlap_y1 - overlap_y0 <= 0:
            return 0

        # 计算重叠矩形的面积以及两个矩形的面积
        size_1 = (x1_1 - x0_1) * (y1_1 - y0_1)
        size_2 = (x1_2 - x0_2) * (y1_2 - y0_2)
        size_intersection = (overlap_x1 - overlap_x0) * (overlap_y1 - overlap_y0)
        size_union = size_1 + size_2 - size_intersection
        # 计算IOU
        return size_intersection / size_union


    def predict(self, frame):
        '''
        检测
        '''
        result = list(self.detection_model(frame, stream=True, conf=self.conf_thresh))[0]  # inference,如果stream=False,返回的是一个列表,如果stream=True,返回的是一个生成器
        boxes = result.boxes  # Boxes object for bbox outputs
        boxes = boxes.cpu().numpy()  # convert to numpy array

        dets = [] # 检测结果
        # 遍历每个框
        for box in boxes.data:
            l,t,r,b = box[:4] # left, top, right, bottom
            conf, class_id = box[4:] # confidence, class
            # 排除不需要追踪的类别
            if class_id not in self.track_classes:
                continue
            dets.append({'bbox': [l,t,r,b], 'score': conf, 'class_id': class_id })
        return dets
    
    
    def main(self):
        '''
        主函数
        '''
        # 读取视频
        cap = cv2.VideoCapture("./media/video.mp4")
  
        # 获取视频帧率、宽、高
        fps = cap.get(cv2.CAP_PROP_FPS)
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        print(f"fps: {fps}, width: {width}, height: {height}")


        tracks_active = [] # 活跃的跟踪器
        frame_id = 1 # 帧ID
        track_idx = 1 # 跟踪器ID
        
        # writer
        out = cv2.VideoWriter("./test_out.mp4", cv2.VideoWriter_fourcc(*'mp4v'), fps, (1280, 720))

        while True:
            # 读取一帧
            start_time = time.time()
            ret, raw_frame = cap.read()
            if ret:
                # 检测
                frame = cv2.resize(raw_frame, (1280, 720))
                raw_frame = frame
                dets = self.predict(raw_frame)
                # 更新后的跟踪器
                updated_tracks = [] 
                # 遍历活跃的跟踪器
                for track in tracks_active:
                    if len(dets) > 0:
                        # 根据最大IOU更新跟踪器,先去explain.ipynb中看一下MAX用法
                        best_match = max(dets, key=lambda x: self.iou(track['bboxes'][-1], x['bbox'])) # 找出dets中与当前跟踪器(track['bboxes'][-1])最匹配的检测框(IOU最大)
                        # 如果最大IOU大于阈值,则将本次检测结果加入跟踪器
                        if self.iou(track['bboxes'][-1], best_match['bbox']) > self.sigma_iou:
                            # 将本次检测结果加入跟踪器
                            track['bboxes'].append(best_match['bbox'])
                            track['max_score'] = max(track['max_score'], best_match['score'])
                            track['frame_ids'].append(frame_id)
                            # 更新跟踪器
                            updated_tracks.append(track)
                            # 删除已经匹配的检测框,避免后续重复匹配以及新建跟踪器
                            del dets[dets.index(best_match)]


                # 如有未分配的目标,创建新的跟踪器
                new_tracks = []
                for det in dets: # 未分配的目标,已经分配的目标已经从dets中删除
                    new_track = {
                        'bboxes': [det['bbox']], # 跟踪目标的矩形框
                        'max_score': det['score'], # 跟踪目标的最大score
                        'start_frame': frame_id,  # 目标出现的 帧id
                        'frame_ids': [frame_id],  # 目标出现的所有帧id
                        'track_id': track_idx,    # 跟踪标号
                        'class_id': det['class_id'], # 类别
                        'is_counted': False       # 是否已经计数
                    }
                    track_idx += 1
                    new_tracks.append(new_track)
                # 最终的跟踪器
                tracks_active = updated_tracks + new_tracks

                cross_line_color = (0,255,0) # 越界线的颜色

                # 绘制跟踪器
                for tracker in tracks_active:
                    # 绘制跟踪器的矩形框
                    l,t,r,b = tracker['bboxes'][-1]
                    # 取整
                    l,t,r,b = int(l), int(t), int(r), int(b)
                    class_id = tracker['class_id']
                    cv2.rectangle(raw_frame, (l,t), (r,b), cross_line_color, 2)
                    # 绘制跟踪器的track_id + class_name + score(99.2%格式)
                    cv2.putText(raw_frame, f"{tracker['track_id']}", (l, t-10), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0,255,0), 2)

                # 设置半透明
                color = (0,0,0)
                alpha = 0.2
                l,t = 0,0
                r,b = l+240,t+40
                raw_frame[t:b,l:r,0] = raw_frame[t:b,l:r,0] * alpha + color[0] * (1-alpha)
                raw_frame[t:b,l:r,1] = raw_frame[t:b,l:r,1] * alpha + color[1] * (1-alpha)
                raw_frame[t:b,l:r,2] = raw_frame[t:b,l:r,2] * alpha + color[2] * (1-alpha)

                # end time
                end_time = time.time()
                # FPS
                fps = 1 / (end_time - start_time)
                # 绘制FPS
                cv2.putText(raw_frame, f"FPS: {fps:.2f}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
                
               
 
                # 显示
                cv2.imshow("frame", raw_frame)
                if cv2.waitKey(1) & 0xFF == ord('q'):
                    break
                    
                out.write(raw_frame)
            else:
                break
                
        out.release()  

# 实例化
iou_tracker = IouTracker()
# 运行
iou_tracker.main()

测试效果,视频是马路上的,如果想要效果好,建议自己训练模型,使用的是yolov8n.pt模型

如有侵权,或需要完整代码,请及时联系博主。

你可能感兴趣的:(AI计算机视觉,YOLO)