123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419 |
- # Copyright (c) OpenMMLab. All rights reserved.
- from typing import List, Optional, Tuple
- import torch
- from torch import Tensor
- from mmdet.registry import MODELS, TASK_UTILS
- from mmdet.structures import DetDataSample, SampleList
- 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 .base_roi_head import BaseRoIHead
- @MODELS.register_module()
- class StandardRoIHead(BaseRoIHead):
- """Simplest base roi head including one bbox head and one mask head."""
- def init_assigner_sampler(self) -> None:
- """Initialize assigner and sampler."""
- self.bbox_assigner = None
- self.bbox_sampler = None
- if self.train_cfg:
- self.bbox_assigner = TASK_UTILS.build(self.train_cfg.assigner)
- self.bbox_sampler = TASK_UTILS.build(
- self.train_cfg.sampler, default_args=dict(context=self))
- 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 init_mask_head(self, mask_roi_extractor: ConfigType,
- mask_head: ConfigType) -> None:
- """Initialize mask head and mask roi extractor.
- Args:
- mask_roi_extractor (dict or ConfigDict): Config of mask roi
- extractor.
- mask_head (dict or ConfigDict): Config of mask in mask head.
- """
- if mask_roi_extractor is not None:
- self.mask_roi_extractor = MODELS.build(mask_roi_extractor)
- self.share_roi_extractor = False
- else:
- self.share_roi_extractor = True
- self.mask_roi_extractor = self.bbox_roi_extractor
- self.mask_head = MODELS.build(mask_head)
- # TODO: Need to refactor later
- def forward(self,
- x: Tuple[Tensor],
- rpn_results_list: InstanceList,
- batch_data_samples: SampleList = None) -> tuple:
- """Network forward process. Usually includes backbone, neck and head
- forward without any post-processing.
- Args:
- x (List[Tensor]): Multi-level features that may have different
- resolutions.
- rpn_results_list (list[:obj:`InstanceData`]): List of region
- proposals.
- batch_data_samples (list[:obj:`DetDataSample`]): Each item contains
- the meta information of each image and corresponding
- annotations.
- Returns
- tuple: A tuple of features from ``bbox_head`` and ``mask_head``
- forward.
- """
- results = ()
- proposals = [rpn_results.bboxes for rpn_results in rpn_results_list]
- rois = bbox2roi(proposals)
- # bbox head
- if self.with_bbox:
- bbox_results = self._bbox_forward(x, rois)
- results = results + (bbox_results['cls_score'],
- bbox_results['bbox_pred'])
- # mask head
- if self.with_mask:
- mask_rois = rois[:100]
- mask_results = self._mask_forward(x, mask_rois)
- results = results + (mask_results['mask_preds'], )
- return 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
- # assign gts and sample proposals
- num_imgs = len(batch_data_samples)
- sampling_results = []
- for i in range(num_imgs):
- # 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],
- feats=[lvl_feat[i][None] for lvl_feat in x])
- 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'])
- # mask head forward and loss
- if self.with_mask:
- mask_results = self.mask_loss(x, sampling_results,
- bbox_results['bbox_feats'],
- batch_gt_instances)
- losses.update(mask_results['loss_mask'])
- return losses
- 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.
- - `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)
- if self.with_shared_head:
- bbox_feats = self.shared_head(bbox_feats)
- cls_score, bbox_pred = self.bbox_head(bbox_feats)
- bbox_results = dict(
- cls_score=cls_score, bbox_pred=bbox_pred, 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)
- 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 mask_loss(self, x: Tuple[Tensor],
- sampling_results: List[SamplingResult], bbox_feats: Tensor,
- batch_gt_instances: InstanceList) -> dict:
- """Perform forward propagation and loss calculation of the mask head on
- the features of the upstream network.
- Args:
- x (tuple[Tensor]): Tuple of multi-level img features.
- sampling_results (list["obj:`SamplingResult`]): Sampling results.
- bbox_feats (Tensor): Extract bbox RoI features.
- batch_gt_instances (list[:obj:`InstanceData`]): Batch of
- gt_instance. It usually includes ``bboxes``, ``labels``, and
- ``masks`` attributes.
- Returns:
- dict: Usually returns a dictionary with keys:
- - `mask_preds` (Tensor): Mask prediction.
- - `mask_feats` (Tensor): Extract mask RoI features.
- - `mask_targets` (Tensor): Mask target of each positive\
- proposals in the image.
- - `loss_mask` (dict): A dictionary of mask loss components.
- """
- if not self.share_roi_extractor:
- pos_rois = bbox2roi([res.pos_priors for res in sampling_results])
- mask_results = self._mask_forward(x, pos_rois)
- else:
- pos_inds = []
- device = bbox_feats.device
- for res in sampling_results:
- pos_inds.append(
- torch.ones(
- res.pos_priors.shape[0],
- device=device,
- dtype=torch.uint8))
- pos_inds.append(
- torch.zeros(
- res.neg_priors.shape[0],
- device=device,
- dtype=torch.uint8))
- pos_inds = torch.cat(pos_inds)
- mask_results = self._mask_forward(
- x, pos_inds=pos_inds, bbox_feats=bbox_feats)
- mask_loss_and_target = self.mask_head.loss_and_target(
- mask_preds=mask_results['mask_preds'],
- sampling_results=sampling_results,
- batch_gt_instances=batch_gt_instances,
- rcnn_train_cfg=self.train_cfg)
- mask_results.update(loss_mask=mask_loss_and_target['loss_mask'])
- return mask_results
- def _mask_forward(self,
- x: Tuple[Tensor],
- rois: Tensor = None,
- pos_inds: Optional[Tensor] = None,
- bbox_feats: Optional[Tensor] = None) -> dict:
- """Mask head forward function used in both training and testing.
- Args:
- x (tuple[Tensor]): Tuple of multi-level img features.
- rois (Tensor): RoIs with the shape (n, 5) where the first
- column indicates batch id of each RoI.
- pos_inds (Tensor, optional): Indices of positive samples.
- Defaults to None.
- bbox_feats (Tensor): Extract bbox RoI features. Defaults to None.
- Returns:
- dict[str, Tensor]: Usually returns a dictionary with keys:
- - `mask_preds` (Tensor): Mask prediction.
- - `mask_feats` (Tensor): Extract mask RoI features.
- """
- assert ((rois is not None) ^
- (pos_inds is not None and bbox_feats is not None))
- if rois is not None:
- mask_feats = self.mask_roi_extractor(
- x[:self.mask_roi_extractor.num_inputs], rois)
- if self.with_shared_head:
- mask_feats = self.shared_head(mask_feats)
- else:
- assert bbox_feats is not None
- mask_feats = bbox_feats[pos_inds]
- mask_preds = self.mask_head(mask_feats)
- mask_results = dict(mask_preds=mask_preds, mask_feats=mask_feats)
- return mask_results
- 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',
- box_type=self.bbox_head.predict_box_type,
- num_classes=self.bbox_head.num_classes,
- score_per_cls=rcnn_test_cfg is None)
- bbox_results = self._bbox_forward(x, rois)
- # split batch bbox prediction back to each image
- 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)
- # some detector with_reg is False, bbox_preds will be None
- if bbox_preds is not None:
- # TODO move this to a sabl_roi_head
- # the bbox prediction of some detectors like SABL is not Tensor
- if isinstance(bbox_preds, torch.Tensor):
- bbox_preds = bbox_preds.split(num_proposals_per_img, 0)
- else:
- bbox_preds = self.bbox_head.bbox_pred_split(
- bbox_preds, num_proposals_per_img)
- 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
- def predict_mask(self,
- x: Tuple[Tensor],
- batch_img_metas: List[dict],
- results_list: InstanceList,
- rescale: bool = False) -> InstanceList:
- """Perform forward propagation of the mask 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.
- results_list (list[:obj:`InstanceData`]): Detection results of
- each image.
- 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).
- - masks (Tensor): Has a shape (num_instances, H, W).
- """
- # don't need to consider aug_test.
- bboxes = [res.bboxes for res in results_list]
- mask_rois = bbox2roi(bboxes)
- if mask_rois.shape[0] == 0:
- results_list = empty_instances(
- batch_img_metas,
- mask_rois.device,
- task_type='mask',
- instance_results=results_list,
- mask_thr_binary=self.test_cfg.mask_thr_binary)
- return results_list
- mask_results = self._mask_forward(x, mask_rois)
- mask_preds = mask_results['mask_preds']
- # split batch mask prediction back to each image
- num_mask_rois_per_img = [len(res) for res in results_list]
- mask_preds = mask_preds.split(num_mask_rois_per_img, 0)
- # TODO: Handle the case where rescale is false
- results_list = self.mask_head.predict_by_feat(
- mask_preds=mask_preds,
- results_list=results_list,
- batch_img_metas=batch_img_metas,
- rcnn_test_cfg=self.test_cfg,
- rescale=rescale)
- return results_list
|