| |
| from abc import ABCMeta, abstractmethod |
| from typing import Dict, List, Tuple, Union |
|
|
| import torch |
| from mmengine.model import BaseModel |
| from torch import Tensor |
|
|
| from mmdet.structures import DetDataSample, OptSampleList, SampleList |
| from mmdet.utils import InstanceList, OptConfigType, OptMultiConfig |
| from ..utils import samplelist_boxtype2tensor |
|
|
| ForwardResults = Union[Dict[str, torch.Tensor], List[DetDataSample], |
| Tuple[torch.Tensor], torch.Tensor] |
|
|
|
|
| class BaseDetector(BaseModel, metaclass=ABCMeta): |
| """Base class for detectors. |
| |
| Args: |
| data_preprocessor (dict or ConfigDict, optional): The pre-process |
| config of :class:`BaseDataPreprocessor`. it usually includes, |
| ``pad_size_divisor``, ``pad_value``, ``mean`` and ``std``. |
| init_cfg (dict or ConfigDict, optional): the config to control the |
| initialization. Defaults to None. |
| """ |
|
|
| def __init__(self, |
| data_preprocessor: OptConfigType = None, |
| init_cfg: OptMultiConfig = None): |
| super().__init__( |
| data_preprocessor=data_preprocessor, init_cfg=init_cfg) |
|
|
| @property |
| def with_neck(self) -> bool: |
| """bool: whether the detector has a neck""" |
| return hasattr(self, 'neck') and self.neck is not None |
|
|
| |
| |
| @property |
| def with_shared_head(self) -> bool: |
| """bool: whether the detector has a shared head in the RoI Head""" |
| return hasattr(self, 'roi_head') and self.roi_head.with_shared_head |
|
|
| @property |
| def with_bbox(self) -> bool: |
| """bool: whether the detector has a bbox head""" |
| return ((hasattr(self, 'roi_head') and self.roi_head.with_bbox) |
| or (hasattr(self, 'bbox_head') and self.bbox_head is not None)) |
|
|
| @property |
| def with_mask(self) -> bool: |
| """bool: whether the detector has a mask head""" |
| return ((hasattr(self, 'roi_head') and self.roi_head.with_mask) |
| or (hasattr(self, 'mask_head') and self.mask_head is not None)) |
|
|
| def forward(self, |
| inputs: torch.Tensor, |
| data_samples: OptSampleList = None, |
| mode: str = 'tensor') -> ForwardResults: |
| """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:`DetDataSample`. |
| - "loss": Forward and return a dict of losses according to the given |
| inputs and data samples. |
| |
| Note that this method doesn't handle either back propagation or |
| parameter update, which are supposed to be done in :meth:`train_step`. |
| |
| Args: |
| inputs (torch.Tensor): The input tensor with shape |
| (N, C, ...) in general. |
| data_samples (list[:obj:`DetDataSample`], optional): A batch of |
| data samples that contain annotations and predictions. |
| Defaults to None. |
| mode (str): Return what kind of value. Defaults to 'tensor'. |
| |
| 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:`DetDataSample`. |
| - If ``mode="loss"``, return a dict of tensor. |
| """ |
| if mode == 'loss': |
| return self.loss(inputs, data_samples) |
| elif mode == 'predict': |
| return self.predict(inputs, data_samples) |
| elif mode == 'tensor': |
| return self._forward(inputs, data_samples) |
| else: |
| raise RuntimeError(f'Invalid mode "{mode}". ' |
| 'Only supports loss, predict and tensor mode') |
|
|
| @abstractmethod |
| def loss(self, batch_inputs: Tensor, |
| batch_data_samples: SampleList) -> Union[dict, tuple]: |
| """Calculate losses from a batch of inputs and data samples.""" |
| pass |
|
|
| @abstractmethod |
| def predict(self, batch_inputs: Tensor, |
| batch_data_samples: SampleList) -> SampleList: |
| """Predict results from a batch of inputs and data samples with post- |
| processing.""" |
| pass |
|
|
| @abstractmethod |
| def _forward(self, |
| batch_inputs: Tensor, |
| batch_data_samples: OptSampleList = None): |
| """Network forward process. |
| |
| Usually includes backbone, neck and head forward without any post- |
| processing. |
| """ |
| pass |
|
|
| @abstractmethod |
| def extract_feat(self, batch_inputs: Tensor): |
| """Extract features from images.""" |
| pass |
|
|
| def add_pred_to_datasample(self, data_samples: SampleList, |
| results_list: InstanceList) -> SampleList: |
| """Add predictions to `DetDataSample`. |
| |
| Args: |
| data_samples (list[:obj:`DetDataSample`], optional): A batch of |
| data samples that contain annotations and predictions. |
| results_list (list[:obj:`InstanceData`]): Detection results of |
| each image. |
| |
| Returns: |
| list[:obj:`DetDataSample`]: Detection results of the |
| input images. Each DetDataSample usually contain |
| 'pred_instances'. And the ``pred_instances`` usually |
| contains following keys. |
| |
| - scores (Tensor): Classification scores, has a shape |
| (num_instance, ) |
| - labels (Tensor): Labels of bboxes, has a shape |
| (num_instances, ). |
| - bboxes (Tensor): Has a shape (num_instances, 4), |
| the last dimension 4 arrange as (x1, y1, x2, y2). |
| """ |
| for data_sample, pred_instances in zip(data_samples, results_list): |
| data_sample.pred_instances = pred_instances |
| samplelist_boxtype2tensor(data_samples) |
| return data_samples |
|
|
