大家好🤗!这次我带着一篇关于我的样本颜色检测和追踪项目的文章回来了。这个项目是我为机器人团队开发的,在这个项目中,我必须从头开始做所有事情,包括准备自己的数据集、注释图像、使用YOLOv8训练检测模型,以及集成Deep SORT追踪算法。我还必须更新部分 Deep SORT 代码,使其能够与 TensorFlow 2.x 兼容,从而避免经常出现的弃用代码问题。这需要一些反复试验,但现在您可以使用这些脚本,而不必担心版本冲突或复杂的错误。您可以通过后台私我访问我创建的数据集。
Yolov8 和 DeepSORT 简介
首先,让我解释一下 YOLO 以及它为何在实时物体检测中如此广泛地应用。YOLO 是“You Only Look Once”(你只看一次)的缩写,由 Joseph Redmon 及其同事提出,是一种检测图像或视频中物体的高级方法。其理念是通过对图像进行单次前向传播来检测物体,而不是使用多个候选阶段,这使得 YOLO 速度极快,同时又不会牺牲太多准确率。随着时间的推移,YOLO 经历了许多改进,最终推出了 YOLOv7 和 YOLOv8 等版本,它们可以更有效地处理物体检测、语义分割和实例分割等任务。我选择 YOLOv8 的一个关键原因是它的灵活性和易用性,尤其是通过 Ultralytics 库,它显著简化了训练和推理的过程。
至于 Deep SORT(简单在线实时追踪),它是一种基于原始 SORT 追踪器构建的高级算法。Deep SORT 不仅可以根据运动预测物体的未来位置,还能利用基于深度学习的外观特征,即使在不同物体重叠或部分遮挡的情况下,也能更稳健地保持追踪身份。换句话说,Deep SORT 通过为每个检测到的物体分配一个一致的 ID 来随时间追踪物体,让您可以监控特定物体在帧与帧之间的移动情况。在本项目中,Deep SORT 依赖于基于 TensorFlow 的神经网络来提取外观嵌入(通常使用名为 MARS 的模型)。由于 TensorFlow 从 1.x 版本到 2.x 版本已经有了显著的改进,一些较旧的 Deep SORT 代码库需要进行一些小的更新,以确保所有代码都能正常运行,避免出现可怕的“无会话”或“无属性”错误。我已经在代码中完成了这些更新,因此您可以直接将 Deep SORT 与 TensorFlow 2.x 集成。
现在我们已经了解了 YOLO 在实时检测中的用途以及 Deep SORT 在跟踪中的实用性,让我们看看我是如何创建和注释数据集的,如何将这些注释转换为 YOLO 格式,以及我最终如何训练自定义 YOLOv8 模型,然后在视频上检测和跟踪样本🤯👩💻🥁。
获取和注释数据集
因此,我首先获取了几个想要检测的颜色样本,主要是红色、蓝色和黄色样本👾🐾。我拍摄了大约 80 张这些样本的照片,它们以不同的配置和背景呈现。有时,图像中只有一个样本;有时,我会一次放置多个样本。收集完所有这些图像后,我需要创建边界框注释,以便 YOLO 模型了解每个样本在图像中的位置以及应该为其分配的颜色标签。
为了注释数据,我使用了CVAT(计算机视觉注释工具)网站,这是一个功能强大且用户友好的在线边界框注释平台。市面上有很多工具,比如 labelImg,但我已经习惯了 CVAT 的工作流程,所以决定继续使用它。
现在,让我向您介绍如何使用 cvat.ai 注释图像:
首先,您需要通过切换到“项目”部分来创建一个项目:
在这里,只需点击“+”并选择“创建新项目”选项。之后会出现这个页面,你可以为你的项目指定任意名称:
您需要逐个添加标签。不用担心,标签不会在您创建时显示,而是在您创建任务时显示,稍后我们会进行创建任务。因此,对于这个项目,我添加了三个标签:“红色”、“蓝色”和“黄色”,然后点击“提交并继续”,您将在屏幕上看到一条通知,提示“项目已创建”。
现在,转到“任务”部分,点击“+”并选择“创建新任务”。然后输入您的任务名称,名称可以随意填写。选择您刚刚创建的项目,其标签将自动用于注释。然后上传所有您想要注释的图像,最后点击“提交并继续”。
然后,当您再次转到“任务”部分时,这次您将看到已创建的任务:
打开它并点击突出显示的#Job以输入要注释的任务:
然后要绘制边界框,请点击矩形形状,选择要注释的对象标签,并选择要从边界框中存储的点数。在本项目中,我选择了 4 个点,然后点击“形状”,再点击对象的 4 个角,即可调整边界框。对图像中的所有对象继续执行此操作,然后不要忘记保存,并移动到下一张图片,并继续重复此过程,直到最后一张图片。
完成后,返回“任务”部分,单击彼此顶部的三个点并说“导出任务数据集”:
选择图像 1.1 的 CVAT 选项,它将为您提供一个 comments.xml 文档,其中包含图像及其图像名称、标签、尺寸和边界框坐标。
现在,是时候开始编码了:
Yolov8检测模型的数据预处理、训练与评估
由于我现在的笔记本电脑没有 GPU,所以我在 Google Colab 上训练了 Yolov8 模型。你只需要将你的 annotations.xml 文件和包含原始图像的文件夹上传到 Google Colab 即可。请确保安装 ultralytics 库,方法如下:
!pip 安装 ultralytics
在训练 YOLOv8 模型之前,我们需要将 CVAT 边界框[xtl, ytl, xbr, ybr]
(左上角和右下角)转换为 YOLO 格式[x_center, y_center, w, h]
(坐标标准化,范围从 0 到 1)。我们还需要将数据集拆分为训练和验证子集(我选择了 80:20 的拆分比例),并生成一个dataset.yaml
文件来告知 YOLO 图像和标签的存储位置 😱💪🤓。如果您希望将 YOLO 文件存储在 Google Colab 中,只需将您的 Google Drive 挂载到 Google Colab,并从 Google Drive 中指定一个目录作为输出目录即可:
# 将您的 Google Drive 安装到 Google Colab
from google.colab import drive
drive.mount( '/content/drive' )
从ultralytics导入YOLO
导入os
导入random
导入shutil
导入xml.etree.ElementTree作为ET
IMAGES_DIR = "images"
XML_FILE = "annotations.xml"
OUTPUT_DIR = "/content/drive/MyDrive/Robotics_Color_Sample_Detection/data_yolo"
TRAIN_RATIO = 0.8
# 数据集的类映射 (label_name -> class_index)。CLASS_MAP
= {
"red" : 0 ,
"blue" : 1 ,
"yellow" : 2
}
在这里,我导入了所有必要的库,并指定了包含原始图像的文件夹目录、文件名annotations.xml
以及OUTPUT_DIR
用于存储处理后数据集的输出目录 ( )。我还将其定义TRAIN_RATIO
为 0.8,这是我想要放入训练集的图像比例,其余部分用于验证,但您可以调整此阈值。最后,我声明了一个名为 的字典,CLASS_MAP
将颜色名称映射到 YOLO 训练所需的数字 ID。
接下来,让我们定义一个名为 的辅助函数convert_to_yolo_bbox
。它将采用 CVAT 格式的边界框(xtl, ytl, xbr, ybr
)加上图像的img_width
和,并计算出 YOLO 兼容的边界框(在标准化坐标系中)。YOLO 格式要求所有值都在 0 到 1 之间,所以我将其除以或。img_height
(x_center, y_center, width, height)
img_width
img_height
接下来,让我们为训练图像、验证图像及其对应的标签文件创建单独的目录。os.makedirs(..., exist_ok=True)
如果这些文件夹尚不存在,则调用这些函数来创建它们。
def convert_to_yolo_bbox ( xtl, ytl, xbr, ybr, img_width, img_height ):
"""
将 CVAT 边界框 [xtl, ytl, xbr, ybr] 转换
为相对坐标中的 YOLO 边界框 [x_center, y_center, width, height]。
"""
x_center = ((xtl + xbr) / 2.0 ) / img_width
y_center = ((ytl + ybr) / 2.0 ) / img_height
w = (xbr - xtl) / img_width
h = (ybr - ytl) / img_height
return x_center, y_center, w, h
# 为图像和标签创建 train/val 文件夹
train_img_dir = os.path.join(OUTPUT_DIR, "images" , "train" )
val_img_dir = os.path.join(OUTPUT_DIR,“图像”,“val”)
train_lbl_dir = os.path.join(OUTPUT_DIR,“标签”,“训练”)
val_lbl_dir = os.path.join(OUTPUT_DIR,“标签”,“val”)
os.makedirs(train_img_dir,exist_ok = True)
os.makedirs(val_img_dir,exist_ok = True)
os.makedirs(train_lbl_dir,exist_ok = True)
os.makedirs(val_lbl_dir,exist_ok = True)
之后,我们需要使用 加载 XML 文件ET.parse(XML_FILE)
并获取其根元素。每个<image>
标签对应一张带注释的图像,因此我将它们收集到 中。打印出在 XML 中找到了多少张图像后,我将它们打乱顺序(这样我们的训练和验证分割会更加随机 😁😎),然后根据 将image_elements
它们分成两个列表train_elements
和。val_elements
TRAIN_RATIO
tree = ET.parse(XML_FILE)
root = tree.getroot()
image_elements = root.findall( 'image' )
print ( "XML 中的图像数量:" , len (image_elements))
# 随机播放
image_elements = list (image_elements)
random.shuffle(image_elements)
train_count = int ( len (image_elements) * TRAIN_RATIO)
train_elements = image_elements[:train_count]
val_elements = image_elements[train_count:]
print ( f "训练图像:{ len (train_elements)} , 验证图像:{ len (val_elements)} " )
好的,现在把这些放在一起:
def process_images ( image_subset, subset_name ):
"""
对于子集中的每个图像,将图像文件复制到相应的 train/val 文件夹,
创建带有边界框的 YOLO 标签文件,并将其放在 label/train 或 label/val 中。
"""
if subset_name == "train" :
img_out_dir = train_img_dir
lbl_out_dir = train_lbl_dir
else :
img_out_dir = val_img_dir
lbl_out_dir = val_lbl_dir
for img_elem in image_subset:
file_name = img_elem.attrib[ 'name' ]
width = float (img_elem.attrib[ 'width' ])
height = float (img_elem.attrib[ 'height' ])
src_img_path = os.path.join(IMAGES_DIR, file_name)
dst_img_path = os.path.join(img_out_dir, file_name)
if not os.path.exists(src_img_path):
print ( f"Warning: {src_img_path} not found. Skipping." )
continue
# 复制图像
shutil.copy2(src_img_path, dst_img_path)
# 准备标签线
boxes = img_elem.findall( 'box' )
label_lines = []
for b in boxes:
label_str = b.attrib[ 'label' ]
xtl = float (b.attrib[ 'xtl' ])
ytl = float (b.attrib[ 'ytl' ])
xbr = float (b.attrib[ 'xbr' ])
ybr = float (b.attrib[ 'ybr' ])
# 将标签转换为类索引
if label_str not in CLASS_MAP:
print ( f"Warning: Label ' {label_str} ' 不在 CLASS_MAP 中。跳过。" )
继续
class_idx = CLASS_MAP[label_str]
x_center, y_center, w, h = convert_to_yolo_bbox(
xtl, ytl, xbr, ybr, width, height
)
label_line = f" {class_idx} {x_center: .6 f} {y_center: .6 f} {w: .6 f} {h: .6f} "
label_lines.append(label_line)
txt_file_name = os.path.splitext(file_name)[ 0 ] + ".txt"
txt_out_path = os.path.join(lbl_out_dir, txt_file_name)
with open (txt_out_path, "w" ) as f:
for line in label_lines:
f.write(line + "\n" )
process_images(train_elements, "train" )
process_images(val_elements, "val" )
print ( "图像和标签处理完成。" )
这里,我定义了一个名为 的函数process_images
,它遍历image_elements
(训练部分或验证部分)的子集。它将每个图像文件复制到正确的文件夹,然后循环遍历<box>
该图像的元素以检索边界框坐标。它还会检查 是否label_str
已启用,CLASS_MAP
以便将类名(例如“red”)转换为其 YOLO 整数 ID(例如 0)。之后,它调用convert_to_yolo_bbox
来获取 YOLO 格式的边界框,并最终将这些行写入.txt
与图像同名的文件中。
最后,由于 YOLOv8 需要一个 YAML 文件,该文件至少包含两个字段“train”和“val”,指向包含图像的目录,以及一个类名列表,我们需要dataset.yaml
通过写出训练和验证集的绝对路径来动态创建,然后指定“names”数组,其中包括按正确顺序排列的“red”、“blue”、“yellow”。
sorted_classes = sorted(CLASS_MAP.items(),key= lambda x:x[ 1 ])
data_yaml_path = os.path.join(OUTPUT_DIR,“dataset.yaml”)
with open(data_yaml_path,“w”)as f:
f.write(f“train:{os.path.abspath(train_img_dir)} \n”)
f.write(f“val:{os.path.abspath(val_img_dir)} \n”)
f.write(“names:\n”)
for k,v in sorted_classes:
f.write(f“ {v}:{k} \n”)
print(f“dataset.yaml created at {data_yaml_path} ”)
model = YOLO('yolov8s.pt')
results = model.train(
data=data_yaml_path,
epochs= 1000,
imgsz = 640,
batch = 32,
name = “robotics_model”
)
print(“训练完成。检查'runs / detect / robotics_model'以获取日志和权重。”)
该文件准备就绪后,我们可以从其较小的预训练权重中加载 YOLOv8 模型'yolov8s.pt'
,并model.train
使用数据集路径调用 epoch 数。您也可以选择 Yolov8 的纳米、中或大尺寸,但这样我获得了良好的性能,所以我没有使用更大的模型。我们将针对图像大小、批量大小和自定义运行名称(“robotics_model”)进行训练。Yolo 的优点还在于它使用了早期停止技术,因此如果您不确定应该训练模型多少个 epoch,您可以给出一个像 1000 这样的大数字,因为它已经可以提前停止,防止过度消耗资源而没有改进🤖💪。训练需要一些时间,具体取决于您的 GPU 或 CPU。完成后,YOLO 会将日志、结果和最佳模型权重保存在 内的文件夹中runs/detect/
,这很神奇,因为我们不需要自己绘制所有这些指标🥳🥳。
现在,让我们看看我们的模型的结果:
我们可以看到,混淆矩阵表明,总体而言,该模型正确区分了三种颜色(红色、蓝色和黄色)以及“背景”类别,预测落在对角线上(13 个正确分类为红色,7 个正确分类为蓝色,12 个正确分类为黄色)。
同时,训练曲线显示,训练和验证集的边界框、类别和 DFL 损失均随时间稳步下降,验证集损失在第 40-60 轮和第 80-90 轮左右偶尔出现峰值。这些峰值可能反映了不稳定或过拟合的情况,但损失最终会稳定在较低值。精确度和召回率指标(mAP50 和 mAP50-95)收敛到较高水平,表明一旦模型从波动中恢复,它就能在检测和正确分类三种颜色的物体方面取得强劲的表现。综合起来,这些指标表明最终模型在区分红色、蓝色、黄色和背景方面既准确又相当稳健,这意味着在对图像进行一些预测后,我们最终可以使用 DeepSORT 😻😻 集成物体追踪功能。
为了进行预测,我们可以加载最佳检查点(通常保存在名称为 下best.pt
)。然后,您可以调用model.predict
所需的任何图像或图像文件夹。通过指定conf=0.7
,我们会滤除置信度低于 70% 的检测结果,您也可以将其提高到 0.8 或任何阈值。此设置save=True
可确保 YOLO 生成的输出图像带有围绕检测到的样本绘制的边界框,并标注其置信度和类别标签。结果默认保存在“runs/predict”文件夹中。

完成 YOLO 训练过程并准备好最佳模型权重后,我们自然希望在视频中随时间推移跟踪检测到的物体。这时 Deep SORT 便应运而生:虽然 YOLO 能够出色地显示物体在每一帧中的位置,但它无法随着时间的推移为每个物体保持一致的身份。Deep SORT 解决了这个问题,它为物体分配唯一的 ID,并预测它们如何逐帧移动,即使它们重叠或部分超出视野范围。这在机器人技术领域非常有用,我们需要检测并跟踪不同的物体,并前往其位置将其拾取。
然而,您在 GitHub 上找到的默认Deep SORT 代码通常针对旧版本的 scikit-learn、SciPy 和 TensorFlow 进行了配置。这种不匹配导致了我遇到的这些错误,我想让您知道,这样您就不会遇到这些类型的错误:“没有名为 的模块sklearn.utils.linear_assignment_
”、“元组索引必须是整数或切片,而不是元组”以及“模块‘tensorflow’没有属性‘Session’”。
下面,我们将看到每个错误是如何产生的,以及如何通过修改几行来修复它们。
由于我的本地机器上没有专用 GPU,因此我在 Google Colab 中执行了这些步骤。我简单地上传了一个压缩目录,其中包含我自定义的 Deep SORT 代码、best.pt
模型权重、一个vid.mp4
包含散落在地板上的颜色样本的测试视频 ( ),以及基于 MARS 的外观编码器(您可以在我的 GitHub 上找到)。这个设置让我可以直接在云端运行 YOLO 和 Deep SORT。
旧版 Deep SORT 代码中的三个主要问题是由于库的变更造成的。首先,新版本的 scikit-learn 删除了sklearn.utils.linear_assignment_
,所以我们必须切换到,这是在 deep_sort 文件夹内的linear_assignment.py文件from scipy.optimize import linear_sum_assignment as linear_assignment
中导入的。
其次,SciPy 现在(row_indices, col_indices)
在调用时返回一个元组linear_sum_assignment
,而旧版 Deep SORT 需要一个匹配对的数组;为了解决这个问题,我们通过堆叠两个数组,这也是文件linear_assignment.py的一个变化:
row_ind,col_ind = linear_assignment(cost_matrix)
索引= np.stack((row_ind,col_ind),轴= 1)
为了更容易理解,我给出更新后的linear_assignment.py代码:
# vim: expandtab:ts=4:sw=4
从__future__导入absolute_import
导入numpy作为np
从scipy.optimize导入linear_sum_assignment作为linear_assignment
从.导入kalman_filter
INFTY_COST = 1e+5
def min_cost_matching (
distance_metric, max_distance, tracks,detections, track_indices= None ,
detection_indices= None ):
"""解决线性分配问题。
参数
----------
distance_metric : Callable[List[Track], List[Detection], List[int], List[int]) -> ndarray
距离度量给定一个轨迹和检测列表,以及
一个包含 N 个轨迹索引和 M 个检测索引的列表。该度量应
返回 NxM 维成本矩阵,其中元素 (i, j) 是
给定轨迹索引中第 i 个轨迹与
给定detection_indices中第 j 个检测之间的关联成本。max_distance
: float
门控阈值。成本大于此值的关联将被
忽略。tracks
: List[track.Track]
当前时间步长的预测轨迹列表。detections
: List[detection.Detection]
当前时间步长的检测列表。
track_indices : List[int]
将 `cost_matrix` 中的行映射到 `tracks` 中的轨道的轨道索引列表
(参见上面的描述)。detection_indices
: List[int]
将 `cost_matrix` 中的列映射到
`detections` 中的检测的检测索引列表(参见上面的描述)。
返回
-------
(List[(int, int)], List[int], List[int])
返回包含以下三个条目的元组:
* 匹配的轨道和检测索引列表。
* 不匹配的轨道索引列表。
* 不匹配的检测索引列表。
"""
如果track_indices为 None:
track_indices = np.arange( len (tracks))
如果detection_indices为 None:
detection_indices = np.arange( len (detections))
如果 len (detection_indices) == 0 或 len(track_indices) == 0 :
返回[], track_indices,detection_indices # 没有匹配项。
cost_matrix = distance_metric(
轨迹, 检测, track_indices,detection_indices)
cost_matrix[cost_matrix > max_distance] = max_distance + 1e-5
row_ind,col_ind = linear_assignment(cost_matrix)
索引 = np.stack((row_ind,col_ind),axis= 1 )
匹配,unmatched_tracks,unmatched_detections = [],[],[]
对于col,枚举中的 detection_idx (detection_indices):如果col不在索引[:,1 ]中: unmatched_detections.append(detection_idx)对于行,枚举中的track_idx (track_indices):如果行不在索引[:,0 ]中: unmatched_tracks.append(track_idx)对于行,索引中的col : track_idx = track_indices[row] detection_idx =detection_indices[col] if cost_matrix[row, col] > max_distance: unmatched_tracks.append(track_idx) unmatched_detections.append(detection_idx) else : matches.append((track_idx,detection_idx)) return matches, unmatched_tracks, unmatched_detections def matching_cascade ( distance_metric, max_distance, cascade_depth, tracks,detections, track_indices= None ,detection_indices= None ): """运行匹配级联。 参数 ---------- distance_metric : Callable[List[Track], List[Detection], List[int], List[int]) -> ndarray 距离度量给出了轨迹和检测列表以及 N 个轨迹索引和 M 个检测索引的列表。该指标应 返回 NxM 维成本矩阵,其中元素 (i, j) 表示 给定轨道索引中第 i 个轨道与 给定检测索引中第 j 个检测之间的关联成本。max_distance :浮点型 门控阈值。成本大于此值的关联将被 忽略。cascade_depth :int 级联深度,应设置为最大轨道年龄。tracks :List[track.Track]
当前时间步长的预测轨迹列表。detections
:List[detection.Detection]
当前时间步长的检测列表。track_indices
:Optional[List[int]]
将 `cost_matrix` 中的行映射到 `tracks` 中的轨迹的轨迹索引列表
(参见上面的说明)。默认为所有轨迹。detection_indices
:Optional[List[int]]将 `cost_matrix` 中的列映射到 `detections` 中的检测的
检测索引列表(参见上面的说明)。默认为所有 检测。 返回 ------- (List[(int, int)], List[int], List[int]) 返回包含以下三个条目的元组: * 匹配的轨迹和检测索引列表。 * 不匹配的轨迹索引列表。 * 不匹配的检测索引列表。 """
如果track_indices为 None:
track_indices = list(range(len( tracks ) ) )
如果detection_indices为 None:
detection_indices = list(range(len(detections ) ) )
unmatched_detections = detection_indices
matches = []
for level in range (cascade_depth):
如果 len (unmatched_detections) == 0: # 没有剩余检测
break
track_indices_l = [
k for k in track_indices
if tracks[k].time_since_update == 1 + level
]
如果 len (track_indices_l) == 0: # 此级别没有可匹配的内容
continue
matches_l, _, unmatched_detections = \
min_cost_matching(
distance_metric, max_distance, tracks, Detections,
track_indices_l, unmatched_detections)
matches += matches_l
unmatched_tracks = list(set (track_indices) -设置(k为k,_在匹配中)
返回匹配、不匹配的轨迹、不匹配的检测
def gate_cost_matrix(
kf、cost_matrix、tracks、detections、track_indices、detection_indices、
gated_cost=INFTY_COST、only_position= False ):
"""根据
卡尔曼滤波获得的状态分布,使成本矩阵中不可行的条目无效。
参数
----------
kf:卡尔曼滤波器。cost_matrix
:ndarray
NxM 维成本矩阵,其中 N 是轨迹索引的数量
,M 是检测索引的数量,使得条目 (i, j) 是
`tracks[track_indices[i]]` 和
`detections[detection_indices[j]]` 之间的关联成本。tracks
:List[track.Track]
当前时间步长的预测轨迹列表。detections
:List[detection.Detection]
当前时间步长的检测列表。track_indices
:List[int]
将 `cost_matrix` 中的行映射到轨迹的轨迹索引列表在
`tracks`中(参见上面的描述)。detection_indices
:List[int]
检测索引列表,将`cost_matrix`中的列映射到
`detections`中的检测(参见上面的描述)。gated_cost
:Optional[float]
成本矩阵中与不可行关联相对应的条目
设置为此值。默认为一个非常大的值。only_position
:Optional[bool]
如果为 True,则在门控期间仅考虑状态分布的 x、y 位置
。默认为 False。
返回
-------
ndarray
返回修改后的成本矩阵。
"""
gating_dim = 2 if only_position else 4
gating_threshold = kalman_filter.chi2inv95[gating_dim]
measurements = np.asarray(
[detections[i].to_xyah() for i indetection_indices ])
for row, track_idx in enumerate (track_indices):
track = tracks[track_idx]
gating_distance = kf.gating_distance(
track.mean, track.covariance, measurements, only_position)
cost_matrix[row, gating_distance > gating_threshold] = gated_cost
return cost_matrix
最后, deep_sort 代码库tools文件夹中的generate_detections.py文件出现了“no attribute 'Session'”错误,这是因为原始 Deep SORT 代码引用了,而 TensorFlow 1.x 的功能。在 TensorFlow 2.x 中,sessions 被 Eager Execution 取代,因此我们启用了兼容模式:tf.Session()
导入tensorflow.compat.v1作为tf
tf.disable_v2_behavior()
您可以在导入之后、_run_in_batches 函数之前添加此函数。
完成这些更改后,Deep SORT 可以继续使用旧的基于会话的方法,不会出现任何问题。如果您不想自己修改代码库,可以克隆我已修复的版本(您可以在我的 GitHub 上找到),将其压缩,然后上传到 Google Colab。我就是这么做的:上传后deep_sort.zip
,我运行了程序!unzip deep_sort.zip
,一切就绪了。现在,让我们来看看我们的检测和跟踪代码,它精确地展示了 YOLO 和 Deep SORT 如何在测试视频上协同工作 😵😀。
!解压缩deep_sort.zip
导入cv2
导入os
导入随机
导入numpy作为np
从ultralytics导入YOLO
从deep_sort.deep_sort.tracker导入Tracker作为DeepSortTracker
从deep_sort.tools导入generate_detections作为gdet
从deep_sort.deep_sort导入nn_matching
从deep_sort.deep_sort.detection导入Detection
打印(“导入完成。”)
好的,导入完要使用的库和函数后,让我们创建两个类来管理追踪。DeepSortWrapper
构造函数会加载基于 MARS 的外观编码器(model_data/mars-small128.pb
),并设置一个具有特定最大余弦距离的追踪器(max_cosine_distance=0.4
),这会影响基于外观的追踪匹配的严格程度。每次调用 时update(frame, detections)
,它都会处理来自 YOLO 的任何新检测,预测现有追踪可能移动到的位置,并更新它们的边界框。之后,它会重新填充self.tracks
当前已确认的追踪,以便我们可以将它们绘制在视频帧上。通过将此逻辑打包到包装器中,我们可以保持主循环简洁,并将注意力集中在 YOLO 和最终的绘制例程上。
什么是 MARS 以及它在 DeepSORT 中如何使用?
在 DeepSORT 中,基于 MARS 的外观编码器从物体检测器提供的裁剪检测结果中提取高维特征向量(嵌入向量),在本例中,YOLOv8 负责这项工作。这些嵌入向量代表每个物体的独特外观,并与卡尔曼滤波器的运动信息一起使用,将新的检测结果与现有的轨迹关联起来。该编码器通过使用余弦相似度或类似的距离度量来比较嵌入向量,帮助 DeepSORT 跨帧甚至在临时遮挡或场景重新进入后保持物体身份。
类 Track:
def __init__(self,track_id,bbox):
self.track_id = track_id
self.bbox = bbox
类 DeepSortWrapper:
def __init __ (self,model_filename = 'model_data / mars-small128.pb',max_cosine_distance = 0.4,nn_budget = None):
metric = nn_matching.NearestNeighborDistanceMetric(“cosine”,max_cosine_distance,nn_budget)
self.tracker = DeepSortTracker(metric)
self.encoder = gdet.create_box_encoder(model_filename,batch_size = 1)
self.tracks = []
def update(self,frame,detections):
# 步骤1:如果没有检测,则使用空列表运行预测更新循环。
if len (detections) == 0 :
self.tracker.predict()
self.tracker.update([])
self._update_tracks()
return
# 步骤2:将[x1, y1, x2, y2]转换为[x, y, w, h]以进行Deep SORT
bboxes = np.array([d[: 4 ] for d indetections ])
scores = [d[ 4 ] for d indetections ]
bboxes[:, 2 :] = bboxes[:, 2 :] - bboxes[:, : 2 ]
# 步骤3:为每个边界框生成外观特征
features = self.encoder(frame, bboxes)
# 步骤4:将所有内容包装在Deep SORT的Detection对象中
dets = []
for bbox_id, bbox in enumerate (bboxes):
dets.append(Detection(bbox, scores[bbox_id], features[bbox_id]))
# 步骤 5:预测并更新跟踪器
self.tracker.predict()
self.tracker.update(dets)
self._update_tracks()
def _update_tracks ( self ):
active_tracks = []
for track in self.tracker.tracks:
if not track.is_confirmed() or track.time_since_update > 1 :
continue
bbox = track.to_tlbr() # 返回 [x1, y1, x2, y2]
track_id = track.track_id
active_tracks.append(Track(track_id,bbox))
self.tracks = active_tracks
现在,让我们加载模型best.pt
并创建一个实例DeepSortWrapper
。这基本上是设置阶段,确认我们的检测和跟踪组件已准备好用于任何输入。
yolo_model_path = "best.pt"
model = YOLO(yolo_model_path)
deepsort = DeepSortWrapper(
model_filename= 'mars-small128.pb' ,
max_cosine_distance= 0.4 ,
nn_budget= None
)
print ( "YOLO 模型和 Deep SORT 包装器已初始化。" )
最后,让我们介绍一下在视频中进行样本跟踪的最终主循环:
首先,我们vid.mp4
用 OpenCV 打开一个名为 的视频文件,并收集一些基本元数据,例如宽度、高度和每秒帧数。然后,我们设置一个输出文件,output_with_detections.mp4
将处理后的帧写入其中,这样我们就可以获得包含检测结果的视频了,是不是很酷💪😻🥳!然后,对每一帧进行如下操作:
- 我们调用
model(frame)
来获取YOLO检测结果。 - 我们会解析每个检测结果,并根据指定的置信度阈值(本例中为 0.8)进行过滤。对于每个符合条件的检测结果,我们会将其边界框坐标
[x1, y1, x2, y2]
与检测分数一起以 格式存储。我们还会跟踪检测到的红色、蓝色或黄色物体数量,并生成可选的叠加文本。 - 我们将所有这些边界框交给
deepsort.update(frame, all_detections)
,它将它们转换为[x, y, w, h]
格式,运行外观编码器,并更新跟踪器,以便我们知道哪些检测对应于现有轨道。 - 我们迭代
deepsort.tracks
绘制一个矩形,并为每个活动轨道添加一个标签(“ID: track_id”)。为了使边界框在视觉上有所区分,我们生成一个随机颜色列表,并根据轨道 ID 选择一种颜色。 - 我们将检测到的每个类别的当前计数叠加在左上角。
- 最后,我们将带注释的帧写入
out
VideoWriter,以便保存最终结果。
当视频播放结束时,我们会释放资源,并print(f"Video saved as {output_file}")
确认已获得完整处理的输出。最终生成的视频会为每个颜色样本显示边界框,并绘制一致的 ID,这些 ID 将始终与每个唯一对象绑定,直到该对象离开帧或以其他方式丢失。
总的来说,这个流程对于机器人或自动化任务来说非常强大,因为检测物体只是任务的一半,另一半是持续可靠地跟踪它们。YOLOv8 的速度和 Deep SORT 的强大跟踪能力相结合,打造出一个精简的系统,能够适应各种用例,从颜色编码样本检测(比如我的)到人员跟踪或库存管理解决方案,应有尽有。希望你能将这个流程用于你自己的创意项目,期待下篇文章再见,祝你编程愉快🤓👩💻👏💪!