[Improvements] Refactor unittest folder structre

mmedit/datasets/pipelines/matting_aug.py

@@ -380,8 +380,9 @@ def __call__(self, results):
         trimap = results['trimap']
 
         # generete segmentation mask
-        kernel = cv2.getStructuringElement(
-            cv2.MORPH_ELLIPSE, (self.kernel_size, self.kernel_size))
+        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE,
+                                           (self.kernel_size,
+                                            self.kernel_size))
         seg = (alpha > 0.5).astype(np.float32)
         seg = cv2.erode(
             seg, kernel, iterations=np.random.randint(*self.erode_iter_range))
@@ -539,10 +540,12 @@ def __call__(self, results):
                                  cv2.BORDER_REPLICATE)
 
         # erode/dilate segmentation mask
-        erode_kernel = cv2.getStructuringElement(
-            cv2.MORPH_ELLIPSE, (self.erode_ksize, self.erode_ksize))
-        dilate_kernel = cv2.getStructuringElement(
-            cv2.MORPH_ELLIPSE, (self.dilate_ksize, self.dilate_ksize))
+        erode_kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE,
+                                                 (self.erode_ksize,
+                                                  self.erode_ksize))
+        dilate_kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE,
+                                                  (self.dilate_ksize,
+                                                   self.dilate_ksize))
         seg = cv2.erode(
             seg,
             erode_kernel,
@@ -616,8 +619,8 @@ def __call__(self, results):
                                                  cv2.DIST_L2, 0)**2
                 dt_mask = dt_mask[..., None]
                 L = 320
-                trimap_trans[..., 3 * k:3 * k + 3] = np.exp(
-                    dt_mask / (2 * ((factor * L)**2)))
+                trimap_trans[..., 3 * k:3 * k +
+                             3] = np.exp(dt_mask / (2 * ((factor * L)**2)))
 
         results['transformed_trimap'] = trimap_trans
         results['two_channel_trimap'] = trimap2

mmedit/models/backbones/__init__.py

@@ -1,4 +1,3 @@
-# yapf: disable
 from .encoder_decoders import (VGG16, ContextualAttentionNeck, DeepFillDecoder,
                                DeepFillEncoder, DeepFillEncoderDecoder,
                                DepthwiseIndexBlock, FBADecoder,
@@ -10,7 +9,6 @@
                                ResGCADecoder, ResGCAEncoder, ResNetDec,
                                ResNetEnc, ResShortcutDec, ResShortcutEnc,
                                SimpleEncoderDecoder)
-# yapf: enable
 from .generation_backbones import ResnetGenerator, UnetGenerator
 from .sr_backbones import (EDSR, RDN, SRCNN, BasicVSRNet, DICNet, EDVRNet,
                            GLEANStyleGANv2, IconVSR, MSRResNet, RRDBNet,

setup.cfg

@@ -11,6 +11,8 @@ addopts=tests/
 based_on_style = pep8
 blank_line_before_nested_class_or_def = true
 split_before_expression_after_opening_paren = true
+split_penalty_import_names=0
+SPLIT_PENALTY_AFTER_OPENING_BRACKET=888
 
 [isort]
 line_length = 79

tests/test_data/test_datasets/test_generation_datasets.py

@@ -0,0 +1,256 @@
+from pathlib import Path
+
+import pytest
+from mmcv.utils.testing import assert_dict_has_keys
+
+from mmedit.datasets import (BaseGenerationDataset, GenerationPairedDataset,
+                             GenerationUnpairedDataset)
+
+
+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 set(file_paths).issubset(set(result))
+        result = toy_dataset.scan_folder(str(self.data_prefix))
+        assert set(file_paths).issubset(set(result))
+
+        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 assert_dict_has_keys(result, target_keys)
+        assert assert_dict_has_keys(result['meta'].data, 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 assert_dict_has_keys(result, target_keys)
+        assert assert_dict_has_keys(result['meta'].data, 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 assert_dict_has_keys(result, target_keys)
+        assert assert_dict_has_keys(result['meta'].data, 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 assert_dict_has_keys(result, target_keys)
+        assert assert_dict_has_keys(result['meta'].data, 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 assert_dict_has_keys(result, target_keys)
+        assert assert_dict_has_keys(result['meta'].data, 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 assert_dict_has_keys(result, target_keys)
+        assert assert_dict_has_keys(result['meta'].data, 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 assert_dict_has_keys(result, target_keys)
+        assert assert_dict_has_keys(result['meta'].data, 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 assert_dict_has_keys(result, target_keys)
+        assert assert_dict_has_keys(result['meta'].data, 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