mirror of https://github.com/open-mmlab/mmyolo.git
72 lines
2.4 KiB
Python
72 lines
2.4 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
import unittest
|
|
|
|
from mmyolo.datasets import YOLOv5CocoDataset
|
|
|
|
|
|
class TestYOLOv5CocoDataset(unittest.TestCase):
|
|
|
|
def test_batch_shapes_cfg(self):
|
|
batch_shapes_cfg = dict(
|
|
type='BatchShapePolicy',
|
|
batch_size=2,
|
|
img_size=640,
|
|
size_divisor=32,
|
|
extra_pad_ratio=0.5)
|
|
|
|
# test serialize_data=True
|
|
dataset = YOLOv5CocoDataset(
|
|
data_prefix=dict(img='imgs'),
|
|
ann_file='tests/data/coco_sample.json',
|
|
filter_cfg=dict(filter_empty_gt=False, min_size=0),
|
|
pipeline=[],
|
|
serialize_data=True,
|
|
batch_shapes_cfg=batch_shapes_cfg,
|
|
)
|
|
|
|
expected_img_ids = [3, 0, 2, 1]
|
|
expected_batch_shapes = [[512, 672], [512, 672], [672, 672],
|
|
[672, 672]]
|
|
for i, data in enumerate(dataset):
|
|
assert data['img_id'] == expected_img_ids[i]
|
|
assert data['batch_shape'].tolist() == expected_batch_shapes[i]
|
|
|
|
# test serialize_data=True
|
|
dataset = YOLOv5CocoDataset(
|
|
data_prefix=dict(img='imgs'),
|
|
ann_file='tests/data/coco_sample.json',
|
|
filter_cfg=dict(filter_empty_gt=False, min_size=0),
|
|
pipeline=[],
|
|
serialize_data=False,
|
|
batch_shapes_cfg=batch_shapes_cfg,
|
|
)
|
|
|
|
expected_img_ids = [3, 0, 2, 1]
|
|
expected_batch_shapes = [[512, 672], [512, 672], [672, 672],
|
|
[672, 672]]
|
|
for i, data in enumerate(dataset):
|
|
assert data['img_id'] == expected_img_ids[i]
|
|
assert data['batch_shape'].tolist() == expected_batch_shapes[i]
|
|
|
|
def test_prepare_data(self):
|
|
dataset = YOLOv5CocoDataset(
|
|
data_prefix=dict(img='imgs'),
|
|
ann_file='tests/data/coco_sample.json',
|
|
filter_cfg=dict(filter_empty_gt=False, min_size=0),
|
|
pipeline=[],
|
|
serialize_data=True,
|
|
batch_shapes_cfg=None,
|
|
)
|
|
for data in dataset:
|
|
assert 'dataset' in data
|
|
|
|
# test with test_mode = True
|
|
dataset = YOLOv5CocoDataset(
|
|
data_prefix=dict(img='imgs'),
|
|
ann_file='tests/data/coco_sample.json',
|
|
test_mode=True,
|
|
pipeline=[])
|
|
|
|
for data in dataset:
|
|
assert 'dataset' not in data
|