点击下方卡片,关注“小白玩转Python”公众号
使用检测和跟踪技术一步步指导如何计算树上漫游的蚂蚁数量。
引言
在视频中计数对象是一项具有挑战性的计算机视觉任务。与在静态图像中计数对象不同,视频涉及额外的复杂性,因为对象可以移动,被遮挡,或在不同时间出现和消失,这使得计数过程变得复杂。在本教程中,我们将演示如何使用对象检测和跟踪技术来计算沿树移动的蚂蚁数量。我们将利用Ultralytics平台集成YOLOv8模型进行检测,BoT-SORT用于跟踪,以及一个计数器来计算蚂蚁数量。
流程概述
在一个典型的视频对象计数流程中,每一帧都会经历一系列过程:检测、跟踪和计数。以下是每个步骤的简要概述:
检测:对象检测器在每一帧中识别并定位对象,在其周围生成边界框。
跟踪:跟踪器跨帧跟踪这些对象,为每个对象分配唯一ID,以确保它们只被计数一次。
计数:计数模块聚合这些信息,并添加每个新对象以提供准确的结果。
连接一个对象检测器、一个跟踪器和一个计数器可能需要大量的编码。幸运的是,Ultralytics库[1]通过提供方便的流程简化了这一过程,这些组件可以无缝集成。
1. 使用YOLOv8检测对象
第一步是在每一帧中检测蚂蚁并围绕它们生成边界框。在本教程中,我们将使用我事先训练好的YOLOv8检测器来检测蚂蚁。我使用了Grounding DINO[2]来标记数据,然后使用注释过的数据来训练YOLOv8模型。如果你想更多地了解如何训练YOLO模型,请参考我之前关于训练YOLOv5的帖子,因为概念是相似的。对于你的应用,你可以使用预训练模型或训练自己的自定义模型。要开始,我们需要使用预训练权重初始化检测器:
from ultralytics import YOLO
# Initialize YOLOv8 model with pre-trained weights
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’文件中定义。你可以调整这些参数以最好地适应你的需求。要将跟踪器更改为ByteTrack,只需将‘bytetrack.yaml’传递给跟踪器参数。
确保交集比值(IoU)值符合你的应用要求;IoU阈值(用于非最大抑制)决定了检测必须多么接近才能被视为同一对象。persist=True参数告诉跟踪器当前帧是序列的一部分,并期望前一帧的轨迹持续到当前帧。
3. 计数对象
现在我们已经检测并跟踪了蚂蚁,最后一步是计算穿过视频中指定线的独特的蚂蚁数量。Ultralytics库中的ObjectCounter类允许我们定义一个计数区域,可以是一条线或一个多边形。对于本教程,我们将使用一条简单的线作为我们的计数区域。这种方法通过确保蚂蚁在穿过线时只被计数一次(即使由于跟踪错误其唯一ID发生变化),从而减少错误。首先,在视频循环之前初始化ObjectCounter:
counter = solutions.ObjectCounter(
view_img=True, # Display the image during processing
reg_pts=[(512, 320), (512, 1850)], # Region of interest points
classes_names=model.names, # Class names from the YOLO model
draw_tracks=True, # Draw tracking lines for objects
line_thickness=2, # Thickness of the lines drawn
)
在视频循环内部,ObjectCounter将计算由跟踪器产生的轨迹。线的点以[(x1, y1), (x2, y2)]格式传递给计数器的reg_pts参数。当蚂蚁的边界框中心点首次穿过线时,根据其轨迹方向将其添加到计数中。向某个方向移动的对象被计数为‘In’,向另一个方向移动的对象被计数为‘Out’。
# Use the Object Counter to count new objects
frame = counter.start_counting(frame, tracks)
完整代码
现在我们已经看到了计数组件,让我们将代码与视频循环集成并保存结果视频。
# Install and import Required Libraries
%pip install ultralytics
import cv2
from ultralytics import YOLO, solutions
# Define paths:
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"
# Initialize YOLOv8 Detection Model
model = YOLO(path_model)
# Initialize Object Counter
counter = solutions.ObjectCounter(
view_img=True, # Display the image during processing
reg_pts=[(512, 320), (512, 1850)], # Region of interest points
classes_names=model.names, # Class names from the YOLO model
draw_tracks=True, # Draw tracking lines for objects
line_thickness=2, # Thickness of the lines drawn
)
# Open the Video File
cap = cv2.VideoCapture(path_input_video)
assert cap.isOpened(), "Error reading video file"
# Initialize the Video Writer to save resulted video
video_writer = cv2.VideoWriter(path_output_video, cv2.VideoWriter_fourcc(*"mp4v"), 30, (1080, 1920))
# itterate over video frames:
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
# Perform object tracking on the current frame
tracks = model.track(frame, persist=True, tracker='botsort.yaml', iou=0.2)
# Use the Object Counter to count objects in the frame and get the annotated image
frame = counter.start_counting(frame, tracks)
# Write the annotated frame to the output video
video_writer.write(frame)
frame_count += 1
# Release all Resources:
cap.release()
video_writer.release()
cv2.destroyAllWindows()
# Print counting results:
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,我们打开输入视频并设置输出的视频写入器。在每一帧中,我们使用BoTSORT进行对象跟踪,计数对象,并注释帧。包括边界框、唯一ID、轨迹和‘in’和‘out’计数的注释帧被保存到输出视频中。‘in’和‘out’计数可以从counter.in_counts和counter.out_counts中分别检索,并且也打印在输出视频上。
注释过的帧。每只蚂蚁都被赋予了边界框和唯一ID。蚂蚁在穿过粉色线时被计数。移动‘in’和‘out’的蚂蚁数量显示在图像的角落。
结束语
在注释视频中,我们正确地计算了总共85只蚂蚁,其中34只进入,51只退出。为了精确计数,检测器的性能良好和跟踪器配置得当至关重要。一个配置良好的跟踪器可以弥补检测器的遗漏,确保跟踪的连续性。
在注释视频中,我们可以看到跟踪器很好地处理了丢失的检测,这可以从蚂蚁周围的边界框消失以及随后帧中以正确ID返回的情况中看出。此外,将不同的ID分配给同一对象的跟踪错误(例如,蚂蚁#42变成#48)并没有影响计数,因为只有穿过线的蚂蚁才被计数。
在本教程中,我们探讨了如何使用高级对象检测和跟踪技术来计算视频中的对象。我们使用YOLOv8来检测蚂蚁,使用BoT-SORT进行稳健的跟踪,所有这些都与Ultralytics库无缝集成。
· END ·
HAPPY LIFE
本文仅供学习交流使用,如有侵权请联系作者删除