📚 课程目录第4章多类别目标跟踪与类别过滤

4. 多类别目标跟踪与类别过滤

任务目标

在YOLO与DeepSORT框架中，默认跟踪器对所有检测目标统一分配ID。本案例要求实现：对检测结果中的多个类别（如person、car、bicycle）分别维护独立的跟踪状态，并支持通过配置参数过滤，仅对指定类别执行跟踪与ID分配，未选中的类别不参与跟踪。

核心实现思路

DeepSORT的update接口接收检测框及特征，本身不区分目标类别。实现多类别跟踪的关键是：按类别拆分检测结果，为每个类别创建独立的DeepSORT跟踪器实例。类别过滤则在送入跟踪器之前完成——只将允许的类别检测框传入对应跟踪器，其余类别直接丢弃。这样每个跟踪器内部的ID空间相互独立，同一ID可出现在不同类别中，符合多类别跟踪的语义。

关键完整代码

import numpy as np
from deep_sort import DeepSort
from collections import defaultdict

class MultiClassDeepSort:
    """多类别跟踪器：每个类别独立维护一个DeepSORT实例，支持类别过滤"""
    def __init__(self, model_path, max_cos_dist=0.3, max_age=30, n_init=3,
                 allowed_classes=None, device='cpu'):
        """
        Args:
            model_path: DeepSORT特征提取模型权重路径
            allowed_classes: 允许跟踪的类别列表，如 ['person', 'car']；None表示全部
        """
        self.allowed_classes = set(allowed_classes) if allowed_classes else None
        self.trackers = {}  # class_name -> DeepSort实例
        self.model_path = model_path
        self.max_cos_dist = max_cos_dist
        self.max_age = max_age
        self.n_init = n_init
        self.device = device

    def _get_tracker(self, class_name):
        """按需创建或获取某类别对应的DeepSORT实例"""
        if class_name not in self.trackers:
            self.trackers[class_name] = DeepSort(
                self.model_path,
                max_cos_dist=self.max_cos_dist,
                max_age=self.max_age,
                n_init=self.n_init,
                nn_budget=None,
                device=self.device
            )
        return self.trackers[class_name]

    def update(self, detections, class_names, ori_img):
        """
        Args:
            detections: np.ndarray, shape (N, 5)，每行为 [x1, y1, x2, y2, conf]
            class_names: list[str], 长度N，每个检测框的类别名
            ori_img: 原始帧图像(BGR)
        Returns:
            results: list[dict]，每个元素为
                {'bbox': [x1,y1,x2,y2], 'track_id': int, 'class': str, 'conf': float}
        """
        # 1. 按类别分组检测结果
        grouped = defaultdict(list)
        for bbox, cls in zip(detections, class_names):
            # 类别过滤：不在允许列表中的类别直接跳过
            if self.allowed_classes is not None and cls not in self.allowed_classes:
                continue
            grouped[cls].append(bbox)

        # 2. 每个类别独立更新跟踪器
        results = []
        for cls, cls_dets in grouped.items():
            tracker = self._get_tracker(cls)
            cls_dets = np.array(cls_dets)
            # DeepSORT要求输入为 [x1,y1,x2,y2,conf]，且conf为0~1
            tracked = tracker.update(cls_dets, ori_img)
            for t in tracked:
                results.append({
                    'bbox': [int(v) for v in t[:4]],
                    'track_id': int(t[4]),
                    'class': cls,
                    'conf': float(t[5]) if len(t) > 5 else 1.0
                })
        return results

使用方式

# 初始化：仅跟踪person和car两类
tracker = MultiClassDeepSort(
    model_path='ckpt.t7',
    allowed_classes=['person', 'car'],
    device='cuda'
)

# 每帧调用（以YOLO检测输出为例）
# detections: (N,5) 格式 [x1,y1,x2,y2,conf]
# class_names: 与detections行对应的类别字符串列表
results = tracker.update(detections, class_names, frame_bgr)

# results中只包含person和car的跟踪结果，其他类别被过滤
for r in results:
    print(f"ID:{r['track_id']} 类别:{r['class']} 框:{r['bbox']}")

验证方法

类别过滤验证：输入包含person、car、bicycle三类检测框，设置allowed_classes=['person']，观察输出结果中应只出现person类别的跟踪ID，bicycle和car的检测框被完全丢弃。
多类别独立ID验证：不设置过滤（allowed_classes=None），同时输入person和car。若画面中同时存在一个person和一个car，两者应各自获得ID（如person-1、car-1），且互不干扰。当person离开再出现时，其ID可能变化，但car的ID不受person进出影响。
运行通过条件：程序无报错，输出结果中每个跟踪目标的类别与输入检测类别一致；过滤后仅保留指定类别；同一类别内同一目标在连续帧中ID保持稳定（除非目标丢失超过max_age帧）。

关键说明

每个类别独立使用一个DeepSORT实例意味着内存占用随类别数线性增长，实际部署时建议限制allowed_classes数量。
若某个类别在某一帧无检测框，则跳过该帧的更新，跟踪器内部会自行处理轨迹预测与超时删除。
DeepSORT的update返回的tracked数组格式为[x1,y1,x2,y2,track_id,conf]，本代码已做解析。