统计视频中的蚂蚁数量

APPLICATION Nov 10, 2024

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

在本教程中，我们将演示如何使用物体检测和跟踪技术计算沿树移动的蚂蚁数量。我们将利用 Ultralytics 平台集成 YOLOv8 模型进行检测、BoT-SORT 进行跟踪以及行计数器来计算蚂蚁数量。

1、处理管道概述

在典型的视频对象计数管道中，每个帧都会经历一系列过程：检测、跟踪和计数。以下是每个步骤的简要概述：

检测：物体检测器识别并定位每个帧中的物体，并在它们周围生成边界框。
跟踪：跟踪器跨帧跟踪这些物体，为每个物体分配唯一的 ID，以确保它们只被计数一次。
计数：计数模块汇总这些信息并添加每个新对象以提供准确的结果。

连接对象检测器、跟踪器和计数器可能需要大量编码。幸运的是，Ultralytics 库 [1] 通过提供无缝集成这些组件的便捷管道简化了此过程。

2、使用 YOLOv8 检测对象

第一步是检测每帧中的蚂蚁并在它们周围产生边界框。在本教程中，我们将使用我预先训练过的 YOLOv8 检测器来检测蚂蚁。我使用 Grounding DINO [2] 标记数据，然后使用标注好的数据训练 YOLOv8 模型。对于你的应用程序，可以使用预先训练的模型或训练你自己的自定义模型。

首先，我们需要使用预训练权重初始化检测器：

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 参数处的计数器。当蚂蚁边界框的中心点第一次越过线时，它会根据其轨迹方向添加到计数中。朝某个方向移动的物体计为“入”，朝另一个方向移动的物体计为“出”。

  # Use the Object Counter to count new objects 
  frame = counter.start_counting(frame, tracks)

4、完整代码

现在我们已经看到了计数组件，让我们将代码与视频循环集成并保存生成的视频。

# 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、轨迹以及“进”和“出”计数，都保存到输出视频中。“进”和“出”计数可分别从 counter.in_counts 和 counter.out_counts 中检索，并打印在输出视频上。