[Improvements] Refactor unittest folder structre

tests/test_data/test_datasets/test_generation_datasets.py

@@ -0,0 +1,261 @@
+from pathlib import Path
+
+import pytest
+
+# yapf: disable
+from mmedit.datasets import (BaseGenerationDataset, GenerationPairedDataset,
+                             GenerationUnpairedDataset)
+
+
+def check_keys_contain(result_keys, target_keys):
+    """Check if all elements in target_keys is in result_keys."""
+    return set(target_keys).issubset(set(result_keys))
+
+
+class TestGenerationDatasets:
+
+    @classmethod
+    def setup_class(cls):
+        cls.data_prefix = Path(__file__).parent.parent.parent / 'data'
+
+    def test_base_generation_dataset(self):
+
+        class ToyDataset(BaseGenerationDataset):
+            """Toy dataset for testing Generation Dataset."""
+
+            def load_annotations(self):
+                pass
+
+        toy_dataset = ToyDataset(pipeline=[])
+        file_paths = [
+            'paired/test/3.jpg', 'paired/train/1.jpg', 'paired/train/2.jpg'
+        ]
+        file_paths = [str(self.data_prefix / v) for v in file_paths]
+
+        # test scan_folder
+        result = toy_dataset.scan_folder(self.data_prefix)
+        assert check_keys_contain(result, file_paths)
+        result = toy_dataset.scan_folder(str(self.data_prefix))
+        assert check_keys_contain(result, file_paths)
+
+        with pytest.raises(TypeError):
+            toy_dataset.scan_folder(123)
+
+        # test evaluate
+        toy_dataset.data_infos = file_paths
+        with pytest.raises(TypeError):
+            _ = toy_dataset.evaluate(1)
+        test_results = [dict(saved_flag=True), dict(saved_flag=True)]
+        with pytest.raises(AssertionError):
+            _ = toy_dataset.evaluate(test_results)
+        test_results = [
+            dict(saved_flag=True),
+            dict(saved_flag=True),
+            dict(saved_flag=False)
+        ]
+        eval_results = toy_dataset.evaluate(test_results)
+        assert eval_results['val_saved_number'] == 2
+
+    def test_generation_paired_dataset(self):
+        # setup
+        img_norm_cfg = dict(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
+        pipeline = [
+            dict(
+                type='LoadPairedImageFromFile',
+                io_backend='disk',
+                key='pair',
+                flag='color'),
+            dict(
+                type='Resize',
+                keys=['img_a', 'img_b'],
+                scale=(286, 286),
+                interpolation='bicubic'),
+            dict(
+                type='FixedCrop',
+                keys=['img_a', 'img_b'],
+                crop_size=(256, 256)),
+            dict(type='Flip', keys=['img_a', 'img_b'], direction='horizontal'),
+            dict(type='RescaleToZeroOne', keys=['img_a', 'img_b']),
+            dict(
+                type='Normalize',
+                keys=['img_a', 'img_b'],
+                to_rgb=True,
+                **img_norm_cfg),
+            dict(type='ImageToTensor', keys=['img_a', 'img_b']),
+            dict(
+                type='Collect',
+                keys=['img_a', 'img_b'],
+                meta_keys=['img_a_path', 'img_b_path'])
+        ]
+        target_keys = ['img_a', 'img_b', 'meta']
+        target_meta_keys = ['img_a_path', 'img_b_path']
+        pair_folder = self.data_prefix / 'paired'
+
+        # input path is Path object
+        generation_paried_dataset = GenerationPairedDataset(
+            dataroot=pair_folder, pipeline=pipeline, test_mode=True)
+        data_infos = generation_paried_dataset.data_infos
+        assert data_infos == [
+            dict(pair_path=str(pair_folder / 'test' / '3.jpg'))
+        ]
+        result = generation_paried_dataset[0]
+        assert (len(generation_paried_dataset) == 1)
+        assert check_keys_contain(result.keys(), target_keys)
+        assert check_keys_contain(result['meta'].data.keys(), target_meta_keys)
+        assert (result['meta'].data['img_a_path'] == str(pair_folder / 'test' /
+                                                         '3.jpg'))
+        assert (result['meta'].data['img_b_path'] == str(pair_folder / 'test' /
+                                                         '3.jpg'))
+
+        # input path is str
+        generation_paried_dataset = GenerationPairedDataset(
+            dataroot=str(pair_folder), pipeline=pipeline, test_mode=True)
+        data_infos = generation_paried_dataset.data_infos
+        assert data_infos == [
+            dict(pair_path=str(pair_folder / 'test' / '3.jpg'))
+        ]
+        result = generation_paried_dataset[0]
+        assert (len(generation_paried_dataset) == 1)
+        assert check_keys_contain(result.keys(), target_keys)
+        assert check_keys_contain(result['meta'].data.keys(), target_meta_keys)
+        assert (result['meta'].data['img_a_path'] == str(pair_folder / 'test' /
+                                                         '3.jpg'))
+        assert (result['meta'].data['img_b_path'] == str(pair_folder / 'test' /
+                                                         '3.jpg'))
+
+        # test_mode = False
+        generation_paried_dataset = GenerationPairedDataset(
+            dataroot=str(pair_folder), pipeline=pipeline, test_mode=False)
+        data_infos = generation_paried_dataset.data_infos
+        assert data_infos == [
+            dict(pair_path=str(pair_folder / 'train' / '1.jpg')),
+            dict(pair_path=str(pair_folder / 'train' / '2.jpg'))
+        ]
+        assert (len(generation_paried_dataset) == 2)
+        result = generation_paried_dataset[0]
+        assert check_keys_contain(result.keys(), target_keys)
+        assert check_keys_contain(result['meta'].data.keys(), target_meta_keys)
+        assert (result['meta'].data['img_a_path'] == str(pair_folder /
+                                                         'train' / '1.jpg'))
+        assert (result['meta'].data['img_b_path'] == str(pair_folder /
+                                                         'train' / '1.jpg'))
+        result = generation_paried_dataset[1]
+        assert check_keys_contain(result.keys(), target_keys)
+        assert check_keys_contain(result['meta'].data.keys(), target_meta_keys)
+        assert (result['meta'].data['img_a_path'] == str(pair_folder /
+                                                         'train' / '2.jpg'))
+        assert (result['meta'].data['img_b_path'] == str(pair_folder /
+                                                         'train' / '2.jpg'))
+
+    def test_generation_unpaired_dataset(self):
+        # setup
+        img_norm_cfg = dict(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
+        pipeline = [
+            dict(
+                type='LoadImageFromFile',
+                io_backend='disk',
+                key='img_a',
+                flag='color'),
+            dict(
+                type='LoadImageFromFile',
+                io_backend='disk',
+                key='img_b',
+                flag='color'),
+            dict(
+                type='Resize',
+                keys=['img_a', 'img_b'],
+                scale=(286, 286),
+                interpolation='bicubic'),
+            dict(
+                type='Crop',
+                keys=['img_a', 'img_b'],
+                crop_size=(256, 256),
+                random_crop=True),
+            dict(type='Flip', keys=['img_a'], direction='horizontal'),
+            dict(type='Flip', keys=['img_b'], direction='horizontal'),
+            dict(type='RescaleToZeroOne', keys=['img_a', 'img_b']),
+            dict(
+                type='Normalize',
+                keys=['img_a', 'img_b'],
+                to_rgb=True,
+                **img_norm_cfg),
+            dict(type='ImageToTensor', keys=['img_a', 'img_b']),
+            dict(
+                type='Collect',
+                keys=['img_a', 'img_b'],
+                meta_keys=['img_a_path', 'img_b_path'])
+        ]
+        target_keys = ['img_a', 'img_b', 'meta']
+        target_meta_keys = ['img_a_path', 'img_b_path']
+        unpair_folder = self.data_prefix / 'unpaired'
+
+        # input path is Path object
+        generation_unpaired_dataset = GenerationUnpairedDataset(
+            dataroot=unpair_folder, pipeline=pipeline, test_mode=True)
+        data_infos_a = generation_unpaired_dataset.data_infos_a
+        data_infos_b = generation_unpaired_dataset.data_infos_b
+        assert data_infos_a == [
+            dict(path=str(unpair_folder / 'testA' / '5.jpg'))
+        ]
+        assert data_infos_b == [
+            dict(path=str(unpair_folder / 'testB' / '6.jpg'))
+        ]
+        result = generation_unpaired_dataset[0]
+        assert (len(generation_unpaired_dataset) == 1)
+        assert check_keys_contain(result.keys(), target_keys)
+        assert check_keys_contain(result['meta'].data.keys(), target_meta_keys)
+        assert (result['meta'].data['img_a_path'] == str(unpair_folder /
+                                                         'testA' / '5.jpg'))
+        assert (result['meta'].data['img_b_path'] == str(unpair_folder /
+                                                         'testB' / '6.jpg'))
+
+        # input path is str
+        generation_unpaired_dataset = GenerationUnpairedDataset(
+            dataroot=str(unpair_folder), pipeline=pipeline, test_mode=True)
+        data_infos_a = generation_unpaired_dataset.data_infos_a
+        data_infos_b = generation_unpaired_dataset.data_infos_b
+        assert data_infos_a == [
+            dict(path=str(unpair_folder / 'testA' / '5.jpg'))
+        ]
+        assert data_infos_b == [
+            dict(path=str(unpair_folder / 'testB' / '6.jpg'))
+        ]
+        result = generation_unpaired_dataset[0]
+        assert (len(generation_unpaired_dataset) == 1)
+        assert check_keys_contain(result.keys(), target_keys)
+        assert check_keys_contain(result['meta'].data.keys(), target_meta_keys)
+        assert (result['meta'].data['img_a_path'] == str(unpair_folder /
+                                                         'testA' / '5.jpg'))
+        assert (result['meta'].data['img_b_path'] == str(unpair_folder /
+                                                         'testB' / '6.jpg'))
+
+        # test_mode = False
+        generation_unpaired_dataset = GenerationUnpairedDataset(
+            dataroot=str(unpair_folder), pipeline=pipeline, test_mode=False)
+        data_infos_a = generation_unpaired_dataset.data_infos_a
+        data_infos_b = generation_unpaired_dataset.data_infos_b
+        assert data_infos_a == [
+            dict(path=str(unpair_folder / 'trainA' / '1.jpg')),
+            dict(path=str(unpair_folder / 'trainA' / '2.jpg'))
+        ]
+        assert data_infos_b == [
+            dict(path=str(unpair_folder / 'trainB' / '3.jpg')),
+            dict(path=str(unpair_folder / 'trainB' / '4.jpg'))
+        ]
+        assert (len(generation_unpaired_dataset) == 2)
+        img_b_paths = [
+            str(unpair_folder / 'trainB' / '3.jpg'),
+            str(unpair_folder / 'trainB' / '4.jpg')
+        ]
+        result = generation_unpaired_dataset[0]
+        assert check_keys_contain(result.keys(), target_keys)
+        assert check_keys_contain(result['meta'].data.keys(), target_meta_keys)
+        assert (result['meta'].data['img_a_path'] == str(unpair_folder /
+                                                         'trainA' / '1.jpg'))
+        assert result['meta'].data['img_b_path'] in img_b_paths
+        result = generation_unpaired_dataset[1]
+        assert check_keys_contain(result.keys(), target_keys)
+        assert check_keys_contain(result['meta'].data.keys(), target_meta_keys)
+        assert (result['meta'].data['img_a_path'] == str(unpair_folder /
+                                                         'trainA' / '2.jpg'))
+        assert result['meta'].data['img_b_path'] in img_b_paths

tests/test_data/test_datasets/test_matting_datasets.py

@@ -0,0 +1,54 @@
+import os.path as osp
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+from mmedit.datasets import AdobeComp1kDataset
+
+
+class TestMattingDatasets:
+
+    @classmethod
+    def setup_class(cls):
+        # create para for creating a dataset.
+        cls.data_prefix = Path(__file__).parent.parent.parent / 'data'
+        cls.ann_file = osp.join(cls.data_prefix, 'test_list.json')
+        cls.pipeline = [
+            dict(type='LoadImageFromFile', key='alpha', flag='grayscale')
+        ]
+
+    def test_comp1k_dataset(self):
+        comp1k_dataset = AdobeComp1kDataset(self.ann_file, self.pipeline,
+                                            self.data_prefix)
+        first_data = comp1k_dataset[0]
+
+        assert 'alpha' in first_data
+        assert isinstance(first_data['alpha'], np.ndarray)
+        assert first_data['alpha'].shape == (552, 800)
+
+    def test_comp1k_evaluate(self):
+        comp1k_dataset = AdobeComp1kDataset(self.ann_file, self.pipeline,
+                                            self.data_prefix)
+
+        with pytest.raises(TypeError):
+            comp1k_dataset.evaluate('Not a list object')
+
+        results = [{
+            'pred_alpha': None,
+            'eval_result': {
+                'SAD': 26,
+                'MSE': 0.006
+            }
+        }, {
+            'pred_alpha': None,
+            'eval_result': {
+                'SAD': 24,
+                'MSE': 0.004
+            }
+        }]
+
+        eval_result = comp1k_dataset.evaluate(results)
+        assert set(eval_result.keys()) == set(['SAD', 'MSE'])
+        assert eval_result['SAD'] == 25
+        assert eval_result['MSE'] == 0.005

tests/test_data/test_datasets/test_repeat_dataset.py

@@ -0,0 +1,24 @@
+from torch.utils.data import Dataset
+
+from mmedit.datasets import RepeatDataset
+
+
+def test_repeat_dataset():
+
+    class ToyDataset(Dataset):
+
+        def __init__(self):
+            super().__init__()
+            self.members = [1, 2, 3, 4, 5]
+
+        def __len__(self):
+            return len(self.members)
+
+        def __getitem__(self, idx):
+            return self.members[idx % 5]
+
+    toy_dataset = ToyDataset()
+    repeat_dataset = RepeatDataset(toy_dataset, 2)
+    assert len(repeat_dataset) == 10
+    assert repeat_dataset[2] == 3
+    assert repeat_dataset[8] == 4