| |
| from detectron2.layers import batched_nms |
| from detectron2.modeling import ROI_HEADS_REGISTRY, StandardROIHeads |
| from detectron2.modeling.roi_heads.roi_heads import Res5ROIHeads |
| from detectron2.structures import Instances |
|
|
|
|
| def merge_branch_instances(instances, num_branch, nms_thresh, topk_per_image): |
| """ |
| Merge detection results from different branches of TridentNet. |
| Return detection results by applying non-maximum suppression (NMS) on bounding boxes |
| and keep the unsuppressed boxes and other instances (e.g mask) if any. |
| |
| Args: |
| instances (list[Instances]): A list of N * num_branch instances that store detection |
| results. Contain N images and each image has num_branch instances. |
| num_branch (int): Number of branches used for merging detection results for each image. |
| nms_thresh (float): The threshold to use for box non-maximum suppression. Value in [0, 1]. |
| topk_per_image (int): The number of top scoring detections to return. Set < 0 to return |
| all detections. |
| |
| Returns: |
| results: (list[Instances]): A list of N instances, one for each image in the batch, |
| that stores the topk most confidence detections after merging results from multiple |
| branches. |
| """ |
| if num_branch == 1: |
| return instances |
|
|
| batch_size = len(instances) // num_branch |
| results = [] |
| for i in range(batch_size): |
| instance = Instances.cat([instances[i + batch_size * j] for j in range(num_branch)]) |
|
|
| |
| keep = batched_nms( |
| instance.pred_boxes.tensor, instance.scores, instance.pred_classes, nms_thresh |
| ) |
| keep = keep[:topk_per_image] |
| result = instance[keep] |
|
|
| results.append(result) |
|
|
| return results |
|
|
|
|
| @ROI_HEADS_REGISTRY.register() |
| class TridentRes5ROIHeads(Res5ROIHeads): |
| """ |
| The TridentNet ROIHeads in a typical "C4" R-CNN model. |
| See :class:`Res5ROIHeads`. |
| """ |
|
|
| def __init__(self, cfg, input_shape): |
| super().__init__(cfg, input_shape) |
|
|
| self.num_branch = cfg.MODEL.TRIDENT.NUM_BRANCH |
| self.trident_fast = cfg.MODEL.TRIDENT.TEST_BRANCH_IDX != -1 |
|
|
| def forward(self, images, features, proposals, targets=None): |
| """ |
| See :class:`Res5ROIHeads.forward`. |
| """ |
| num_branch = self.num_branch if self.training or not self.trident_fast else 1 |
| all_targets = targets * num_branch if targets is not None else None |
| pred_instances, losses = super().forward(images, features, proposals, all_targets) |
| del images, all_targets, targets |
|
|
| if self.training: |
| return pred_instances, losses |
| else: |
| pred_instances = merge_branch_instances( |
| pred_instances, |
| num_branch, |
| self.box_predictor.test_nms_thresh, |
| self.box_predictor.test_topk_per_image, |
| ) |
|
|
| return pred_instances, {} |
|
|
|
|
| @ROI_HEADS_REGISTRY.register() |
| class TridentStandardROIHeads(StandardROIHeads): |
| """ |
| The `StandardROIHeads` for TridentNet. |
| See :class:`StandardROIHeads`. |
| """ |
|
|
| def __init__(self, cfg, input_shape): |
| super(TridentStandardROIHeads, self).__init__(cfg, input_shape) |
|
|
| self.num_branch = cfg.MODEL.TRIDENT.NUM_BRANCH |
| self.trident_fast = cfg.MODEL.TRIDENT.TEST_BRANCH_IDX != -1 |
|
|
| def forward(self, images, features, proposals, targets=None): |
| """ |
| See :class:`Res5ROIHeads.forward`. |
| """ |
| |
| num_branch = self.num_branch if self.training or not self.trident_fast else 1 |
| |
| all_targets = targets * num_branch if targets is not None else None |
| pred_instances, losses = super().forward(images, features, proposals, all_targets) |
| del images, all_targets, targets |
|
|
| if self.training: |
| return pred_instances, losses |
| else: |
| pred_instances = merge_branch_instances( |
| pred_instances, |
| num_branch, |
| self.box_predictor.test_nms_thresh, |
| self.box_predictor.test_topk_per_image, |
| ) |
|
|
| return pred_instances, {} |
|
|
