File size: 1,572 Bytes
6c9ac8f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 |
# Copyright (c) OpenMMLab. All rights reserved.
from unittest import TestCase
import torch
from mmengine import ConfigDict
from mmdet.models import DetTTAModel
from mmdet.registry import MODELS
from mmdet.structures import DetDataSample
from mmdet.testing import get_detector_cfg
from mmdet.utils import register_all_modules
class TestDetTTAModel(TestCase):
def setUp(self):
register_all_modules()
def test_det_tta_model(self):
detector_cfg = get_detector_cfg(
'retinanet/retinanet_r18_fpn_1x_coco.py')
cfg = ConfigDict(
type='DetTTAModel',
module=detector_cfg,
tta_cfg=dict(
nms=dict(type='nms', iou_threshold=0.5), max_per_img=100))
model: DetTTAModel = MODELS.build(cfg)
imgs = []
data_samples = []
directions = ['horizontal', 'vertical']
for i in range(12):
flip_direction = directions[0] if i % 3 == 0 else directions[1]
imgs.append(torch.randn(1, 3, 100 + 10 * i, 100 + 10 * i))
data_samples.append([
DetDataSample(
metainfo=dict(
ori_shape=(100, 100),
img_shape=(100 + 10 * i, 100 + 10 * i),
scale_factor=((100 + 10 * i) / 100,
(100 + 10 * i) / 100),
flip=(i % 2 == 0),
flip_direction=flip_direction), )
])
model.test_step(dict(inputs=imgs, data_samples=data_samples))
|