# Copyright (c) OpenMMLab. All rights reserved. from typing import List, Tuple from torch import Tensor from mmdet.registry import MODELS from mmdet.structures import DetDataSample from mmdet.structures.bbox import bbox2roi from mmdet.utils import ConfigType, InstanceList from ..task_modules.samplers import SamplingResult from ..utils import empty_instances, unpack_gt_instances from .standard_roi_head import StandardRoIHead @MODELS.register_module() class MultiInstanceRoIHead(StandardRoIHead): """The roi head for Multi-instance prediction.""" def __init__(self, num_instance: int = 2, *args, **kwargs) -> None: self.num_instance = num_instance super().__init__(*args, **kwargs) def init_bbox_head(self, bbox_roi_extractor: ConfigType, bbox_head: ConfigType) -> None: """Initialize box head and box roi extractor. Args: bbox_roi_extractor (dict or ConfigDict): Config of box roi extractor. bbox_head (dict or ConfigDict): Config of box in box head. """ self.bbox_roi_extractor = MODELS.build(bbox_roi_extractor) self.bbox_head = MODELS.build(bbox_head) def _bbox_forward(self, x: Tuple[Tensor], rois: Tensor) -> dict: """Box head forward function used in both training and testing. Args: x (tuple[Tensor]): List of multi-level img features. rois (Tensor): RoIs with the shape (n, 5) where the first column indicates batch id of each RoI. Returns: dict[str, Tensor]: Usually returns a dictionary with keys: - `cls_score` (Tensor): Classification scores. - `bbox_pred` (Tensor): Box energies / deltas. - `cls_score_ref` (Tensor): The cls_score after refine model. - `bbox_pred_ref` (Tensor): The bbox_pred after refine model. - `bbox_feats` (Tensor): Extract bbox RoI features. """ # TODO: a more flexible way to decide which feature maps to use bbox_feats = self.bbox_roi_extractor( x[:self.bbox_roi_extractor.num_inputs], rois) bbox_results = self.bbox_head(bbox_feats) if self.bbox_head.with_refine: bbox_results = dict( cls_score=bbox_results[0], bbox_pred=bbox_results[1], cls_score_ref=bbox_results[2], bbox_pred_ref=bbox_results[3], bbox_feats=bbox_feats) else: bbox_results = dict( cls_score=bbox_results[0], bbox_pred=bbox_results[1], bbox_feats=bbox_feats) return bbox_results def bbox_loss(self, x: Tuple[Tensor], sampling_results: List[SamplingResult]) -> dict: """Perform forward propagation and loss calculation of the bbox head on the features of the upstream network. Args: x (tuple[Tensor]): List of multi-level img features. sampling_results (list["obj:`SamplingResult`]): Sampling results. Returns: dict[str, Tensor]: Usually returns a dictionary with keys: - `cls_score` (Tensor): Classification scores. - `bbox_pred` (Tensor): Box energies / deltas. - `bbox_feats` (Tensor): Extract bbox RoI features. - `loss_bbox` (dict): A dictionary of bbox loss components. """ rois = bbox2roi([res.priors for res in sampling_results]) bbox_results = self._bbox_forward(x, rois) # If there is a refining process, add refine loss. if 'cls_score_ref' in bbox_results: bbox_loss_and_target = self.bbox_head.loss_and_target( cls_score=bbox_results['cls_score'], bbox_pred=bbox_results['bbox_pred'], rois=rois, sampling_results=sampling_results, rcnn_train_cfg=self.train_cfg) bbox_results.update(loss_bbox=bbox_loss_and_target['loss_bbox']) bbox_loss_and_target_ref = self.bbox_head.loss_and_target( cls_score=bbox_results['cls_score_ref'], bbox_pred=bbox_results['bbox_pred_ref'], rois=rois, sampling_results=sampling_results, rcnn_train_cfg=self.train_cfg) bbox_results['loss_bbox']['loss_rcnn_emd_ref'] = \ bbox_loss_and_target_ref['loss_bbox']['loss_rcnn_emd'] else: bbox_loss_and_target = self.bbox_head.loss_and_target( cls_score=bbox_results['cls_score'], bbox_pred=bbox_results['bbox_pred'], rois=rois, sampling_results=sampling_results, rcnn_train_cfg=self.train_cfg) bbox_results.update(loss_bbox=bbox_loss_and_target['loss_bbox']) return bbox_results def loss(self, x: Tuple[Tensor], rpn_results_list: InstanceList, batch_data_samples: List[DetDataSample]) -> dict: """Perform forward propagation and loss calculation of the detection roi on the features of the upstream network. Args: x (tuple[Tensor]): List of multi-level img features. rpn_results_list (list[:obj:`InstanceData`]): List of region proposals. 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 """ assert len(rpn_results_list) == len(batch_data_samples) outputs = unpack_gt_instances(batch_data_samples) batch_gt_instances, batch_gt_instances_ignore, _ = outputs sampling_results = [] for i in range(len(batch_data_samples)): # rename rpn_results.bboxes to rpn_results.priors rpn_results = rpn_results_list[i] rpn_results.priors = rpn_results.pop('bboxes') assign_result = self.bbox_assigner.assign( rpn_results, batch_gt_instances[i], batch_gt_instances_ignore[i]) sampling_result = self.bbox_sampler.sample( assign_result, rpn_results, batch_gt_instances[i], batch_gt_instances_ignore=batch_gt_instances_ignore[i]) sampling_results.append(sampling_result) losses = dict() # bbox head loss if self.with_bbox: bbox_results = self.bbox_loss(x, sampling_results) losses.update(bbox_results['loss_bbox']) return losses def predict_bbox(self, x: Tuple[Tensor], batch_img_metas: List[dict], rpn_results_list: InstanceList, rcnn_test_cfg: ConfigType, rescale: bool = False) -> InstanceList: """Perform forward propagation of the bbox head and predict detection results on the features of the upstream network. Args: x (tuple[Tensor]): Feature maps of all scale level. batch_img_metas (list[dict]): List of image information. rpn_results_list (list[:obj:`InstanceData`]): List of region proposals. rcnn_test_cfg (obj:`ConfigDict`): `test_cfg` of R-CNN. rescale (bool): If True, return boxes in original image space. Defaults to False. Returns: list[:obj:`InstanceData`]: Detection results of each image after the post process. Each item 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). """ proposals = [res.bboxes for res in rpn_results_list] rois = bbox2roi(proposals) if rois.shape[0] == 0: return empty_instances( batch_img_metas, rois.device, task_type='bbox') bbox_results = self._bbox_forward(x, rois) # split batch bbox prediction back to each image if 'cls_score_ref' in bbox_results: cls_scores = bbox_results['cls_score_ref'] bbox_preds = bbox_results['bbox_pred_ref'] else: cls_scores = bbox_results['cls_score'] bbox_preds = bbox_results['bbox_pred'] num_proposals_per_img = tuple(len(p) for p in proposals) rois = rois.split(num_proposals_per_img, 0) cls_scores = cls_scores.split(num_proposals_per_img, 0) if bbox_preds is not None: bbox_preds = bbox_preds.split(num_proposals_per_img, 0) else: bbox_preds = (None, ) * len(proposals) result_list = self.bbox_head.predict_by_feat( rois=rois, cls_scores=cls_scores, bbox_preds=bbox_preds, batch_img_metas=batch_img_metas, rcnn_test_cfg=rcnn_test_cfg, rescale=rescale) return result_list