# Copyright (c) OpenMMLab. All rights reserved. from typing import List, Optional, Tuple import numpy as np import torch from mmengine.structures import InstanceData try: import motmetrics from motmetrics.lap import linear_sum_assignment except ImportError: motmetrics = None from torch import Tensor from mmdet.registry import MODELS, TASK_UTILS from mmdet.structures import DetDataSample from mmdet.structures.bbox import bbox_overlaps, bbox_xyxy_to_cxcyah from mmdet.utils import OptConfigType from ..utils import imrenormalize from .base_tracker import BaseTracker @MODELS.register_module() class SORTTracker(BaseTracker): """Tracker for SORT/DeepSORT. Args: obj_score_thr (float, optional): Threshold to filter the objects. Defaults to 0.3. motion (dict): Configuration of motion. Defaults to None. reid (dict, optional): Configuration for the ReID model. - num_samples (int, optional): Number of samples to calculate the feature embeddings of a track. Default to 10. - image_scale (tuple, optional): Input scale of the ReID model. Default to (256, 128). - img_norm_cfg (dict, optional): Configuration to normalize the input. Default to None. - match_score_thr (float, optional): Similarity threshold for the matching process. Default to 2.0. match_iou_thr (float, optional): Threshold of the IoU matching process. Defaults to 0.7. num_tentatives (int, optional): Number of continuous frames to confirm a track. Defaults to 3. """ def __init__(self, motion: Optional[dict] = None, obj_score_thr: float = 0.3, reid: dict = dict( num_samples=10, img_scale=(256, 128), img_norm_cfg=None, match_score_thr=2.0), match_iou_thr: float = 0.7, num_tentatives: int = 3, **kwargs): if motmetrics is None: raise RuntimeError('motmetrics is not installed,\ please install it by: pip install motmetrics') super().__init__(**kwargs) if motion is not None: self.motion = TASK_UTILS.build(motion) assert self.motion is not None, 'SORT/Deep SORT need KalmanFilter' self.obj_score_thr = obj_score_thr self.reid = reid self.match_iou_thr = match_iou_thr self.num_tentatives = num_tentatives @property def confirmed_ids(self) -> List: """Confirmed ids in the tracker.""" ids = [id for id, track in self.tracks.items() if not track.tentative] return ids def init_track(self, id: int, obj: Tuple[Tensor]) -> None: """Initialize a track.""" super().init_track(id, obj) self.tracks[id].tentative = True bbox = bbox_xyxy_to_cxcyah(self.tracks[id].bboxes[-1]) # size = (1, 4) assert bbox.ndim == 2 and bbox.shape[0] == 1 bbox = bbox.squeeze(0).cpu().numpy() self.tracks[id].mean, self.tracks[id].covariance = self.kf.initiate( bbox) def update_track(self, id: int, obj: Tuple[Tensor]) -> None: """Update a track.""" super().update_track(id, obj) if self.tracks[id].tentative: if len(self.tracks[id]['bboxes']) >= self.num_tentatives: self.tracks[id].tentative = False bbox = bbox_xyxy_to_cxcyah(self.tracks[id].bboxes[-1]) # size = (1, 4) assert bbox.ndim == 2 and bbox.shape[0] == 1 bbox = bbox.squeeze(0).cpu().numpy() self.tracks[id].mean, self.tracks[id].covariance = self.kf.update( self.tracks[id].mean, self.tracks[id].covariance, bbox) def pop_invalid_tracks(self, frame_id: int) -> None: """Pop out invalid tracks.""" invalid_ids = [] for k, v in self.tracks.items(): # case1: disappeared frames >= self.num_frames_retrain case1 = frame_id - v['frame_ids'][-1] >= self.num_frames_retain # case2: tentative tracks but not matched in this frame case2 = v.tentative and v['frame_ids'][-1] != frame_id if case1 or case2: invalid_ids.append(k) for invalid_id in invalid_ids: self.tracks.pop(invalid_id) def track(self, model: torch.nn.Module, img: Tensor, data_sample: DetDataSample, data_preprocessor: OptConfigType = None, rescale: bool = False, **kwargs) -> InstanceData: """Tracking forward function. Args: model (nn.Module): MOT model. img (Tensor): of shape (T, C, H, W) encoding input image. Typically these should be mean centered and std scaled. The T denotes the number of key images and usually is 1 in SORT method. data_sample (:obj:`TrackDataSample`): The data sample. It includes information such as `pred_det_instances`. data_preprocessor (dict or ConfigDict, optional): The pre-process config of :class:`TrackDataPreprocessor`. it usually includes, ``pad_size_divisor``, ``pad_value``, ``mean`` and ``std``. rescale (bool, optional): If True, the bounding boxes should be rescaled to fit the original scale of the image. Defaults to False. Returns: :obj:`InstanceData`: Tracking results of the input images. Each InstanceData usually contains ``bboxes``, ``labels``, ``scores`` and ``instances_id``. """ metainfo = data_sample.metainfo bboxes = data_sample.pred_instances.bboxes labels = data_sample.pred_instances.labels scores = data_sample.pred_instances.scores frame_id = metainfo.get('frame_id', -1) if frame_id == 0: self.reset() if not hasattr(self, 'kf'): self.kf = self.motion if self.with_reid: if self.reid.get('img_norm_cfg', False): img_norm_cfg = dict( mean=data_preprocessor['mean'], std=data_preprocessor['std'], to_bgr=data_preprocessor['rgb_to_bgr']) reid_img = imrenormalize(img, img_norm_cfg, self.reid['img_norm_cfg']) else: reid_img = img.clone() valid_inds = scores > self.obj_score_thr bboxes = bboxes[valid_inds] labels = labels[valid_inds] scores = scores[valid_inds] if self.empty or bboxes.size(0) == 0: num_new_tracks = bboxes.size(0) ids = torch.arange( self.num_tracks, self.num_tracks + num_new_tracks, dtype=torch.long).to(bboxes.device) self.num_tracks += num_new_tracks if self.with_reid: crops = self.crop_imgs(reid_img, metainfo, bboxes.clone(), rescale) if crops.size(0) > 0: embeds = model.reid(crops, mode='tensor') else: embeds = crops.new_zeros((0, model.reid.head.out_channels)) else: ids = torch.full((bboxes.size(0), ), -1, dtype=torch.long).to(bboxes.device) # motion self.tracks, costs = self.motion.track(self.tracks, bbox_xyxy_to_cxcyah(bboxes)) active_ids = self.confirmed_ids if self.with_reid: crops = self.crop_imgs(reid_img, metainfo, bboxes.clone(), rescale) embeds = model.reid(crops, mode='tensor') # reid if len(active_ids) > 0: track_embeds = self.get( 'embeds', active_ids, self.reid.get('num_samples', None), behavior='mean') reid_dists = torch.cdist(track_embeds, embeds) # support multi-class association track_labels = torch.tensor([ self.tracks[id]['labels'][-1] for id in active_ids ]).to(bboxes.device) cate_match = labels[None, :] == track_labels[:, None] cate_cost = (1 - cate_match.int()) * 1e6 reid_dists = (reid_dists + cate_cost).cpu().numpy() valid_inds = [list(self.ids).index(_) for _ in active_ids] reid_dists[~np.isfinite(costs[valid_inds, :])] = np.nan row, col = linear_sum_assignment(reid_dists) for r, c in zip(row, col): dist = reid_dists[r, c] if not np.isfinite(dist): continue if dist <= self.reid['match_score_thr']: ids[c] = active_ids[r] active_ids = [ id for id in self.ids if id not in ids and self.tracks[id].frame_ids[-1] == frame_id - 1 ] if len(active_ids) > 0: active_dets = torch.nonzero(ids == -1).squeeze(1) track_bboxes = self.get('bboxes', active_ids) ious = bbox_overlaps(track_bboxes, bboxes[active_dets]) # support multi-class association track_labels = torch.tensor([ self.tracks[id]['labels'][-1] for id in active_ids ]).to(bboxes.device) cate_match = labels[None, active_dets] == track_labels[:, None] cate_cost = (1 - cate_match.int()) * 1e6 dists = (1 - ious + cate_cost).cpu().numpy() row, col = linear_sum_assignment(dists) for r, c in zip(row, col): dist = dists[r, c] if dist < 1 - self.match_iou_thr: ids[active_dets[c]] = active_ids[r] new_track_inds = ids == -1 ids[new_track_inds] = torch.arange( self.num_tracks, self.num_tracks + new_track_inds.sum(), dtype=torch.long).to(bboxes.device) self.num_tracks += new_track_inds.sum() self.update( ids=ids, bboxes=bboxes, scores=scores, labels=labels, embeds=embeds if self.with_reid else None, frame_ids=frame_id) # update pred_track_instances pred_track_instances = InstanceData() pred_track_instances.bboxes = bboxes pred_track_instances.labels = labels pred_track_instances.scores = scores pred_track_instances.instances_id = ids return pred_track_instances