本篇文章是对之前YOLOv4 tensorrt推理项目的更新，在trt推理，多进程语音报警的功能中又新添加了目标跟踪【可用于目标计数】
采用IOU进行跟踪【也可以选用DIOU】。
yolov4+deepsort可以参考我另一篇：yolov4+deepsort
说明：
本项目做的是单目标的跟踪，可以实现对特定类别的跟踪。iou目标跟踪效果还是比较差。适用于场景变换不大，物体移动速度慢的场景，只是检测速度会比deepsort的要快一些，如果对跟踪效果要求比较高，还是建议采用deepsort。
后期会更新CIOU进行目标跟踪。
利用iou进行跟踪是比较简单的，也是目标检测中最为常用的手段，但与目标检测还是不一样，在目标检测的iou过滤中【指预测部分】，是将大于阈值的目标过滤掉，保留小于阈值的，在训练阶段是将target box与output的box iou大的保留下来表示正样本的匹配。
而在目标跟踪中，是将目标前一帧和后一帧的box计算iou，如果大于阈值说明是同一目标。
iou对box的处理可以参考我另一篇文章：采用iou过滤预测框
这里直接附上利用iou做跟踪的代码。
主要是返回两个值，
一个是target_dict:包含了目标的id以及box信息【x1,y1,x2,y2,score,iou的值】。
另一个是target_index:目标的ID,可以用来做计数。
当判断大于iou阈值的时候会将信息记录在reg_dict中，并设置is_track为True.如果是新目标的时候，target_index + 1.代码中的输入box_data值的是当前帧【也就是model的输出结果】，target_dict中的box信息是前一帧的信息，用这两个box计算iou。
代码注解我也都标注出来了。
'''
采用iou进行目标的跟踪，对前一帧和当前帧iou判断是否为同一目标
'''
import copy
def compute_iou(box_a, box_b, diou=True):
S_rec1 = (box_a[2] - box_a[0]) * (box_a[3] - box_a[1])
S_rec2 = (box_b[2] - box_b[0]) * (box_b[3] - box_b[1])
# 并集面积
sum_area = S_rec1 + S_rec2
left_X = max(box_a[1], box_b[1]) # 相交的左上角
left_Y = max(box_a[0], box_b[0])
right_X = min(box_a[3], box_b[3]) # 相交的右下角
right_Y = min(box_a[2], box_b[2])
# ----------------用来算diou内容--------------------
center_a_x = (box_a[3] - box_a[1]) / 2 + box_a[1]
center_a_y = (box_a[2] - box_a[0]) / 2 + box_a[0]
center_b_x = (box_b[3] - box_b[1]) / 2 + box_b[1]
center_b_y = (box_b[2] - box_b[0]) / 2 + box_b[0]
# 计算两个框中心点的欧式距离
center_distance = (center_b_x-center_a_x)**2 + (center_b_y-center_a_y) ** 2
# 计算两个框最小矩形的左上角和右下角
closebox_min_x = min(box_a[1], box_b[1])
closebox_min_y = min(box_a[0], box_b[0])
closebox_max_x = max(box_a[3], box_b[3])
closebox_max_y = max(box_a[2], box_b[2])
# 计算两个框最小矩形的对角线距离
closebox_distance = (closebox_min_x-closebox_max_x)**2 + (closebox_min_y- closebox_max_y)**2
# -------------------------------------------------
if left_X >= right_X or left_Y >= right_Y:
return 0.
else:
# 相交的面积
inter_area = (right_X - left_X) * (right_Y - left_Y)
iou = inter_area / (sum_area - inter_area)
return iou - (center_distance/closebox_distance) if diou else iou
def tracking_box(box_data, target_dict, track_index, iou_thre=0.8):
reg_dict = {} # 用来记录识别id以及信息
# box_data是列表形式
for bbox in box_data: # box_data会保存目标坐标，每个元素bbox是元组的形式，包含内容(x1.y1,x2,y2,score)
xa0, ya0, xa1, ya1, score = bbox # 当前帧目标的box
is_tracke = False # 状态初始化
for k_ in target_dict.keys():
xb0, yb0, xb1, yb1, s, _ = target_dict[k_] # 前一帧的box
iou_ = compute_iou((ya0, xa0, ya1, xa1), (yb0, xb0, yb1, xb1))
if iou_ > iou_thre: # 跟踪目标成功
reg_dict[k_] = (xa0, ya0, xa1, ya1, score, iou_)
is_tracke = True # 跟踪√
# print('iou_', iou_)
if not is_tracke: # 表示新的目标
reg_dict[track_index] = (xa0, ya0, xa1, ya1, score, 0.) # 记录目标id以及信息
track_index += 1
if track_index >= 65535: # 越界归零
track_index = 0
if track_index >= 100:
track_index = 0
target_dict = copy.deepcopy(reg_dict) # 记录target信息
# print(target_dict)
return target_dict, track_index
def target_tracking(box_data, target_dict, track_index):
target_dict, track_index = tracking_box(box_data, target_dict, track_index)
return target_dict, track_index
代码：
git clone --branch v3.0 https://github.com/YINYIPENG-EN/YOLOv4_tensorRT_pytorch.git
python predict.py --weights [weight path] --video --video_path 0 --engine --track --track_class person
本文内容由网友自发贡献，版权归原作者所有，本站不承担相应法律责任。如您发现有涉嫌抄袭侵权的内容，请联系:hwhale#tublm.com(使用前将#替换为@)
