| | |
| | 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 |
| |
|
