Spaces:
Running
on
T4
Running
on
T4
# 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 | |