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