Spaces:
Runtime error
Runtime error
File size: 4,343 Bytes
3b96cb1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 |
# Copyright (c) OpenMMLab. All rights reserved.
from typing import Optional
from torch import Tensor
from mmdet.registry import MODELS
from mmdet.structures import TrackSampleList
from mmdet.utils import OptConfigType
from .base import BaseMOTModel
@MODELS.register_module()
class DeepSORT(BaseMOTModel):
"""Simple online and realtime tracking with a deep association metric.
Details can be found at `DeepSORT<https://arxiv.org/abs/1703.07402>`_.
Args:
detector (dict): Configuration of detector. Defaults to None.
reid (dict): Configuration of reid. Defaults to None
tracker (dict): Configuration of tracker. Defaults to None.
data_preprocessor (dict or ConfigDict, optional): The pre-process
config of :class:`TrackDataPreprocessor`. it usually includes,
``pad_size_divisor``, ``pad_value``, ``mean`` and ``std``.
init_cfg (dict or list[dict]): Configuration of initialization.
Defaults to None.
"""
def __init__(self,
detector: Optional[dict] = None,
reid: Optional[dict] = None,
tracker: Optional[dict] = None,
data_preprocessor: OptConfigType = None,
init_cfg: OptConfigType = None):
super().__init__(data_preprocessor, init_cfg)
if detector is not None:
self.detector = MODELS.build(detector)
if reid is not None:
self.reid = MODELS.build(reid)
if tracker is not None:
self.tracker = MODELS.build(tracker)
self.preprocess_cfg = data_preprocessor
def loss(self, inputs: Tensor, data_samples: TrackSampleList,
**kwargs) -> dict:
"""Calculate losses from a batch of inputs and data samples."""
raise NotImplementedError(
'Please train `detector` and `reid` models firstly, then \
inference with SORT/DeepSORT.')
def predict(self,
inputs: Tensor,
data_samples: TrackSampleList,
rescale: bool = True,
**kwargs) -> TrackSampleList:
"""Predict results from a video and data samples with post- processing.
Args:
inputs (Tensor): of shape (N, T, C, H, W) encoding
input images. The N denotes batch size.
The T denotes the number of key frames
and reference frames.
data_samples (list[:obj:`TrackDataSample`]): The batch
data samples. It usually includes information such
as `gt_instance`.
rescale (bool, Optional): If False, then returned bboxes and masks
will fit the scale of img, otherwise, returned bboxes and masks
will fit the scale of original image shape. Defaults to True.
Returns:
TrackSampleList: List[TrackDataSample]
Tracking results of the input videos.
Each DetDataSample usually contains ``pred_track_instances``.
"""
assert inputs.dim() == 5, 'The img must be 5D Tensor (N, T, C, H, W).'
assert inputs.size(0) == 1, \
'SORT/DeepSORT inference only support ' \
'1 batch size per gpu for now.'
assert len(data_samples) == 1, \
'SORT/DeepSORT inference only support ' \
'1 batch size per gpu for now.'
track_data_sample = data_samples[0]
video_len = len(track_data_sample)
if track_data_sample[0].frame_id == 0:
self.tracker.reset()
for frame_id in range(video_len):
img_data_sample = track_data_sample[frame_id]
single_img = inputs[:, frame_id].contiguous()
# det_results List[DetDataSample]
det_results = self.detector.predict(single_img, [img_data_sample])
assert len(det_results) == 1, 'Batch inference is not supported.'
pred_track_instances = self.tracker.track(
model=self,
img=single_img,
feats=None,
data_sample=det_results[0],
data_preprocessor=self.preprocess_cfg,
rescale=rescale,
**kwargs)
img_data_sample.pred_track_instances = pred_track_instances
return [track_data_sample]
|