| |
| import copy |
| import warnings |
|
|
| import torch |
| from torch import Tensor |
|
|
| from mmdet.registry import MODELS |
| from mmdet.structures import SampleList |
| from mmdet.utils import ConfigType, OptConfigType, OptMultiConfig |
| from .single_stage import SingleStageDetector |
|
|
|
|
| @MODELS.register_module() |
| class RPN(SingleStageDetector): |
| """Implementation of Region Proposal Network. |
| |
| Args: |
| backbone (:obj:`ConfigDict` or dict): The backbone config. |
| neck (:obj:`ConfigDict` or dict): The neck config. |
| bbox_head (:obj:`ConfigDict` or dict): The bbox head config. |
| train_cfg (:obj:`ConfigDict` or dict, optional): The training config. |
| test_cfg (:obj:`ConfigDict` or dict, optional): The testing config. |
| data_preprocessor (:obj:`ConfigDict` or dict, optional): Config of |
| :class:`DetDataPreprocessor` to process the input data. |
| Defaults to None. |
| init_cfg (:obj:`ConfigDict` or list[:obj:`ConfigDict`] or dict or |
| list[dict], optional): Initialization config dict. |
| Defaults to None. |
| """ |
|
|
| def __init__(self, |
| backbone: ConfigType, |
| neck: ConfigType, |
| rpn_head: ConfigType, |
| train_cfg: ConfigType, |
| test_cfg: ConfigType, |
| data_preprocessor: OptConfigType = None, |
| init_cfg: OptMultiConfig = None, |
| **kwargs) -> None: |
| super(SingleStageDetector, self).__init__( |
| data_preprocessor=data_preprocessor, init_cfg=init_cfg) |
| self.backbone = MODELS.build(backbone) |
| self.neck = MODELS.build(neck) if neck is not None else None |
| rpn_train_cfg = train_cfg['rpn'] if train_cfg is not None else None |
| rpn_head_num_classes = rpn_head.get('num_classes', 1) |
| if rpn_head_num_classes != 1: |
| warnings.warn('The `num_classes` should be 1 in RPN, but get ' |
| f'{rpn_head_num_classes}, please set ' |
| 'rpn_head.num_classes = 1 in your config file.') |
| rpn_head.update(num_classes=1) |
| rpn_head.update(train_cfg=rpn_train_cfg) |
| rpn_head.update(test_cfg=test_cfg['rpn']) |
| self.bbox_head = MODELS.build(rpn_head) |
| self.train_cfg = train_cfg |
| self.test_cfg = test_cfg |
|
|
| def loss(self, batch_inputs: Tensor, |
| batch_data_samples: SampleList) -> dict: |
| """Calculate losses from a batch of inputs and data samples. |
| |
| Args: |
| batch_inputs (Tensor): Input images of shape (N, C, H, W). |
| These should usually be mean centered and std scaled. |
| batch_data_samples (list[:obj:`DetDataSample`]): The batch |
| data samples. It usually includes information such |
| as `gt_instance` or `gt_panoptic_seg` or `gt_sem_seg`. |
| |
| Returns: |
| dict[str, Tensor]: A dictionary of loss components. |
| """ |
| x = self.extract_feat(batch_inputs) |
|
|
| |
| rpn_data_samples = copy.deepcopy(batch_data_samples) |
| for data_sample in rpn_data_samples: |
| data_sample.gt_instances.labels = \ |
| torch.zeros_like(data_sample.gt_instances.labels) |
|
|
| losses = self.bbox_head.loss(x, rpn_data_samples) |
| return losses |
|
|
