YOLOv8实战:针对蚂蚁样本的检测、跟踪与计数 原创

发布于 2024-7-11 08:24
浏览
0收藏

本文使用YOLOv8模型,并借助开源Ultralytics框架和BoT-SORT跟踪器,实现了对树上漫步的蚂蚁进行计数。

简介

计算视频中的物体是一项具有挑战性的计算机视觉任务。与静态图像中的物体计数不同,视频信息涉及额外的复杂性,因为物体可能在不同的时间移动、被遮挡或出现和消失,这使得计数过程变得更为复杂。

在本教程中,我们将演示如何使用对象检测和跟踪技术来计数沿着树移动的蚂蚁。具体地说,我们将利用开源的Ultralytics平台,并集成用于检测任务的YOLOv8模型、用于跟踪任务的BoT-SORT跟踪器和用于对蚂蚁数量进行计数的行计数器。

流程概述

在典型的视频对象计数流程中,每一帧都要经历一系列的过程,包括检测、跟踪和计数。以下给出的是针对每个步骤的简述:

  • 检测:对象检测器识别并定位每个帧中的对象,并在其周围生成边界框。
  • 跟踪:跟踪器跨帧跟踪这些对象,为每个对象指定唯一的ID,以确保它们只计数一次。
  • 计数:计数模块聚合这些信息并添加每个新对象以提供准确的结果。

总体流程示意图如下:

YOLOv8实战:针对蚂蚁样本的检测、跟踪与计数-AI.x社区

连接对象检测器、跟踪器和计数器可能需要大量编码。幸运的是,开源的Ultralytics库(参考文献1)提供了无缝集成这些组件的方案,从而简化了这一过程。

1.使用YOLOv8检测对象

第一步是检测每一帧中的蚂蚁,在它们周围产生边界框。在本教程中,我们将使用我预先训练的YOLOv8检测器来检测蚂蚁。我使用Grounding DINO模型(引文2:https://towardsdatascience.com/automatic-labeling-of-object-detection-datasets-using-groundingdino-b66c486656fe?sk=7c98df89b60ea49a6de9efd5278f645e)来标注数据,然后使用标注数据来训练YOLOv8模型。如果你想了解更多关于训练YOLO模型的信息,请参阅我之前关于训练YOLOv5(https://towardsdatascience.com/the-practical-guide-for-object-detection-with-yolov5-algorithm-74c04aac4843?sk=00d2a9d6dd84d6ac4de153cab3dba7c0)的文章,因为这些概念是相似的。对于您的应用程序,您可以使用预先训练的模型或训练自己的自定义模型。

首先,我们需要使用预先训练的权重值来初始化检测器:

from ultralytics import YOLO
# 使用预先训练的权重值来初始化YOLOv8模型
model = YOLO("/path/to/your/yolo_model.pt")

稍后,我们将使用检测器检测视频循环中每帧中的蚂蚁,将检测与跟踪过程相结合。

2.使用BoT-SORT跟踪器跟踪对象

由于蚂蚁在视频帧中多次出现,因此必须跟踪每只蚂蚁并为其分配一个唯一的ID,以确保每只蚂蚁只被计数一次。Ultralytics框架同时支持BoT-SORT跟踪器(见引文3)和ByteTrack跟踪器(见引文4)进行跟踪。

  • ByteTrack跟踪器:提供准确性和速度之间的平衡,降低计算复杂度。它可能无法像BoT-SORT跟踪器那样处理遮挡和相机运动。
  • BoT-SORT跟踪器:与ByteTrack跟踪器相比,提供了更好的跟踪准确性和稳健性,尤其是在具有遮挡和相机运动的挑战性场景中。它是以更高的计算复杂度和更低的帧速率为代价的。

当然,在这些算法之间的选择取决于应用程序的特定要求。

BoT-SORT的工作原理是:BoT-SORT是一个多对象跟踪器,可以同时跟踪多个对象。它将运动和外观信息与相机运动补偿相结合。物体的位置是使用卡尔曼滤波器预测的,与现有轨迹的匹配是基于它们的位置和视觉特征。这种方法允许BoT-SORT跟踪器即使在存在遮挡或相机移动的情况下也能保持准确的轨迹。

配置良好的跟踪器可以补偿探测器的轻微故障。例如,如果对象检测器暂时无法检测到蚂蚁,跟踪器可以使用运动和外观线索来维持蚂蚁的轨迹。

检测器和跟踪器在视频循环内的每个帧上迭代使用以产生轨迹。以下是将其集成到视频处理循环中的方法:

tracks = model.track(frame, persist=True, tracker=’botsort.yaml’, iou=0.2)

跟踪器配置是在“botsort.yaml”文件中定义。您可以调整这些参数以最适合您的需求。要将tracker更改为ByteTrack,只需将“bytrack.yaml”传递给tracker参数即可。

注意,请确保并集上的交集(IoU)值符合您的应用程序要求;IoU阈值(用于非最大值抑制)确定将多接近的检测视为同一对象。persistent=True参数告诉跟踪器当前帧是序列的一部分,并期望前一帧中的轨迹持续到当前帧中。

3.计数对象

现在,我们已经检测到并跟踪了蚂蚁,最后一步是统计视频中穿过指定线的唯一蚂蚁。Ultralytics库中的ObjectCounter类允许我们定义计数区域,该区域可以是直线或多边形。在本教程中,我们将使用一条简单的线作为计数区域。这种方法通过确保蚂蚁过线时只被计数一次来减少错误,即使它的唯一ID因跟踪错误而改变也没有问题。

首先,我们在视频循环之前初始化ObjectCounter类:

counter = solutions.ObjectCounter( 
view_img=True,                     # 处理过程中显示图像 
reg_pts=[(512, 320), (512, 1850)], # 兴趣点区域 
classes_names=model.names,         # YOLO模型的类名 
draw_tracks=True,                  # 为对象绘制跟踪线 
line_thickness=2,                  # 绘制的线的厚度 
)

在视频循环中,ObjectCounter将对跟踪器生成的轨迹进行计数。线的点以[(x1,y1),(x2,y2)]格式传递给reg_pts参数处的计数器。当蚂蚁边界框的中心点第一次与线相交时,它会根据其轨迹方向添加到计数中。在某个方向上移动的对象被计数为“in”,而在另一个方向上运动的对象被计算为“Out”。

# 使用对象计数器对新对象进行计数
frame = counter.start_counting(frame, tracks)

完整的代码

现在,我们已经实现了计数组件。接下来,让我们将代码与视频循环集成到一起,并保存生成的视频。

#安装和导入所需库
%pip install ultralytics
import cv2
from ultralytics import YOLO, solutions

# 定义路径
path_input_video = '/path/to/your/input_video.mp4' 
path_output_video = "/path/to/your/output_video.avi"
path_model = "/path/to/your/yolo_model.pt"

#初始化YOLOv8检测模型
model = YOLO(path_model)

#初始化对象计数器
counter = solutions.ObjectCounter( 
view_img=True,                     #处理过程中显示图像
reg_pts=[(512, 320), (512, 1850)], # 兴趣点区域 
classes_names=model.names,         # YOLO模型的类名
draw_tracks=True,                  # 为对象绘制跟踪线 
line_thickness=2,                  # 绘制的线的厚度 
)

#打开视频文件
cap = cv2.VideoCapture(path_input_video) 
assert cap.isOpened(), "Error reading video file"

#初始化视频写入程序以保存生成的视频
video_writer = cv2.VideoWriter(path_output_video, cv2.VideoWriter_fourcc(*"mp4v"), 30, (1080, 1920))

# 迭代视频帧
frame_count = 0 
while cap.isOpened(): 
success, frame = cap.read() 
if not success: 
print("Video frame is empty or video processing has been successfully completed.") 
break 

# 对当前帧执行对象跟踪 
tracks = model.track(frame, persist=True, tracker='botsort.yaml', iou=0.2) 

# 使用对象计数器对帧中的对象进行计数,并获得标注图像
frame = counter.start_counting(frame, tracks) 

# 将带标注的帧写入输出视频
video_writer.write(frame) 
frame_count += 1

#释放所有资源
cap.release() 
video_writer.release() 
cv2.destroyAllWindows()

# 打印计数结果
print(f'In: {counter.in_counts}\nOut: {counter.out_counts}\nTotal: {counter.in_counts + counter.out_counts}')
print(f'Saves output video to {path_output_video}')

上面的代码将对象检测和跟踪集成到视频处理循环中,以保存标注的视频。通过使用开源视频库OpenCV,我们打开输入视频,并为输出设置一个视频编写器。在每一帧中,我们使用BoT-SORT执行对象跟踪,对对象进行计数,并对帧进行标注。带标注的帧,包括边界框、唯一ID、轨迹以及“入”和“出”计数,都将保存到输出视频中。“in”和“out”计数可以分别从counter.in_counts和counter.out_counts中检索,也可以打印在输出视频中。

YOLOv8实战:针对蚂蚁样本的检测、跟踪与计数-AI.x社区

上图展示了本实验中一个带标注的框架。图中,每只蚂蚁都被分配了一个边界框和一个唯一的ID。蚂蚁穿过粉线时会被计数。图像的一角显示了蚂蚁“进”和“出”的数量。

结束语

在上面带标注的视频中,我们正确地统计了总共85只蚂蚁,其中34只进入,51只退出。对于精确计数,至关重要的是检测器性能良好,跟踪器配置良好。配置良好的跟踪器可以补偿探测器的失误,确保跟踪的连续性。

在带标注的视频中,我们可以看到跟踪器很好地处理了丢失的检测,蚂蚁周围的边界框消失了,并在随后的帧中返回了正确的ID。此外,为同一对象分配不同ID的跟踪错误(例如,蚂蚁#42变成#48)不会影响计数,因为只有越过线的蚂蚁才会被计数。

总之,在本教程中,我们探讨了如何使用先进的对象检测和跟踪技术对视频中的对象进行计数。我们使用YOLOv8模型检测蚂蚁数据,并使用BoT-SORT跟踪器进行稳健跟踪,所有这些部分都与开源Ultralytics库无缝地集成在一起。

参考文献

[1]Ultralytics GitHub(Ultralytics开源框架的代码仓库):https://github.com/ultralytics/ultralytics。

[2]Grounding DINO:Grounding DINO: Marrying DINO with Grounded Pre-Training for Open-Set Object Detection(将DINO与接地预训练结合起来进行开放式物体检测):https://arxiv.org/pdf/2303.05499。

[3]BoT-SORT: Robust Associations Multi-Pedestrian Tracking(BoT-SORT:鲁棒关联多行人跟踪):https://arxiv.org/pdf/2206.14651。

[4]ByteTrack: Multi-Object Tracking by Associating Every Detection Box(ByteTrack:逐个检测框关联法进行多目标跟踪):https://arxiv.org/pdf/2110.06864。

译者介绍

朱先忠,51CTO社区编辑,51CTO专家博客、讲师,潍坊一所高校计算机教师,自由编程界老兵一枚。

原文标题:Mastering Object Counting in Videos,作者:Lihi Gur Arie

链接:

https://towardsdatascience.com/mastering-object-counting-in-videos-3d49a9230bd2​。

©著作权归作者所有,如需转载,请注明出处,否则将追究法律责任
收藏
回复
举报
回复
相关推荐