1. bytetrack简介
ByteTrack是一种基于检测的目标追踪算法,它在YOLOv8检测器的基础上进行了改进,实现了更高效的目标追踪,具有简单、高效和通用的特点。相较于传统的多目标跟踪方法,ByteTrack不依赖于ReID模型,而是通过关联每个检测框来进行跟踪。这种方法可以有效地解决低分检测框被简单丢弃的问题,从而减少漏检和碎片化轨迹的情况。
ByteTrack算法流程如下:首先,使用目标检测器对当前帧进行检测,得到一系列候选目标框。然后,利用卡尔曼滤波对目标框进行预测,并利用匈牙利算法进行数据关联,将检测框与历史轨迹进行匹配。对于得分较高的目标框,直接与历史轨迹匹配;对于得分较低的目标框,则与第一次没有匹配上的轨迹进行匹配,用于检测目标遮挡的情形。
为了实现高效的实时多目标跟踪,ByteTrack还采用了一些优化策略。例如,对轨迹进行分类,避免在代码阅读时出现混淆的情形;同时,对于连续两帧都未匹配上的轨迹,将其标记为即将删除的轨迹,从而及时清理无效轨迹。
在实际应用中,ByteTrack能够轻松应用到各种多目标跟踪框架中,并取得显著的性能提升。在MOT17测试集上,ByteTrack实现了80.3 MOTA、77.3 IDF1和63.1 HOTA等优异性能指标,同时在单个V100 GPU上运行速度达到了30 FPS。这表明ByteTrack具有高效、准确和实时性强的特点,能够满足实际应用的需求。
2. 实现流程
使用了 YOLOv8 和 ByteTrack 进行目标识别与跟踪。实现流程:如下面代码所示
-
导入必要的库:
cv2
:用于处理视频和图像。os
:用于处理文件路径。ultralytics.YOLO
:用于加载 YOLOv8 模型进行目标检测和跟踪。
-
加载 YOLOv8 模型:
pythonmodel = YOLO(r'track/pt/best.pt')
这行代码加载了预训练的 YOLOv8 模型,该模型用于目标检测和跟踪。
-
设置输入视频文件夹路径和输出视频文件夹路径:
pythoninput_video_folder = r"track/input" output_video_folder = r"track/output"
这里定义了输入视频文件夹路径和输出视频文件夹路径。
-
获取视频文件列表并按文件名排序:
pythonvideo_files = [file for file in os.listdir(input_video_folder) if file.endswith(".mp4")] video_files.sort()
这行代码获取指定文件夹中所有以
.mp4
结尾的视频文件,并按文件名排序。 -
循环处理每个视频文件:
pythonfor video_file in video_files:
这是一个循环,遍历所有视频文件。
-
打开视频文件并设置输出视频路径:
pythoninput_video_path = os.path.join(input_video_folder, video_file) output_video_path = os.path.join(output_video_folder, video_file)
这里根据当前视频文件构建输入视频文件路径和输出视频文件路径。
-
获取视频的帧率和尺寸:
pythonfps = cap.get(cv2.CAP_PROP_FPS) width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
这行代码获取视频的帧率和尺寸。
-
创建视频写入对象:
pythonout = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height))
这行代码创建了一个视频写入对象,用于写入处理后的视频帧。
-
循环处理视频的每一帧:
pythonwhile cap.isOpened():
这是一个循环,用于处理视频的每一帧。
-
读取视频的下一帧:
pythonsuccess, frame = cap.read()
这行代码读取视频的下一帧,并将帧存储在变量
frame
中。 -
使用 YOLOv8 模型进行目标检测和跟踪:
pythonresults = model.track(frame, tracker="ultralytics/cfg/trackers/bytetrack.yaml", persist=True)
这行代码使用 YOLOv8 模型对当前帧进行目标检测和跟踪,使用 ByteTrack 跟踪器,并设置
persist=True
以保持跟踪。 -
将检测和跟踪结果可视化并写入输出视频:
pythonannotated_frame = results[0].plot() out.write(annotated_frame)
这行代码将检测和跟踪结果可视化在当前帧上,并将结果写入输出视频。
-
显示处理后的视频帧:
pythoncv2.imshow("YOLOv8 Tracking", annotated_frame)
这行代码显示处理后的视频帧。
-
检测是否按下 'q' 键:
pythonif cv2.waitKey(1) & 0xFF == ord("q"): break
如果用户按下键盘上的 'q' 键,则跳出循环。
-
释放视频捕获对象和视频写入对象:
pythoncap.release() out.release()
这行代码释放视频捕获对象和视频写入对象,释放视频资源。
-
关闭显示窗口:
pythoncv2.destroyAllWindows()
这行代码关闭 OpenCV 显示的所有窗口。
-
显示处理完成信息:
pythonprint("所有视频处理完成!")
这行代码打印输出所有视频处理完成的消息。
3. 总体代码:
import cv2
from ultralytics import YOLO
import os
# Load the YOLOv8 model
model = YOLO(r'track/pt/best.pt')
# 输入视频文件夹路径和输出视频文件夹路径
input_video_folder = r"track/input"
output_video_folder = r"track/output"
# 获取视频文件夹中的所有视频文件名,并按文件名排序
video_files = [file for file in os.listdir(input_video_folder) if file.endswith(".mp4")]
video_files.sort()
# 循环处理每个视频文件
for video_file in video_files:
# 打开视频文件
input_video_path = os.path.join(input_video_folder, video_file)
cap = cv2.VideoCapture(input_video_path)
# 获取输出视频文件名
output_video_path = os.path.join(output_video_folder, video_file)
# 获取视频的帧率和尺寸
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))
# 创建视频写入对象
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
out = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height))
# Loop through the video frames
while cap.isOpened():
# Read a frame from the video
success, frame = cap.read()
if success:
# Run YOLOv8 tracking on the frame, persisting tracks between frames
results = model.track(frame, tracker="ultralytics/cfg/trackers/bytetrack.yaml", persist=True)
# Visualize the results on the frame
annotated_frame = results[0].plot()
# Write the annotated frame to the output video
out.write(annotated_frame)
# Display the annotated frame
cv2.imshow("YOLOv8 Tracking", annotated_frame)
# Break the loop if 'q' is pressed
if cv2.waitKey(1) & 0xFF == ord("q"):
break
else:
# Break the loop if the end of the video is reached
break
# Release the video capture and video write objects
cap.release()
out.release()
# Close the display window
cv2.destroyAllWindows()
print("所有视频处理完成!")