Spaces:
Build error
Build error
File size: 6,072 Bytes
28c256d |
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 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 |
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
from abc import ABCMeta, abstractmethod
from typing import Dict, List, Tuple, Union
from mmengine.model import BaseModel
from torch import Tensor
from mmdet.registry import MODELS
from mmdet.structures import OptTrackSampleList, TrackSampleList
from mmdet.utils import OptConfigType, OptMultiConfig
@MODELS.register_module()
class BaseMOTModel(BaseModel, metaclass=ABCMeta):
"""Base class for multiple object tracking.
Args:
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]): Initialization config dict.
"""
def __init__(self,
data_preprocessor: OptConfigType = None,
init_cfg: OptMultiConfig = None) -> None:
super().__init__(
data_preprocessor=data_preprocessor, init_cfg=init_cfg)
def freeze_module(self, module: Union[List[str], Tuple[str], str]) -> None:
"""Freeze module during training."""
if isinstance(module, str):
modules = [module]
else:
if not (isinstance(module, list) or isinstance(module, tuple)):
raise TypeError('module must be a str or a list.')
else:
modules = module
for module in modules:
m = getattr(self, module)
m.eval()
for param in m.parameters():
param.requires_grad = False
@property
def with_detector(self) -> bool:
"""bool: whether the framework has a detector."""
return hasattr(self, 'detector') and self.detector is not None
@property
def with_reid(self) -> bool:
"""bool: whether the framework has a reid model."""
return hasattr(self, 'reid') and self.reid is not None
@property
def with_motion(self) -> bool:
"""bool: whether the framework has a motion model."""
return hasattr(self, 'motion') and self.motion is not None
@property
def with_track_head(self) -> bool:
"""bool: whether the framework has a track_head."""
return hasattr(self, 'track_head') and self.track_head is not None
@property
def with_tracker(self) -> bool:
"""bool: whether the framework has a tracker."""
return hasattr(self, 'tracker') and self.tracker is not None
def forward(self,
inputs: Dict[str, Tensor],
data_samples: OptTrackSampleList = None,
mode: str = 'predict',
**kwargs):
"""The unified entry for a forward process in both training and test.
The method should accept three modes: "tensor", "predict" and "loss":
- "tensor": Forward the whole network and return tensor or tuple of
tensor without any post-processing, same as a common nn.Module.
- "predict": Forward and return the predictions, which are fully
processed to a list of :obj:`TrackDataSample`.
- "loss": Forward and return a dict of losses according to the given
inputs and data samples.
Note that this method doesn't handle neither back propagation nor
optimizer updating, which are done in the :meth:`train_step`.
Args:
inputs (Dict[str, Tensor]): of shape (N, T, C, H, W)
encoding input images. Typically these should be mean centered
and std scaled. The N denotes batch size. The T denotes the
number of key/reference frames.
- img (Tensor) : The key images.
- ref_img (Tensor): The reference images.
data_samples (list[:obj:`TrackDataSample`], optional): The
annotation data of every samples. Defaults to None.
mode (str): Return what kind of value. Defaults to 'predict'.
Returns:
The return type depends on ``mode``.
- If ``mode="tensor"``, return a tensor or a tuple of tensor.
- If ``mode="predict"``, return a list of :obj:`TrackDataSample`.
- If ``mode="loss"``, return a dict of tensor.
"""
if mode == 'loss':
return self.loss(inputs, data_samples, **kwargs)
elif mode == 'predict':
return self.predict(inputs, data_samples, **kwargs)
elif mode == 'tensor':
return self._forward(inputs, data_samples, **kwargs)
else:
raise RuntimeError(f'Invalid mode "{mode}". '
'Only supports loss, predict and tensor mode')
@abstractmethod
def loss(self, inputs: Dict[str, Tensor], data_samples: TrackSampleList,
**kwargs) -> Union[dict, tuple]:
"""Calculate losses from a batch of inputs and data samples."""
pass
@abstractmethod
def predict(self, inputs: Dict[str, Tensor], data_samples: TrackSampleList,
**kwargs) -> TrackSampleList:
"""Predict results from a batch of inputs and data samples with post-
processing."""
pass
def _forward(self,
inputs: Dict[str, Tensor],
data_samples: OptTrackSampleList = None,
**kwargs):
"""Network forward process. Usually includes backbone, neck and head
forward without any post-processing.
Args:
inputs (Dict[str, Tensor]): of shape (N, T, C, H, W).
data_samples (List[:obj:`TrackDataSample`], optional): The
Data Samples. It usually includes information such as
`gt_instance`.
Returns:
tuple[list]: A tuple of features from ``head`` forward.
"""
raise NotImplementedError(
"_forward function (namely 'tensor' mode) is not supported now")
|