Spaces:
Running
on
Zero
Running
on
Zero
| # Copyright (c) Facebook, Inc. and its affiliates. | |
| import itertools | |
| import logging | |
| from typing import Dict, List | |
| import torch | |
| from detectron2.config import configurable | |
| from detectron2.layers import ShapeSpec, batched_nms_rotated, cat | |
| from detectron2.structures import Instances, RotatedBoxes, pairwise_iou_rotated | |
| from detectron2.utils.memory import retry_if_cuda_oom | |
| from ..box_regression import Box2BoxTransformRotated | |
| from .build import PROPOSAL_GENERATOR_REGISTRY | |
| from .proposal_utils import _is_tracing | |
| from .rpn import RPN | |
| logger = logging.getLogger(__name__) | |
| def find_top_rrpn_proposals( | |
| proposals, | |
| pred_objectness_logits, | |
| image_sizes, | |
| nms_thresh, | |
| pre_nms_topk, | |
| post_nms_topk, | |
| min_box_size, | |
| training, | |
| ): | |
| """ | |
| For each feature map, select the `pre_nms_topk` highest scoring proposals, | |
| apply NMS, clip proposals, and remove small boxes. Return the `post_nms_topk` | |
| highest scoring proposals among all the feature maps if `training` is True, | |
| otherwise, returns the highest `post_nms_topk` scoring proposals for each | |
| feature map. | |
| Args: | |
| proposals (list[Tensor]): A list of L tensors. Tensor i has shape (N, Hi*Wi*A, 5). | |
| All proposal predictions on the feature maps. | |
| pred_objectness_logits (list[Tensor]): A list of L tensors. Tensor i has shape (N, Hi*Wi*A). | |
| image_sizes (list[tuple]): sizes (h, w) for each image | |
| nms_thresh (float): IoU threshold to use for NMS | |
| pre_nms_topk (int): number of top k scoring proposals to keep before applying NMS. | |
| When RRPN is run on multiple feature maps (as in FPN) this number is per | |
| feature map. | |
| post_nms_topk (int): number of top k scoring proposals to keep after applying NMS. | |
| When RRPN is run on multiple feature maps (as in FPN) this number is total, | |
| over all feature maps. | |
| min_box_size(float): minimum proposal box side length in pixels (absolute units wrt | |
| input images). | |
| training (bool): True if proposals are to be used in training, otherwise False. | |
| This arg exists only to support a legacy bug; look for the "NB: Legacy bug ..." | |
| comment. | |
| Returns: | |
| proposals (list[Instances]): list of N Instances. The i-th Instances | |
| stores post_nms_topk object proposals for image i. | |
| """ | |
| num_images = len(image_sizes) | |
| device = proposals[0].device | |
| # 1. Select top-k anchor for every level and every image | |
| topk_scores = [] # #lvl Tensor, each of shape N x topk | |
| topk_proposals = [] | |
| level_ids = [] # #lvl Tensor, each of shape (topk,) | |
| batch_idx = torch.arange(num_images, device=device) | |
| for level_id, proposals_i, logits_i in zip( | |
| itertools.count(), proposals, pred_objectness_logits | |
| ): | |
| Hi_Wi_A = logits_i.shape[1] | |
| if isinstance(Hi_Wi_A, torch.Tensor): # it's a tensor in tracing | |
| num_proposals_i = torch.clamp(Hi_Wi_A, max=pre_nms_topk) | |
| else: | |
| num_proposals_i = min(Hi_Wi_A, pre_nms_topk) | |
| topk_scores_i, topk_idx = logits_i.topk(num_proposals_i, dim=1) | |
| # each is N x topk | |
| topk_proposals_i = proposals_i[batch_idx[:, None], topk_idx] # N x topk x 5 | |
| topk_proposals.append(topk_proposals_i) | |
| topk_scores.append(topk_scores_i) | |
| level_ids.append(torch.full((num_proposals_i,), level_id, dtype=torch.int64, device=device)) | |
| # 2. Concat all levels together | |
| topk_scores = cat(topk_scores, dim=1) | |
| topk_proposals = cat(topk_proposals, dim=1) | |
| level_ids = cat(level_ids, dim=0) | |
| # 3. For each image, run a per-level NMS, and choose topk results. | |
| results = [] | |
| for n, image_size in enumerate(image_sizes): | |
| boxes = RotatedBoxes(topk_proposals[n]) | |
| scores_per_img = topk_scores[n] | |
| lvl = level_ids | |
| valid_mask = torch.isfinite(boxes.tensor).all(dim=1) & torch.isfinite(scores_per_img) | |
| if not valid_mask.all(): | |
| if training: | |
| raise FloatingPointError( | |
| "Predicted boxes or scores contain Inf/NaN. Training has diverged." | |
| ) | |
| boxes = boxes[valid_mask] | |
| scores_per_img = scores_per_img[valid_mask] | |
| lvl = lvl[valid_mask] | |
| boxes.clip(image_size) | |
| # filter empty boxes | |
| keep = boxes.nonempty(threshold=min_box_size) | |
| if _is_tracing() or keep.sum().item() != len(boxes): | |
| boxes, scores_per_img, lvl = (boxes[keep], scores_per_img[keep], lvl[keep]) | |
| keep = batched_nms_rotated(boxes.tensor, scores_per_img, lvl, nms_thresh) | |
| # In Detectron1, there was different behavior during training vs. testing. | |
| # (https://github.com/facebookresearch/Detectron/issues/459) | |
| # During training, topk is over the proposals from *all* images in the training batch. | |
| # During testing, it is over the proposals for each image separately. | |
| # As a result, the training behavior becomes batch-dependent, | |
| # and the configuration "POST_NMS_TOPK_TRAIN" end up relying on the batch size. | |
| # This bug is addressed in Detectron2 to make the behavior independent of batch size. | |
| keep = keep[:post_nms_topk] | |
| res = Instances(image_size) | |
| res.proposal_boxes = boxes[keep] | |
| res.objectness_logits = scores_per_img[keep] | |
| results.append(res) | |
| return results | |
| class RRPN(RPN): | |
| """ | |
| Rotated Region Proposal Network described in :paper:`RRPN`. | |
| """ | |
| def __init__(self, *args, **kwargs): | |
| super().__init__(*args, **kwargs) | |
| if self.anchor_boundary_thresh >= 0: | |
| raise NotImplementedError( | |
| "anchor_boundary_thresh is a legacy option not implemented for RRPN." | |
| ) | |
| def from_config(cls, cfg, input_shape: Dict[str, ShapeSpec]): | |
| ret = super().from_config(cfg, input_shape) | |
| ret["box2box_transform"] = Box2BoxTransformRotated(weights=cfg.MODEL.RPN.BBOX_REG_WEIGHTS) | |
| return ret | |
| def label_and_sample_anchors(self, anchors: List[RotatedBoxes], gt_instances: List[Instances]): | |
| """ | |
| Args: | |
| anchors (list[RotatedBoxes]): anchors for each feature map. | |
| gt_instances: the ground-truth instances for each image. | |
| Returns: | |
| list[Tensor]: | |
| List of #img tensors. i-th element is a vector of labels whose length is | |
| the total number of anchors across feature maps. Label values are in {-1, 0, 1}, | |
| with meanings: -1 = ignore; 0 = negative class; 1 = positive class. | |
| list[Tensor]: | |
| i-th element is a Nx5 tensor, where N is the total number of anchors across | |
| feature maps. The values are the matched gt boxes for each anchor. | |
| Values are undefined for those anchors not labeled as 1. | |
| """ | |
| anchors = RotatedBoxes.cat(anchors) | |
| gt_boxes = [x.gt_boxes for x in gt_instances] | |
| del gt_instances | |
| gt_labels = [] | |
| matched_gt_boxes = [] | |
| for gt_boxes_i in gt_boxes: | |
| """ | |
| gt_boxes_i: ground-truth boxes for i-th image | |
| """ | |
| match_quality_matrix = retry_if_cuda_oom(pairwise_iou_rotated)(gt_boxes_i, anchors) | |
| matched_idxs, gt_labels_i = retry_if_cuda_oom(self.anchor_matcher)(match_quality_matrix) | |
| # Matching is memory-expensive and may result in CPU tensors. But the result is small | |
| gt_labels_i = gt_labels_i.to(device=gt_boxes_i.device) | |
| # A vector of labels (-1, 0, 1) for each anchor | |
| gt_labels_i = self._subsample_labels(gt_labels_i) | |
| if len(gt_boxes_i) == 0: | |
| # These values won't be used anyway since the anchor is labeled as background | |
| matched_gt_boxes_i = torch.zeros_like(anchors.tensor) | |
| else: | |
| # TODO wasted indexing computation for ignored boxes | |
| matched_gt_boxes_i = gt_boxes_i[matched_idxs].tensor | |
| gt_labels.append(gt_labels_i) # N,AHW | |
| matched_gt_boxes.append(matched_gt_boxes_i) | |
| return gt_labels, matched_gt_boxes | |
| def predict_proposals(self, anchors, pred_objectness_logits, pred_anchor_deltas, image_sizes): | |
| pred_proposals = self._decode_proposals(anchors, pred_anchor_deltas) | |
| return find_top_rrpn_proposals( | |
| pred_proposals, | |
| pred_objectness_logits, | |
| image_sizes, | |
| self.nms_thresh, | |
| self.pre_nms_topk[self.training], | |
| self.post_nms_topk[self.training], | |
| self.min_box_size, | |
| self.training, | |
| ) | |