|
from mmdet.core import bbox2result |
|
from ..builder import DETECTORS |
|
from .single_stage import SingleStageDetector |
|
|
|
|
|
@DETECTORS.register_module() |
|
class DETR(SingleStageDetector): |
|
r"""Implementation of `DETR: End-to-End Object Detection with |
|
Transformers <https://arxiv.org/pdf/2005.12872>`_""" |
|
|
|
def __init__(self, |
|
backbone, |
|
bbox_head, |
|
train_cfg=None, |
|
test_cfg=None, |
|
pretrained=None, |
|
init_cfg=None): |
|
super(DETR, self).__init__(backbone, None, bbox_head, train_cfg, |
|
test_cfg, pretrained, init_cfg) |
|
|
|
def simple_test(self, img, img_metas, rescale=False): |
|
"""Test function without test time augmentation. |
|
|
|
Args: |
|
imgs (list[torch.Tensor]): List of multiple images |
|
img_metas (list[dict]): List of image information. |
|
rescale (bool, optional): Whether to rescale the results. |
|
Defaults to False. |
|
|
|
Returns: |
|
list[list[np.ndarray]]: BBox results of each image and classes. |
|
The outer list corresponds to each image. The inner list |
|
corresponds to each class. |
|
""" |
|
batch_size = len(img_metas) |
|
assert batch_size == 1, 'Currently only batch_size 1 for inference ' \ |
|
f'mode is supported. Found batch_size {batch_size}.' |
|
x = self.extract_feat(img) |
|
outs = self.bbox_head(x, img_metas) |
|
bbox_list = self.bbox_head.get_bboxes( |
|
*outs, img_metas, rescale=rescale) |
|
|
|
bbox_results = [ |
|
bbox2result(det_bboxes, det_labels, self.bbox_head.num_classes) |
|
for det_bboxes, det_labels in bbox_list |
|
] |
|
return bbox_results |
|
|
