deepsort_tracker.py 中文注释版代码如下:
import time
import logging
from collections.abc import Iterable
import cv2
import numpy as np
from deep_sort_realtime.deep_sort import nn_matching
from deep_sort_realtime.deep_sort.detection import Detection
from deep_sort_realtime.deep_sort.tracker import Tracker
from deep_sort_realtime.utils.nms import non_max_suppression
logger = logging.getLogger(__name__) # 创建一个日志记录器实例,用于记录 DeepSort 跟踪器的日志信息。`__self__` 表示当前模块的名称。
EMBEDDER_CHOICES = [ # 定义一个列表,包含可用于 DeepSort 跟踪器的特征嵌入模型的名称。
"mobilenet", # 轻量级的深度学习模型,常用于特征嵌入。
"torchreid", # 基于 PyTorch 的 ReID(行人重识别)模型。
"clip_RN50", # CLIP 模型的一个变体,使用 ResNet-50 作为骨干网络。
"clip_RN101", # CLIP 模型的一个变体,使用 ResNet-101 作为骨干网络。
"clip_RN50x4", # CLIP 模型的一个变体,使用 4 倍 ResNet-50 作为骨干网络。
"clip_RN50x16", # CLIP 模型的一个变体,使用 16 倍 ResNet-50 作为骨干网络。
"clip_ViT-B/32", # CLIP 模型的一个变体,使用 Vision Transformer (ViT) B/32 作为骨干网络。
"clip_ViT-B/16", # CLIP 模型的一个变体,使用 Vision Transformer (ViT) B/16 作为骨干网络。
] # 这些嵌入模型可以用于提取跟踪对象的特征向量,以便进行有效的目标跟踪和识别。
class DeepSort(object):
"""
class DeepSort(object):
Parameters
----------
max_iou_distance : Optional[float] = 0.7
# 最大IoU距离:设置IoU(交并比)的最大阈值。大于这个值的关联成本将被忽略。这是 deep_sort_realtime.deep_sort.tracker.Tracker 的参数。
max_age : Optional[int] = 30
# 最大年龄:跟踪在丢失多少次帧后才被删除的最大次数。这是 deep_sort_realtime.deep_sort.tracker.Tracker 的参数。
n_init : int
# 初始化帧数:跟踪在初始化阶段保持的帧数,默认为3。这是 deep_sort_realtime.deep_sort.tracker.Tracker 的参数。
nms_max_overlap : Optional[float] = 1.0
# NMS最大重叠:非极大值抑制(NMS)的阈值。如果设置为1.0,将禁用NMS。
max_cosine_distance : Optional[float] = 0.2
# 最大余弦距离:余弦距离的最大门限值。
nn_budget : Optional[int] = None
# 特征描述符的最大大小,如果为None,则不强制执行预算。
gating_only_position : Optional[bool]
# 在门限过程中仅使用位置:在门限过程中,比较卡尔曼滤波器(KF)预测和测量状态时使用。如果为True,则仅考虑状态分布的x, y位置。默认为False,会考虑x,y,宽高比和高度。
override_track_class : Optional[object] = None
# 覆盖跟踪类:提供这个参数将覆盖默认的跟踪类,必须继承Track类。这是 deep_sort_realtime.deep_sort.tracker.Tracker 的参数。
embedder : Optional[str] = 'mobilenet'
# 嵌入器:是否使用内置的嵌入器。如果为None,则在更新时必须提供嵌入特征。可选的嵌入器有 'mobilenet', 'torchreid', 'clip_RN50', 'clip_RN101', 'clip_RN50x4', 'clip_RN50x16', 'clip_ViT-B/32', 'clip_ViT-B/16'。
half : Optional[bool] = True
# 使用半精度:是否对深度嵌入器使用半精度(仅适用于mobilenet)。
bgr : Optional[bool] = True
# 是否期望嵌入器接收的帧为BGR格式,否则为RGB。
embedder_gpu: Optional[bool] = True
# 嵌入器是否使用GPU。
embedder_model_name: Optional[str] = None
# 仅当embedder=='torchreid'时使用。这提供了在torchreid库中使用哪个模型。请查看torchreid的模型库。
embedder_wts: Optional[str] = None
# 嵌入器模型权重的路径的可选指定。默认情况下,会在 `deep_sort_realtime/embedder/weights` 目录下寻找权重。如果deep_sort_realtime作为包安装,并且使用CLIP模型作为嵌入器,最好提供路径。
polygon: Optional[bool] = False
# 是否检测多边形(例如,定向边界框)。
today: Optional[datetime.date]
# 提供今天的日期,用于跟踪的命名。这是 deep_sort_realtime.deep_sort.tracker.Tracker 的参数。"""
def __init__( # 构造函数,用于初始化 DeepSort 对象
self,
max_iou_distance=0.7, # 最大IoU距离,用于关联检测框和跟踪对象
max_age=30, # 最大年龄,跟踪对象在被删除前最多能错过的帧数
n_init=3, # 初始化帧数,新跟踪对象在成为确认跟踪前需要经过的帧数
nms_max_overlap=1.0, # NMS最大重叠阈值,用于抑制非最大值抑制过程中的重叠检测框
max_cosine_distance=0.2, # 最大余弦距离,用于特征匹配时的相似度度量
nn_budget=None, # 特征匹配时的预算,限制考虑的特征数量
gating_only_position=False, # 在门限过程中是否只考虑位置信息
override_track_class=None, # 用于覆盖默认的跟踪类
embedder="mobilenet", # 用于特征嵌入的模型
half=True, # 是否使用半精度浮点数以加快计算速度
bgr=True, # 输入图像是否为BGR格式
embedder_gpu=True, # 嵌入器是否在GPU上运行
embedder_model_name=None, # 嵌入器模型的名称
embedder_wts=None, # 嵌入器模型权重的路径
polygon=False, # 是否处理多边形检测框
today=None, # 当前日期,用于跟踪对象命名
):
# 初始化函数体
self.nms_max_overlap = nms_max_overlap # 设置NMS最大重叠阈值
metric = nn_matching.NearestNeighborDistanceMetric( # 创建最近邻距离度量对象
"cosine", max_cosine_distance, nn_budget
)
self.tracker = Tracker( # 创建Tracker对象
metric, # 使用的距离度量
max_iou_distance=max_iou_distance, # 设置最大IoU距离
max_age=max_age, # 设置最大年龄
n_init=n_init, # 设置初始化帧数
override_track_class=override_track_class, # 可能覆盖默认跟踪类
today=today, # 当前日期
gating_only_position=gating_only_position, # 是否仅在门限过程中使用位置信息
)
if embedder is not None: # 如果提供了嵌入器
if embedder not in EMBEDDER_CHOICES: # 检查嵌入器是否有效
raise Exception(f"Embedder {embedder} is not a valid choice.") # 如果无效,抛出异常
if embedder == "mobilenet": # 如果嵌入器是mobilenet
from deep_sort_realtime.embedder.embedder_pytorch import ( # 导入MobileNetv2_Embedder
MobileNetv2_Embedder as Embedder,
)
self.embedder = Embedder( # 创建MobileNetv2_Embedder对象
half=half, # 是否使用半精度
max_batch_size=16, # 最大批量大小
bgr=bgr, # 输入图像是否为BGR格式
gpu=embedder_gpu, # 是否在GPU上运行
model_wts_path=embedder_wts, # 模型权重路径
)
elif embedder == 'torchreid': # 如果嵌入器是torchreid
from deep_sort_realtime.embedder.embedder_pytorch import TorchReID_Embedder as Embedder
self.embedder = Embedder( # 创建TorchReID嵌入器对象
bgr=bgr, # 输入图像是否为BGR格式
gpu=embedder_gpu, # 是否在GPU上运行
model_name=embedder_model_name, # 模型名称
model_wts_path=embedder_wts, # 模型权重路径
)
elif embedder.startswith('clip_'): # 如果嵌入器是CLIP模型之一
from deep_sort_realtime.embedder.embedder_clip import (
Clip_Embedder as Embedder,
)
model_name = "_".join(embedder.split("_")[1:]) # 获取CLIP模型的名称
self.embedder = Embedder( # 创建CLIP嵌入器对象
model_name=model_name, # 模型名称
model_wts_path=embedder_wts, # 模型权重路径
max_batch_size=16, # 最大批量大小
bgr=bgr, # 输入图像是否为BGR格式
gpu=embedder_gpu, # 是否在GPU上运行
)
else: # 如果没有提供嵌入器
self.embedder = None # 嵌入器对象设置为None
self.polygon = polygon # 是否处理多边形检测框
logger.info("DeepSort Tracker initialised") # 记录日志,DeepSort跟踪器已初始化
logger.info(f"- max age: {max_age}") # 记录日志,最大年龄
logger.info(f"- appearance threshold: {max_cosine_distance}") # 记录日志,外观阈值
logger.info( # 记录日志,NMS阈值
f'- nms threshold: {"OFF" if self.nms_max_overlap == 1.0 else self.nms_max_overlap}'
)
logger.info(f"- max num of appearance features: {nn_budget}") # 记录日志,最大特征数量
logger.info( # 记录日志,是否覆盖跟踪类
f'- overriding track class : {"No" if override_track_class is None else "Yes"}'
)
logger.info( # 记录日志,是否提供了今天的日期
f'- today given : {"No" if today is None else "Yes"}'
)
logger.info( # 记录日志,是否使用了内置嵌入器
f'- in-build embedder : {"No" if self.embedder is None else "Yes"}'
)
logger.info( # 记录日志,是否处理多边形检测框
f'- polygon detections : {"No" if polygon is False else "Yes"}'
)
"""Run multi-target tracker on a particular sequence.
Parameters
----------
raw_detections (horizontal bb) : List[ Tuple[ List[float or int], float, str ] ]
List of detections, each in tuples of ( [left,top,w,h] , confidence, detection_class)
raw_detections (polygon) : List[ List[float], List[int or str], List[float] ]
List of Polygons, Classes, Confidences. All 3 sublists of the same length. A polygon defined as a ndarray-like [x1,y1,x2,y2,...].
embeds : Optional[ List[] ] = None
List of appearance features corresponding to detections
frame : Optional [ np.ndarray ] = None
if embeds not given, Image frame must be given here, in [H,W,C].
today: Optional[datetime.date]
Provide today's date, for naming of tracks
others: Optional[ List ] = None
Other things associated to detections to be stored in tracks, usually, could be corresponding segmentation mask, other associated values, etc. Currently others is ignored with polygon is True.
instance_masks: Optional [ List ] = None
Instance masks corresponding to detections. If given, they are used to filter out background and only use foreground for apperance embedding. Expects numpy boolean mask matrix.
Returns
-------
list of track objects (Look into track.py for more info or see "main" section below in this script to see simple example)
"""
def update_tracks(self, raw_detections, embeds=None, frame=None, today=None, others=None, instance_masks=None):
""" 在特定序列上运行多目标跟踪器
参数
----------
raw_detections (水平边界框) : List[ Tuple[ List[float or int], float, str ] ]
检测结果列表,格式为 ([left,top,w,h] , confidence, detection_class)
raw_detections (多边形) : List[ List[float], List[int or str], List[float] ]
多边形、类别、置信度的列表。所有3个子列表长度相同。多边形定义为类似 ndarray 的 [x1,y1,x2,y2,...]。
embeds : Optional[ List[] ] = None
与检测结果对应的外观特征列表
frame : Optional [ np.ndarray ] = None
如果未提供外观特征列表,则必须在此处提供图像帧,格式为 [H,W,C]。
today: Optional[datetime.date]
提供今天的日期,用于命名跟踪结果
others: Optional[ List ] = None
与检测结果相关的其他信息,通常是对应的分割掩模、其他相关值等。当多边形为 True 时,当前忽略其他信息。
instance_masks: Optional [ List ] = None
与检测结果对应的实例掩模。如果提供,则用于过滤背景并仅使用前景进行外观特征提取。期望是 numpy 的布尔掩模矩阵。
返回
-------
轨迹对象列表(查看 track.py 获取更多信息,或查看此脚本下方的 "main" 部分以查看简单示例)
"""
if embeds is None: # 如果未提供外观特征列表
if self.embedder is None: # 如果在初始化时未创建外观特征提取器
raise Exception(
"Embedder not created during init so embeddings must be given now!"
)
if frame is None: # 如果未提供外观特征列表或图像帧
raise Exception("either embeddings or frame must be given!")
assert isinstance(raw_detections, Iterable) # 确保raw_detections是可迭代对象
if len(raw_detections) > 0: # 如果有检测结果
if not self.polygon: # 如果不是多边形检测
assert len(raw_detections[0][0]) == 4 # 确保检测结果的bbox格式为[x,y,w,h]
raw_detections = [d for d in raw_detections if d[0][2] > 0 and d[0][3] > 0] # 过滤掉宽度或高度小于等于0的检测结果
if embeds is None: # 如果未提供外观特征列表
embeds = self.generate_embeds(frame, raw_detections, instance_masks=instance_masks) # 生成外观特征列表
# Proper deep sort detection objects that consist of bbox, confidence and embedding.
detections = self.create_detections(raw_detections, embeds, instance_masks=instance_masks, others=others) # 创建深度排序的检测对象,包括bbox、置信度和外观特征
else: # 如果是多边形检测
polygons, bounding_rects = self.process_polygons(raw_detections[0]) # 处理多边形检测结果
if embeds is None: # 如果未提供外观特征列表
embeds = self.generate_embeds_poly(frame, polygons, bounding_rects) # 生成多边形检测的外观特征列表
# 创建深度排序的多边形检测对象,包括bbox、置信度和外观特征
detections = self.create_detections_poly(
raw_detections, embeds, bounding_rects,
)
else: # 如果没有检测结果
detections = []
# Run non-maxima suppression.
boxes = np.array([d.ltwh for d in detections]) # 提取检测结果的bbox
scores = np.array([d.confidence for d in detections]) # 提取检测结果的置信度
if self.nms_max_overlap < 1.0: # 如果非极大值抑制的最大重叠阈值小于1.0
# nms_tic = time.perf_counter()
indices = non_max_suppression(boxes, self.nms_max_overlap, scores) # 执行非极大值抑制
# nms_toc = time.perf_counter()
# logger.debug(f'nms time: {nms_toc-nms_tic}s')
detections = [detections[i] for i in indices] # 根据非极大值抑制的结果过滤检测结果
# Update tracker.
self.tracker.predict() # 预测目标轨迹
self.tracker.update(detections, today=today) # 更新目标轨迹
return self.tracker.tracks # 返回更新后的目标轨迹列表
def refresh_track_ids(self):
self.tracker._next_id
def generate_embeds(self, frame, raw_dets, instance_masks=None):
# 裁剪检测框内的图像区域
crops, cropped_inst_masks = self.crop_bb(frame, raw_dets, instance_masks=instance_masks)
# 如果存在实例掩码,则将掩码应用到裁剪的图像区域上
if cropped_inst_masks is not None:
masked_crops = []
for crop, mask in zip(crops, cropped_inst_masks):
# 创建一个与裁剪的图像区域相同大小的零矩阵
masked_crop = np.zeros_like(crop)
# 将图像区域上的像素值减去均值(123.675, 116.28, 103.53)
masked_crop = masked_crop + np.array([123.675, 116.28, 103.53], dtype=crop.dtype)
# 将掩码应用到裁剪的图像区域上
masked_crop[mask] = crop[mask]
masked_crops.append(masked_crop)
# 使用嵌入提取器预测裁剪后的图像区域的外观特征
return self.embedder.predict(masked_crops)
else:
# 如果不存在实例掩码,则直接使用嵌入提取器预测裁剪后的图像区域的外观特征
return self.embedder.predict(crops)
def generate_embeds_poly(self, frame, polygons, bounding_rects):
"""
使用多边形检测结果生成特征嵌入。
参数
----------
frame : numpy.ndarray
当前帧的图像数据。
polygons : List[List[float]]
多边形检测结果的坐标列表。
bounding_rects : List[List[float]]
多边形检测结果对应的边界矩形坐标列表。
返回
------
embeddings : numpy.ndarray
预测得到的特征嵌入。
"""
crops = self.crop_poly_pad_black(frame, polygons, bounding_rects) # 从帧中裁剪并填充多边形区域
return self.embedder.predict(crops) # 使用嵌入器对裁剪的区域进行预测,得到特征嵌入
def create_detections(self, raw_dets, embeds, instance_masks=None, others=None):
"""
根据原始检测结果和特征嵌入创建检测对象列表。
参数
----------
raw_dets : List[Tuple[float]]
原始检测结果,包含边界框、置信度和类别。
embeds : List[numpy.ndarray]
对应的特征嵌入。
instance_masks : Optional[List[numpy.ndarray]]
实例掩码列表,用于指示每个检测的区域。
others : Optional[List[Any]]
其他附加信息,如跟踪ID等。
返回
------
detections : List[Detection]
创建的检测对象列表。
"""
detection_list = []
for i, (raw_det, embed) in enumerate(zip(raw_dets, embeds)):
# 创建检测对象并添加到列表中
detection_list.append(
Detection(
raw_det[0], # 边界框
raw_det[1], # 置信度
embed, # 特征嵌入
class_name=raw_det[2] if len(raw_det) == 3 else None, # 类别名称
instance_mask=instance_masks[i] if isinstance(instance_masks, Iterable) else instance_masks,
others=others[i] if isinstance(others, Iterable) else others,
)
)
return detection_list # 返回创建的检测对象列表
def create_detections_poly(self, dets, embeds, bounding_rects):
"""
根据多边形检测结果和特征嵌入创建检测对象列表。
参数
----------
dets : List[Tuple[float]]
包含多边形、类别和置信度的检测结果。
embeds : List[numpy.ndarray]
对应的特征嵌入。
bounding_rects : List[List[float]]
多边形检测结果对应的边界矩形坐标列表。
返回
------
detections : List[Detection]
创建的检测对象列表。
"""
detection_list = []
dets.extend([embeds, bounding_rects]) # 扩展检测结果列表以包含特征嵌入和边界矩形
for raw_polygon, cl, score, embed, bounding_rect in zip(*dets):
# 从边界矩形中提取坐标
x, y, w, h = bounding_rect
x = max(0, x) # 确保x不小于0
y = max(0, y) # 确保y不小于0
bbox = [x, y, w, h] # 创建边界框列表
# 创建检测对象并添加到列表中
detection_list.append(
Detection(bbox, score, embed, class_name=cl, others=raw_polygon)
)
return detection_list # 返回创建的检测对象列表
@staticmethod
def process_polygons(raw_polygons):
"""
处理多边形检测结果,计算它们的边界矩形并返回多边形和边界矩形列表。
参数
----------
raw_polygons : List[List[float]]
原始多边形检测结果的坐标列表。
返回
------
polygons : List[List[float]]
处理后的多边形坐标列表。
bounding_rects : List[List[float]]
对应的边界矩形坐标列表。
"""
polygons = [
[polygon[x: x + 2] for x in range(0, len(polygon), 2)]
for polygon in raw_polygons
]
# 计算每个多边形的边界矩形
bounding_rects = [
cv2.boundingRect(np.array([polygon]).astype(int)) for polygon in polygons
]
return polygons, bounding_rects
@staticmethod
def crop_bb(frame, raw_dets, instance_masks=None):
"""
从图像帧中裁剪出边界框区域,并可选地裁剪出实例掩码。
参数
----------
frame : numpy.ndarray
图像帧。
raw_dets : List[Tuple[float]]
原始边界框检测结果。
instance_masks : Optional[List[numpy.ndarray]]
实例掩码列表。
返回
------
crops : List[numpy.ndarray]
裁剪后的图像区域列表。
masks : List[numpy.ndarray], optional
裁剪后的实例掩码列表。
"""
crops = []
im_height, im_width = frame.shape[:2]
if instance_masks is not None:
masks = []
else:
masks = None
for i, detection in enumerate(raw_dets):
# 解析边界框坐标
l, t, w, h = [int(x) for x in detection[0]]
r = l + w
b = t + h
# 裁剪图像区域
crop_l = max(0, l)
crop_r = min(im_width, r)
crop_t = max(0, t)
crop_b = min(im_height, b)
crops.append(frame[crop_t:crop_b, crop_l:crop_r])
if instance_masks is not None:
masks.append(instance_masks[i][crop_t:crop_b, crop_l:crop_r])
return crops, masks
@staticmethod
def crop_poly_pad_black(frame, polygons, bounding_rects):
"""
从图像帧中裁剪出多边形区域,并将周围填充为黑色。
参数
----------
frame : numpy.ndarray
图像帧。
polygons : List[List[float]]
多边形坐标列表。
bounding_rects : List[List[float]]
多边形对应的边界矩形坐标列表。
返回
------
masked_polys : List[numpy.ndarray]
裁剪并填充后的图像区域列表。
"""
masked_polys = []
im_height, im_width = frame.shape[:2]
for polygon, bounding_rect in zip(polygons, bounding_rects):
mask = np.zeros(frame.shape, dtype=np.uint8)
polygon_mask = np.array([polygon]).astype(int)
# 在掩码上填充多边形
cv2.fillPoly(mask, polygon_mask, color=(255, 255, 255))
# 应用掩码到图像
masked_image = cv2.bitwise_and(frame, mask)
# 裁剪并填充的图像区域
x, y, w, h = bounding_rect
crop_l = max(0, x)
crop_r = min(im_width, x + w)
crop_t = max(0, y)
crop_b = min(im_height, y + h)
cropped = masked_image[crop_t:crop_b, crop_l:crop_r].copy()
masked_polys.append(np.array(cropped))
return masked_polys
def delete_all_tracks(self):
"""
删除跟踪器中的所有跟踪对象。
返回
------
无
"""
self.tracker.delete_all_tracks()
# 调用跟踪器的删除所有跟踪的方法