mmpretrain/tests/test_data/test_datasets/test_common.py

596 lines
21 KiB
Python

# Copyright (c) OpenMMLab. All rights reserved.
# import os
import os.path as osp
from unittest import TestCase
from unittest.mock import MagicMock
import numpy as np
# import pickle
# import tempfile
from mmengine.registry import TRANSFORMS
from mmcls.datasets import DATASETS
from mmcls.utils import get_root_logger
# import torch
mmcls_logger = get_root_logger()
ASSETS_ROOT = osp.abspath(
osp.join(osp.dirname(__file__), '../../data/dataset'))
class TestBaseDataset(TestCase):
DATASET_TYPE = 'BaseDataset'
DEFAULT_ARGS = dict(data_root=ASSETS_ROOT, ann_file='ann.json')
def test_initialize(self):
dataset_class = DATASETS.get(self.DATASET_TYPE)
# Test loading metainfo from ann_file
cfg = {**self.DEFAULT_ARGS, 'metainfo': None, 'classes': None}
dataset = dataset_class(**cfg)
self.assertEqual(
dataset.CLASSES,
dataset_class.METAINFO.get('CLASSES', ('first', 'second')))
self.assertFalse(dataset.test_mode)
# Test overriding metainfo by `metainfo` argument
cfg = {**self.DEFAULT_ARGS, 'metainfo': {'CLASSES': ('bus', 'car')}}
dataset = dataset_class(**cfg)
self.assertEqual(dataset.CLASSES, ('bus', 'car'))
# Test overriding metainfo by `classes` argument
cfg = {**self.DEFAULT_ARGS, 'classes': ['bus', 'car']}
dataset = dataset_class(**cfg)
self.assertEqual(dataset.CLASSES, ('bus', 'car'))
classes_file = osp.join(ASSETS_ROOT, 'classes.txt')
cfg = {**self.DEFAULT_ARGS, 'classes': classes_file}
dataset = dataset_class(**cfg)
self.assertEqual(dataset.CLASSES, ('bus', 'car'))
self.assertEqual(dataset.class_to_idx, {'bus': 0, 'car': 1})
# Test invalid classes
cfg = {**self.DEFAULT_ARGS, 'classes': dict(classes=1)}
with self.assertRaisesRegex(ValueError, "type <class 'dict'>"):
dataset_class(**cfg)
def test_get_cat_ids(self):
dataset_class = DATASETS.get(self.DATASET_TYPE)
dataset = dataset_class(**self.DEFAULT_ARGS)
cat_ids = dataset.get_cat_ids(0)
self.assertIsInstance(cat_ids, list)
self.assertEqual(len(cat_ids), 1)
self.assertIsInstance(cat_ids[0], int)
def test_repr(self):
dataset_class = DATASETS.get(self.DATASET_TYPE)
cfg = {**self.DEFAULT_ARGS, 'lazy_init': True}
dataset = dataset_class(**cfg)
head = 'Dataset ' + dataset.__class__.__name__
self.assertIn(head, repr(dataset))
if dataset.CLASSES is not None:
num_classes = len(dataset.CLASSES)
self.assertIn(f'Number of categories: \t{num_classes}',
repr(dataset))
else:
self.assertIn('The `CLASSES` meta info is not set.', repr(dataset))
self.assertIn("Haven't been initialized", repr(dataset))
dataset.full_init()
self.assertIn(f'Number of samples: \t{len(dataset)}', repr(dataset))
self.assertIn(f'Annotation file: \t{dataset.ann_file}', repr(dataset))
self.assertIn(f'Prefix of images: \t{dataset.img_prefix}',
repr(dataset))
TRANSFORMS.register_module(name='test_mock', module=MagicMock)
cfg = {**self.DEFAULT_ARGS, 'pipeline': [dict(type='test_mock')]}
dataset = dataset_class(**cfg)
self.assertIn('With transforms', repr(dataset))
del TRANSFORMS.module_dict['test_mock']
"""Temporarily disabled.
class TestMultiLabelDataset(TestBaseDataset):
DATASET_TYPE = 'MultiLabelDataset'
def test_get_cat_ids(self):
dataset_class = DATASETS.get(self.DATASET_TYPE)
fake_ann = [
dict(
img_prefix='',
img_info=dict(),
gt_label=np.array([0, 1, 1, 0], dtype=np.uint8))
]
with patch.object(dataset_class, 'load_annotations') as mock_load:
mock_load.return_value = fake_ann
dataset = dataset_class(**self.DEFAULT_ARGS)
cat_ids = dataset.get_cat_ids(0)
self.assertIsInstance(cat_ids, list)
self.assertEqual(len(cat_ids), 2)
self.assertIsInstance(cat_ids[0], int)
self.assertEqual(cat_ids, [1, 2])
def test_evaluate(self):
dataset_class = DATASETS.get(self.DATASET_TYPE)
fake_ann = [
dict(gt_label=np.array([1, 1, 0, -1], dtype=np.int8)),
dict(gt_label=np.array([1, 1, 0, -1], dtype=np.int8)),
dict(gt_label=np.array([0, -1, 1, -1], dtype=np.int8)),
dict(gt_label=np.array([0, 1, 0, -1], dtype=np.int8)),
dict(gt_label=np.array([0, 1, 0, -1], dtype=np.int8)),
]
with patch.object(dataset_class, 'load_annotations') as mock_load:
mock_load.return_value = fake_ann
dataset = dataset_class(**self.DEFAULT_ARGS)
fake_results = np.array([
[0.9, 0.8, 0.3, 0.2],
[0.1, 0.2, 0.2, 0.1],
[0.7, 0.5, 0.9, 0.3],
[0.8, 0.1, 0.1, 0.2],
[0.8, 0.1, 0.1, 0.2],
])
# the metric must be valid for the dataset
with self.assertRaisesRegex(ValueError,
"{'unknown'} is not supported"):
dataset.evaluate(fake_results, metric='unknown')
# only one metric
eval_results = dataset.evaluate(fake_results, metric='mAP')
self.assertEqual(eval_results.keys(), {'mAP'})
self.assertAlmostEqual(eval_results['mAP'], 67.5, places=4)
# multiple metrics
eval_results = dataset.evaluate(
fake_results, metric=['mAP', 'CR', 'OF1'])
self.assertEqual(eval_results.keys(), {'mAP', 'CR', 'OF1'})
self.assertAlmostEqual(eval_results['mAP'], 67.50, places=2)
self.assertAlmostEqual(eval_results['CR'], 43.75, places=2)
self.assertAlmostEqual(eval_results['OF1'], 42.86, places=2)
"""
class TestCustomDataset(TestBaseDataset):
DATASET_TYPE = 'CustomDataset'
DEFAULT_ARGS = dict(data_root=ASSETS_ROOT, ann_file='ann.txt')
def test_initialize(self):
dataset_class = DATASETS.get(self.DATASET_TYPE)
# Test overriding metainfo by `metainfo` argument
cfg = {**self.DEFAULT_ARGS, 'metainfo': {'CLASSES': ('bus', 'car')}}
dataset = dataset_class(**cfg)
self.assertEqual(dataset.CLASSES, ('bus', 'car'))
# Test overriding metainfo by `classes` argument
cfg = {**self.DEFAULT_ARGS, 'classes': ['bus', 'car']}
dataset = dataset_class(**cfg)
self.assertEqual(dataset.CLASSES, ('bus', 'car'))
classes_file = osp.join(ASSETS_ROOT, 'classes.txt')
cfg = {**self.DEFAULT_ARGS, 'classes': classes_file}
dataset = dataset_class(**cfg)
self.assertEqual(dataset.CLASSES, ('bus', 'car'))
self.assertEqual(dataset.class_to_idx, {'bus': 0, 'car': 1})
# Test invalid classes
cfg = {**self.DEFAULT_ARGS, 'classes': dict(classes=1)}
with self.assertRaisesRegex(ValueError, "type <class 'dict'>"):
dataset_class(**cfg)
def test_load_data_list(self):
dataset_class = DATASETS.get(self.DATASET_TYPE)
# test load without ann_file
cfg = {
**self.DEFAULT_ARGS,
'data_prefix': ASSETS_ROOT,
'ann_file': None,
}
dataset = dataset_class(**cfg)
self.assertEqual(len(dataset), 3)
self.assertEqual(dataset.CLASSES, ('a', 'b')) # auto infer classes
self.assertGreaterEqual(
dataset.get_data_info(0).items(), {
'img_path': osp.join(ASSETS_ROOT, 'a/1.JPG'),
'gt_label': 0
}.items())
self.assertGreaterEqual(
dataset.get_data_info(2).items(), {
'img_path': osp.join(ASSETS_ROOT, 'b/subb/3.jpg'),
'gt_label': 1
}.items())
# test ann_file assertion
cfg = {
**self.DEFAULT_ARGS,
'data_prefix': ASSETS_ROOT,
'ann_file': ['ann_file.txt'],
}
with self.assertRaisesRegex(TypeError, 'expected str'):
dataset_class(**cfg)
# test load with ann_file
cfg = {
**self.DEFAULT_ARGS,
'data_root': ASSETS_ROOT,
'ann_file': 'ann.txt',
}
dataset = dataset_class(**cfg)
self.assertEqual(len(dataset), 3)
# custom dataset won't infer CLASSES from ann_file
self.assertIsNone(dataset.CLASSES, None)
self.assertGreaterEqual(
dataset.get_data_info(0).items(), {
'img_path': osp.join(ASSETS_ROOT, 'a/1.JPG'),
'gt_label': 0,
}.items())
self.assertGreaterEqual(
dataset.get_data_info(2).items(), {
'img_path': osp.join(ASSETS_ROOT, 'b/subb/3.jpg'),
'gt_label': 1
}.items())
np.testing.assert_equal(dataset.get_gt_labels(), np.array([0, 1, 1]))
# test load with absolute ann_file
cfg = {
**self.DEFAULT_ARGS,
'data_root': None,
'data_prefix': None,
'ann_file': osp.join(ASSETS_ROOT, 'ann.txt'),
}
dataset = dataset_class(**cfg)
self.assertEqual(len(dataset), 3)
# custom dataset won't infer CLASSES from ann_file
self.assertIsNone(dataset.CLASSES, None)
self.assertGreaterEqual(
dataset.get_data_info(0).items(), {
'img_path': 'a/1.JPG',
'gt_label': 0,
}.items())
self.assertGreaterEqual(
dataset.get_data_info(2).items(), {
'img_path': 'b/subb/3.jpg',
'gt_label': 1
}.items())
# test extensions filter
cfg = {
**self.DEFAULT_ARGS, 'data_prefix': dict(img_path=ASSETS_ROOT),
'ann_file': None,
'extensions': ('.txt', )
}
with self.assertRaisesRegex(RuntimeError,
'Supported extensions are: .txt'):
dataset_class(**cfg)
cfg = {
**self.DEFAULT_ARGS, 'data_prefix': ASSETS_ROOT,
'ann_file': None,
'extensions': ('.jpeg', )
}
with self.assertLogs(mmcls_logger, 'WARN') as log:
dataset = dataset_class(**cfg)
self.assertIn('Supported extensions are: .jpeg', log.output[0])
self.assertEqual(len(dataset), 1)
self.assertGreaterEqual(
dataset.get_data_info(0).items(), {
'img_path': osp.join(ASSETS_ROOT, 'b/2.jpeg'),
'gt_label': 1
}.items())
# test classes check
cfg = {
**self.DEFAULT_ARGS,
'data_prefix': ASSETS_ROOT,
'classes': ('apple', 'banana'),
'ann_file': None,
}
dataset = dataset_class(**cfg)
self.assertEqual(dataset.CLASSES, ('apple', 'banana'))
cfg['classes'] = ['apple', 'banana', 'dog']
with self.assertRaisesRegex(AssertionError,
r"\(2\) doesn't match .* classes \(3\)"):
dataset_class(**cfg)
class TestImageNet(TestCustomDataset):
DATASET_TYPE = 'ImageNet'
DEFAULT_ARGS = dict(data_root=ASSETS_ROOT, ann_file='ann.txt')
def test_load_data_list(self):
dataset_class = DATASETS.get(self.DATASET_TYPE)
# test classes number
cfg = {
**self.DEFAULT_ARGS,
'data_prefix': ASSETS_ROOT,
'ann_file': None,
}
with self.assertRaisesRegex(
AssertionError, r"\(2\) doesn't match .* classes \(1000\)"):
dataset_class(**cfg)
# test override classes
cfg = {
**self.DEFAULT_ARGS,
'data_prefix': ASSETS_ROOT,
'classes': ['cat', 'dog'],
'ann_file': None,
}
dataset = dataset_class(**cfg)
self.assertEqual(len(dataset), 3)
self.assertEqual(dataset.CLASSES, ('cat', 'dog'))
class TestImageNet21k(TestCustomDataset):
DATASET_TYPE = 'ImageNet21k'
DEFAULT_ARGS = dict(
data_root=ASSETS_ROOT, classes=['cat', 'dog'], ann_file='ann.txt')
def test_load_data_list(self):
super().test_initialize()
dataset_class = DATASETS.get(self.DATASET_TYPE)
# The multi_label option is not implemented not.
cfg = {**self.DEFAULT_ARGS, 'multi_label': True}
with self.assertRaisesRegex(NotImplementedError, 'not supported'):
dataset_class(**cfg)
# Warn about ann_file
cfg = {**self.DEFAULT_ARGS, 'ann_file': None}
with self.assertLogs(mmcls_logger, 'WARN') as log:
dataset_class(**cfg)
self.assertIn('specify the `ann_file`', log.output[0])
# Warn about classes
cfg = {**self.DEFAULT_ARGS, 'classes': None}
with self.assertLogs(mmcls_logger, 'WARN') as log:
dataset_class(**cfg)
self.assertIn('specify the `classes`', log.output[0])
"""Temporarily disabled.
class TestMNIST(TestBaseDataset):
DATASET_TYPE = 'MNIST'
@classmethod
def setUpClass(cls) -> None:
super().setUpClass()
tmpdir = tempfile.TemporaryDirectory()
cls.tmpdir = tmpdir
data_prefix = tmpdir.name
cls.DEFAULT_ARGS = dict(data_prefix=data_prefix, pipeline=[])
dataset_class = DATASETS.get(cls.DATASET_TYPE)
def rm_suffix(s):
return s[:s.rfind('.')]
train_image_file = osp.join(
data_prefix,
rm_suffix(dataset_class.resources['train_image_file'][0]))
train_label_file = osp.join(
data_prefix,
rm_suffix(dataset_class.resources['train_label_file'][0]))
test_image_file = osp.join(
data_prefix,
rm_suffix(dataset_class.resources['test_image_file'][0]))
test_label_file = osp.join(
data_prefix,
rm_suffix(dataset_class.resources['test_label_file'][0]))
cls.fake_img = np.random.randint(0, 255, size=(28, 28), dtype=np.uint8)
cls.fake_label = np.random.randint(0, 10, size=(1, ), dtype=np.uint8)
for file in [train_image_file, test_image_file]:
magic = b'\x00\x00\x08\x03' # num_dims = 3, type = uint8
head = b'\x00\x00\x00\x01' + b'\x00\x00\x00\x1c' * 2 # (1, 28, 28)
data = magic + head + cls.fake_img.flatten().tobytes()
with open(file, 'wb') as f:
f.write(data)
for file in [train_label_file, test_label_file]:
magic = b'\x00\x00\x08\x01' # num_dims = 3, type = uint8
head = b'\x00\x00\x00\x01' # (1, )
data = magic + head + cls.fake_label.tobytes()
with open(file, 'wb') as f:
f.write(data)
def test_load_annotations(self):
dataset_class = DATASETS.get(self.DATASET_TYPE)
with patch.object(dataset_class, 'download'):
# Test default behavior
dataset = dataset_class(**self.DEFAULT_ARGS)
self.assertEqual(len(dataset), 1)
data_info = dataset[0]
np.testing.assert_equal(data_info['img'], self.fake_img)
np.testing.assert_equal(data_info['gt_label'], self.fake_label)
@classmethod
def tearDownClass(cls):
cls.tmpdir.cleanup()
class TestCIFAR10(TestBaseDataset):
DATASET_TYPE = 'CIFAR10'
@classmethod
def setUpClass(cls) -> None:
super().setUpClass()
tmpdir = tempfile.TemporaryDirectory()
cls.tmpdir = tmpdir
data_prefix = tmpdir.name
cls.DEFAULT_ARGS = dict(data_prefix=data_prefix, pipeline=[])
dataset_class = DATASETS.get(cls.DATASET_TYPE)
base_folder = osp.join(data_prefix, dataset_class.base_folder)
os.mkdir(base_folder)
cls.fake_imgs = np.random.randint(
0, 255, size=(6, 3 * 32 * 32), dtype=np.uint8)
cls.fake_labels = np.random.randint(0, 10, size=(6, ))
cls.fake_classes = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
batch1 = dict(
data=cls.fake_imgs[:2], labels=cls.fake_labels[:2].tolist())
with open(osp.join(base_folder, 'data_batch_1'), 'wb') as f:
f.write(pickle.dumps(batch1))
batch2 = dict(
data=cls.fake_imgs[2:4], labels=cls.fake_labels[2:4].tolist())
with open(osp.join(base_folder, 'data_batch_2'), 'wb') as f:
f.write(pickle.dumps(batch2))
test_batch = dict(
data=cls.fake_imgs[4:], labels=cls.fake_labels[4:].tolist())
with open(osp.join(base_folder, 'test_batch'), 'wb') as f:
f.write(pickle.dumps(test_batch))
meta = {dataset_class.meta['key']: cls.fake_classes}
meta_filename = dataset_class.meta['filename']
with open(osp.join(base_folder, meta_filename), 'wb') as f:
f.write(pickle.dumps(meta))
dataset_class.train_list = [['data_batch_1', None],
['data_batch_2', None]]
dataset_class.test_list = [['test_batch', None]]
dataset_class.meta['md5'] = None
def test_load_annotations(self):
dataset_class = DATASETS.get(self.DATASET_TYPE)
# Test default behavior
dataset = dataset_class(**self.DEFAULT_ARGS)
self.assertEqual(len(dataset), 4)
self.assertEqual(dataset.CLASSES, self.fake_classes)
data_info = dataset[0]
fake_img = self.fake_imgs[0].reshape(3, 32, 32).transpose(1, 2, 0)
np.testing.assert_equal(data_info['img'], fake_img)
np.testing.assert_equal(data_info['gt_label'], self.fake_labels[0])
# Test with test_mode=True
cfg = {**self.DEFAULT_ARGS, 'test_mode': True}
dataset = dataset_class(**cfg)
self.assertEqual(len(dataset), 2)
data_info = dataset[0]
fake_img = self.fake_imgs[4].reshape(3, 32, 32).transpose(1, 2, 0)
np.testing.assert_equal(data_info['img'], fake_img)
np.testing.assert_equal(data_info['gt_label'], self.fake_labels[4])
@classmethod
def tearDownClass(cls):
cls.tmpdir.cleanup()
class TestCIFAR100(TestCIFAR10):
DATASET_TYPE = 'CIFAR100'
class TestVOC(TestMultiLabelDataset):
DATASET_TYPE = 'VOC'
DEFAULT_ARGS = dict(data_prefix='VOC2007', pipeline=[])
class TestCUB(TestBaseDataset):
DATASET_TYPE = 'CUB'
@classmethod
def setUpClass(cls) -> None:
super().setUpClass()
tmpdir = tempfile.TemporaryDirectory()
cls.tmpdir = tmpdir
cls.data_prefix = tmpdir.name
cls.ann_file = osp.join(cls.data_prefix, 'ann_file.txt')
cls.image_class_labels_file = osp.join(cls.data_prefix, 'classes.txt')
cls.train_test_split_file = osp.join(cls.data_prefix, 'split.txt')
cls.train_test_split_file2 = osp.join(cls.data_prefix, 'split2.txt')
cls.DEFAULT_ARGS = dict(
data_prefix=cls.data_prefix,
pipeline=[],
ann_file=cls.ann_file,
image_class_labels_file=cls.image_class_labels_file,
train_test_split_file=cls.train_test_split_file)
with open(cls.ann_file, 'w') as f:
f.write('\n'.join([
'1 1.txt',
'2 2.txt',
'3 3.txt',
]))
with open(cls.image_class_labels_file, 'w') as f:
f.write('\n'.join([
'1 2',
'2 3',
'3 1',
]))
with open(cls.train_test_split_file, 'w') as f:
f.write('\n'.join([
'1 0',
'2 1',
'3 1',
]))
with open(cls.train_test_split_file2, 'w') as f:
f.write('\n'.join([
'1 0',
'2 1',
]))
def test_load_annotations(self):
dataset_class = DATASETS.get(self.DATASET_TYPE)
# Test default behavior
dataset = dataset_class(**self.DEFAULT_ARGS)
self.assertEqual(len(dataset), 2)
self.assertEqual(dataset.CLASSES, dataset_class.CLASSES)
data_info = dataset[0]
np.testing.assert_equal(data_info['img_prefix'], self.data_prefix)
np.testing.assert_equal(data_info['img_info'], {'filename': '2.txt'})
np.testing.assert_equal(data_info['gt_label'], 3 - 1)
# Test with test_mode=True
cfg = {**self.DEFAULT_ARGS, 'test_mode': True}
dataset = dataset_class(**cfg)
self.assertEqual(len(dataset), 1)
data_info = dataset[0]
np.testing.assert_equal(data_info['img_prefix'], self.data_prefix)
np.testing.assert_equal(data_info['img_info'], {'filename': '1.txt'})
np.testing.assert_equal(data_info['gt_label'], 2 - 1)
# Test if the numbers of line are not match
cfg = {
**self.DEFAULT_ARGS, 'train_test_split_file':
self.train_test_split_file2
}
with self.assertRaisesRegex(AssertionError, 'should have same length'):
dataset_class(**cfg)
@classmethod
def tearDownClass(cls):
cls.tmpdir.cleanup()
"""