用Python和OpenCV复现SORT算法从卡尔曼滤波预测到匈牙利匹配的完整代码解读在计算机视觉领域多目标跟踪(Multi-Object Tracking, MOT)是一个极具挑战性的任务。SORT(Simple Online and Realtime Tracking)算法以其简洁高效的特点成为该领域的经典基准方法。本文将带您从零开始用Python和OpenCV完整实现SORT算法深入解析卡尔曼滤波预测与匈牙利匹配的代码细节并分享实际开发中的调试经验。1. 环境准备与核心组件实现SORT算法需要以下Python库支持import numpy as np import cv2 from scipy.optimize import linear_sum_assignment from filterpy.kalman import KalmanFilter关键组件说明numpy处理矩阵运算opencv图像处理和可视化scipy.optimize提供匈牙利算法实现filterpy简化卡尔曼滤波器实现提示建议使用Python 3.8环境可通过pip install filterpy scipy opencv-python安装依赖2. 卡尔曼滤波器实现卡尔曼滤波是SORT算法的核心预测组件我们首先实现一个针对边界框跟踪的定制化卡尔曼滤波器class KalmanBoxTracker: count 0 # 类变量用于ID分配 def __init__(self, bbox): # 初始化7维状态向量[u,v,s,r,u,v,s] self.kf KalmanFilter(dim_x7, dim_z4) self.kf.F np.array([[1,0,0,0,1,0,0], [0,1,0,0,0,1,0], [0,0,1,0,0,0,1], [0,0,0,1,0,0,0], [0,0,0,0,1,0,0], [0,0,0,0,0,1,0], [0,0,0,0,0,0,1]]) self.kf.H np.array([[1,0,0,0,0,0,0], [0,1,0,0,0,0,0], [0,0,1,0,0,0,0], [0,0,0,1,0,0,0]]) # 初始化协方差矩阵 self.kf.P[4:,4:] * 1000 self.kf.P * 10 self.kf.Q[-1,-1] * 0.01 self.kf.Q[4:,4:] * 0.01 self.kf.x[:4] self.convert_bbox_to_z(bbox) self.time_since_update 0 self.id KalmanBoxTracker.count KalmanBoxTracker.count 1 self.history [] self.hits 0 self.hit_streak 0 self.age 0状态向量设计要点u,v边界框中心坐标s边界框面积r宽高比u,v,s对应变量的变化率3. 匈牙利匹配与IOU计算匹配检测框与预测框的关键步骤是计算IOU(交并比)矩阵并应用匈牙利算法进行最优分配def iou_batch(bb_test, bb_gt): 计算两组边界框之间的IOU矩阵 :param bb_test: 检测框 [N x 4] :param bb_gt: 预测框 [M x 4] :return: IOU矩阵 [N x M] bb_gt np.expand_dims(bb_gt, 0) bb_test np.expand_dims(bb_test, 1) xx1 np.maximum(bb_test[..., 0], bb_gt[..., 0]) yy1 np.maximum(bb_test[..., 1], bb_gt[..., 1]) xx2 np.minimum(bb_test[..., 2], bb_gt[..., 2]) yy2 np.minimum(bb_test[..., 3], bb_gt[..., 3]) w np.maximum(0., xx2 - xx1) h np.maximum(0., yy2 - yy1) intersection w * h area_test (bb_test[..., 2] - bb_test[..., 0]) * (bb_test[..., 3] - bb_test[..., 1]) area_gt (bb_gt[..., 2] - bb_gt[..., 0]) * (bb_gt[..., 3] - bb_gt[..., 1]) return intersection / (area_test area_gt - intersection 1e-10) def associate_detections_to_trackers(detections, trackers, iou_threshold0.3): 使用匈牙利算法将检测框与预测框匹配 :param detections: 当前帧检测框 :param trackers: 预测框 :param iou_threshold: 匹配阈值 :return: 匹配对、未匹配检测、未匹配预测 if len(trackers) 0: return np.empty((0,2), dtypeint), np.arange(len(detections)), np.empty((0,5), dtypeint) iou_matrix iou_batch(detections, trackers) # 使用匈牙利算法求解最优匹配 row_ind, col_ind linear_sum_assignment(-iou_matrix) matched_indices np.array(list(zip(row_ind, col_ind))) # 过滤低IOU匹配 unmatched_detections [] for d, det in enumerate(detections): if d not in matched_indices[:, 0]: unmatched_detections.append(d) unmatched_trackers [] for t, trk in enumerate(trackers): if t not in matched_indices[:, 1]: unmatched_trackers.append(t) matches [] for m in matched_indices: if iou_matrix[m[0], m[1]] iou_threshold: unmatched_detections.append(m[0]) unmatched_trackers.append(m[1]) else: matches.append(m.reshape(1,2)) if len(matches) 0: matches np.empty((0,2), dtypeint) else: matches np.concatenate(matches, axis0) return matches, np.array(unmatched_detections), np.array(unmatched_trackers)4. SORT主流程实现整合上述组件构建完整的SORT跟踪器class Sort: def __init__(self, max_age1, min_hits3, iou_threshold0.3): self.max_age max_age self.min_hits min_hits self.iou_threshold iou_threshold self.trackers [] self.frame_count 0 def update(self, detsnp.empty((0,5))): self.frame_count 1 # 获取跟踪器的预测结果 trks np.zeros((len(self.trackers), 5)) to_del [] ret [] for t, trk in enumerate(trks): pos self.trackers[t].predict()[0] trk[:] [pos[0], pos[1], pos[2], pos[3], 0] if np.any(np.isnan(pos)): to_del.append(t) trks np.ma.compress_rows(np.ma.masked_invalid(trks)) # 移除无效跟踪器 for t in reversed(to_del): self.trackers.pop(t) # 匹配检测与跟踪器预测 matched, unmatched_dets, unmatched_trks associate_detections_to_trackers(dets, trks, self.iou_threshold) # 更新匹配的跟踪器 for m in matched: self.trackers[m[1]].update(dets[m[0], :]) # 为未匹配的检测创建新跟踪器 for i in unmatched_dets: trk KalmanBoxTracker(dets[i, :]) self.trackers.append(trk) # 输出跟踪结果 i len(self.trackers) for trk in reversed(self.trackers): d trk.get_state()[0] if (trk.time_since_update 1) and (trk.hit_streak self.min_hits or self.frame_count self.min_hits): ret.append(np.concatenate((d, [trk.id1])).reshape(1,-1)) i - 1 # 移除丢失的跟踪器 if trk.time_since_update self.max_age: self.trackers.pop(i) if len(ret) 0: return np.concatenate(ret) return np.empty((0,5))5. 实际应用与调试技巧将SORT跟踪器与目标检测器结合使用时需要注意以下关键点常见问题与解决方案问题现象可能原因解决方案ID频繁切换IOU阈值过低适当提高iou_threshold(0.3-0.5)跟踪框抖动过程噪声设置不当调整卡尔曼滤波器的Q矩阵新目标无法跟踪min_hits设置过高降低min_hits或设为0目标丢失过快max_age设置过小适当增加max_age(3-5帧)性能优化建议对于高帧率视频可以降低卡尔曼滤波的预测频率使用更精确的目标检测器能显著提升跟踪效果对于特定场景可以调整状态转移矩阵F以适应不同的运动模型# 示例与YOLOv3检测器集成 def run_sort_with_yolo(): # 初始化 sort_tracker Sort(max_age5, min_hits3, iou_threshold0.3) yolo YOLOv3Detector() # 假设已实现 cap cv2.VideoCapture(input.mp4) while cap.isOpened(): ret, frame cap.read() if not ret: break # 检测目标 detections yolo.detect(frame) # 更新跟踪器 tracked_objects sort_tracker.update(detections) # 可视化结果 for obj in tracked_objects: x1,y1,x2,y2,obj_id obj cv2.rectangle(frame, (x1,y1), (x2,y2), (0,255,0), 2) cv2.putText(frame, fID:{int(obj_id)}, (x1,y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0,255,0), 2) cv2.imshow(Tracking, frame) if cv2.waitKey(1) 0xFF ord(q): break cap.release() cv2.destroyAllWindows()在实际项目中我们发现当目标遮挡严重时SORT算法的表现会显著下降。这时可以考虑升级到DeepSORT算法它通过引入外观特征匹配增强了遮挡处理能力。
用Python和OpenCV复现SORT算法：从卡尔曼滤波预测到匈牙利匹配的完整代码解读
用Python和OpenCV复现SORT算法从卡尔曼滤波预测到匈牙利匹配的完整代码解读在计算机视觉领域多目标跟踪(Multi-Object Tracking, MOT)是一个极具挑战性的任务。SORT(Simple Online and Realtime Tracking)算法以其简洁高效的特点成为该领域的经典基准方法。本文将带您从零开始用Python和OpenCV完整实现SORT算法深入解析卡尔曼滤波预测与匈牙利匹配的代码细节并分享实际开发中的调试经验。1. 环境准备与核心组件实现SORT算法需要以下Python库支持import numpy as np import cv2 from scipy.optimize import linear_sum_assignment from filterpy.kalman import KalmanFilter关键组件说明numpy处理矩阵运算opencv图像处理和可视化scipy.optimize提供匈牙利算法实现filterpy简化卡尔曼滤波器实现提示建议使用Python 3.8环境可通过pip install filterpy scipy opencv-python安装依赖2. 卡尔曼滤波器实现卡尔曼滤波是SORT算法的核心预测组件我们首先实现一个针对边界框跟踪的定制化卡尔曼滤波器class KalmanBoxTracker: count 0 # 类变量用于ID分配 def __init__(self, bbox): # 初始化7维状态向量[u,v,s,r,u,v,s] self.kf KalmanFilter(dim_x7, dim_z4) self.kf.F np.array([[1,0,0,0,1,0,0], [0,1,0,0,0,1,0], [0,0,1,0,0,0,1], [0,0,0,1,0,0,0], [0,0,0,0,1,0,0], [0,0,0,0,0,1,0], [0,0,0,0,0,0,1]]) self.kf.H np.array([[1,0,0,0,0,0,0], [0,1,0,0,0,0,0], [0,0,1,0,0,0,0], [0,0,0,1,0,0,0]]) # 初始化协方差矩阵 self.kf.P[4:,4:] * 1000 self.kf.P * 10 self.kf.Q[-1,-1] * 0.01 self.kf.Q[4:,4:] * 0.01 self.kf.x[:4] self.convert_bbox_to_z(bbox) self.time_since_update 0 self.id KalmanBoxTracker.count KalmanBoxTracker.count 1 self.history [] self.hits 0 self.hit_streak 0 self.age 0状态向量设计要点u,v边界框中心坐标s边界框面积r宽高比u,v,s对应变量的变化率3. 匈牙利匹配与IOU计算匹配检测框与预测框的关键步骤是计算IOU(交并比)矩阵并应用匈牙利算法进行最优分配def iou_batch(bb_test, bb_gt): 计算两组边界框之间的IOU矩阵 :param bb_test: 检测框 [N x 4] :param bb_gt: 预测框 [M x 4] :return: IOU矩阵 [N x M] bb_gt np.expand_dims(bb_gt, 0) bb_test np.expand_dims(bb_test, 1) xx1 np.maximum(bb_test[..., 0], bb_gt[..., 0]) yy1 np.maximum(bb_test[..., 1], bb_gt[..., 1]) xx2 np.minimum(bb_test[..., 2], bb_gt[..., 2]) yy2 np.minimum(bb_test[..., 3], bb_gt[..., 3]) w np.maximum(0., xx2 - xx1) h np.maximum(0., yy2 - yy1) intersection w * h area_test (bb_test[..., 2] - bb_test[..., 0]) * (bb_test[..., 3] - bb_test[..., 1]) area_gt (bb_gt[..., 2] - bb_gt[..., 0]) * (bb_gt[..., 3] - bb_gt[..., 1]) return intersection / (area_test area_gt - intersection 1e-10) def associate_detections_to_trackers(detections, trackers, iou_threshold0.3): 使用匈牙利算法将检测框与预测框匹配 :param detections: 当前帧检测框 :param trackers: 预测框 :param iou_threshold: 匹配阈值 :return: 匹配对、未匹配检测、未匹配预测 if len(trackers) 0: return np.empty((0,2), dtypeint), np.arange(len(detections)), np.empty((0,5), dtypeint) iou_matrix iou_batch(detections, trackers) # 使用匈牙利算法求解最优匹配 row_ind, col_ind linear_sum_assignment(-iou_matrix) matched_indices np.array(list(zip(row_ind, col_ind))) # 过滤低IOU匹配 unmatched_detections [] for d, det in enumerate(detections): if d not in matched_indices[:, 0]: unmatched_detections.append(d) unmatched_trackers [] for t, trk in enumerate(trackers): if t not in matched_indices[:, 1]: unmatched_trackers.append(t) matches [] for m in matched_indices: if iou_matrix[m[0], m[1]] iou_threshold: unmatched_detections.append(m[0]) unmatched_trackers.append(m[1]) else: matches.append(m.reshape(1,2)) if len(matches) 0: matches np.empty((0,2), dtypeint) else: matches np.concatenate(matches, axis0) return matches, np.array(unmatched_detections), np.array(unmatched_trackers)4. SORT主流程实现整合上述组件构建完整的SORT跟踪器class Sort: def __init__(self, max_age1, min_hits3, iou_threshold0.3): self.max_age max_age self.min_hits min_hits self.iou_threshold iou_threshold self.trackers [] self.frame_count 0 def update(self, detsnp.empty((0,5))): self.frame_count 1 # 获取跟踪器的预测结果 trks np.zeros((len(self.trackers), 5)) to_del [] ret [] for t, trk in enumerate(trks): pos self.trackers[t].predict()[0] trk[:] [pos[0], pos[1], pos[2], pos[3], 0] if np.any(np.isnan(pos)): to_del.append(t) trks np.ma.compress_rows(np.ma.masked_invalid(trks)) # 移除无效跟踪器 for t in reversed(to_del): self.trackers.pop(t) # 匹配检测与跟踪器预测 matched, unmatched_dets, unmatched_trks associate_detections_to_trackers(dets, trks, self.iou_threshold) # 更新匹配的跟踪器 for m in matched: self.trackers[m[1]].update(dets[m[0], :]) # 为未匹配的检测创建新跟踪器 for i in unmatched_dets: trk KalmanBoxTracker(dets[i, :]) self.trackers.append(trk) # 输出跟踪结果 i len(self.trackers) for trk in reversed(self.trackers): d trk.get_state()[0] if (trk.time_since_update 1) and (trk.hit_streak self.min_hits or self.frame_count self.min_hits): ret.append(np.concatenate((d, [trk.id1])).reshape(1,-1)) i - 1 # 移除丢失的跟踪器 if trk.time_since_update self.max_age: self.trackers.pop(i) if len(ret) 0: return np.concatenate(ret) return np.empty((0,5))5. 实际应用与调试技巧将SORT跟踪器与目标检测器结合使用时需要注意以下关键点常见问题与解决方案问题现象可能原因解决方案ID频繁切换IOU阈值过低适当提高iou_threshold(0.3-0.5)跟踪框抖动过程噪声设置不当调整卡尔曼滤波器的Q矩阵新目标无法跟踪min_hits设置过高降低min_hits或设为0目标丢失过快max_age设置过小适当增加max_age(3-5帧)性能优化建议对于高帧率视频可以降低卡尔曼滤波的预测频率使用更精确的目标检测器能显著提升跟踪效果对于特定场景可以调整状态转移矩阵F以适应不同的运动模型# 示例与YOLOv3检测器集成 def run_sort_with_yolo(): # 初始化 sort_tracker Sort(max_age5, min_hits3, iou_threshold0.3) yolo YOLOv3Detector() # 假设已实现 cap cv2.VideoCapture(input.mp4) while cap.isOpened(): ret, frame cap.read() if not ret: break # 检测目标 detections yolo.detect(frame) # 更新跟踪器 tracked_objects sort_tracker.update(detections) # 可视化结果 for obj in tracked_objects: x1,y1,x2,y2,obj_id obj cv2.rectangle(frame, (x1,y1), (x2,y2), (0,255,0), 2) cv2.putText(frame, fID:{int(obj_id)}, (x1,y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0,255,0), 2) cv2.imshow(Tracking, frame) if cv2.waitKey(1) 0xFF ord(q): break cap.release() cv2.destroyAllWindows()在实际项目中我们发现当目标遮挡严重时SORT算法的表现会显著下降。这时可以考虑升级到DeepSORT算法它通过引入外观特征匹配增强了遮挡处理能力。
