Merge pull request open-mmlab#238 from ckkelvinchan/multiple-gt-dataset

Add MultipleGT datasets

mmedit/datasets/__init__.py

@@ -14,14 +14,19 @@
 from .sr_folder_gt_dataset import SRFolderGTDataset
 from .sr_lmdb_dataset import SRLmdbDataset
 from .sr_reds_dataset import SRREDSDataset
+from .sr_reds_multiple_gt_dataset import SRREDSMultipleGTDataset
+from .sr_test_multiple_gt_dataset import SRTestMultipleGTDataset
 from .sr_vid4_dataset import SRVid4Dataset
 from .sr_vimeo90k_dataset import SRVimeo90KDataset
+from .sr_vimeo90k_multiple_gt_dataset import SRVimeo90KMultipleGTDataset
 
 __all__ = [
     'DATASETS', 'PIPELINES', 'build_dataset', 'build_dataloader',
     'BaseDataset', 'BaseMattingDataset', 'ImgInpaintingDataset',
     'AdobeComp1kDataset', 'SRLmdbDataset', 'SRFolderDataset',
     'SRAnnotationDataset', 'BaseSRDataset', 'RepeatDataset', 'SRREDSDataset',
     'SRVimeo90KDataset', 'BaseGenerationDataset', 'GenerationPairedDataset',
-    'GenerationUnpairedDataset', 'SRVid4Dataset', 'SRFolderGTDataset'
+    'GenerationUnpairedDataset', 'SRVid4Dataset', 'SRFolderGTDataset',
+    'SRREDSMultipleGTDataset', 'SRVimeo90KMultipleGTDataset',
+    'SRTestMultipleGTDataset'
 ]

mmedit/datasets/sr_reds_multiple_gt_dataset.py

@@ -0,0 +1,73 @@
+from .base_sr_dataset import BaseSRDataset
+from .registry import DATASETS
+
+
+@DATASETS.register_module()
+class SRREDSMultipleGTDataset(BaseSRDataset):
+    """REDS dataset for video super resolution for recurrent networks.
+
+    The dataset loads several LQ (Low-Quality) frames and GT (Ground-Truth)
+    frames. Then it applies specified transforms and finally returns a dict
+    containing paired data and other information.
+
+    Args:
+        lq_folder (str | :obj:`Path`): Path to a lq folder.
+        gt_folder (str | :obj:`Path`): Path to a gt folder.
+        num_input_frames (int): Number of input frames.
+        pipeline (list[dict | callable]): A sequence of data transformations.
+        scale (int): Upsampling scale ratio.
+        val_partition (str): Validation partition mode. Choices ['official' or
+        'REDS4']. Default: 'official'.
+        test_mode (bool): Store `True` when building test dataset.
+            Default: `False`.
+    """
+
+    def __init__(self,
+                 lq_folder,
+                 gt_folder,
+                 num_input_frames,
+                 pipeline,
+                 scale,
+                 val_partition='official',
+                 test_mode=False):
+        super().__init__(pipeline, scale, test_mode)
+        self.lq_folder = str(lq_folder)
+        self.gt_folder = str(gt_folder)
+        self.num_input_frames = num_input_frames
+        self.val_partition = val_partition
+        self.data_infos = self.load_annotations()
+
+    def load_annotations(self):
+        """Load annoations for REDS dataset.
+
+        Returns:
+            dict: Returned dict for LQ and GT pairs.
+        """
+        # generate keys
+        keys = [f'{i:03d}' for i in range(0, 270)]
+
+        if self.val_partition == 'REDS4':
+            val_partition = ['000', '011', '015', '020']
+        elif self.val_partition == 'official':
+            val_partition = [f'{i:03d}' for i in range(240, 270)]
+        else:
+            raise ValueError(
+                f'Wrong validation partition {self.val_partition}.'
+                f'Supported ones are ["official", "REDS4"]')
+
+        if self.test_mode:
+            keys = [v for v in keys if v in val_partition]
+        else:
+            keys = [v for v in keys if v not in val_partition]
+
+        data_infos = []
+        for key in keys:
+            data_infos.append(
+                dict(
+                    lq_path=self.lq_folder,
+                    gt_path=self.gt_folder,
+                    key=key,
+                    sequence_length=100,  # REDS has 100 frames for each clip
+                    num_input_frames=self.num_input_frames))
+
+        return data_infos

mmedit/datasets/sr_test_multiple_gt_dataset.py

@@ -0,0 +1,53 @@
+import glob
+import os.path as osp
+
+from .base_sr_dataset import BaseSRDataset
+from .registry import DATASETS
+
+
+@DATASETS.register_module()
+class SRTestMultipleGTDataset(BaseSRDataset):
+    """Test dataset for video super resolution for recurrent networks.
+
+    It assumes all video sequences under the root directory is used for test.
+
+    The dataset loads several LQ (Low-Quality) frames and GT (Ground-Truth)
+    frames. Then it applies specified transforms and finally returns a dict
+    containing paired data and other information.
+
+    Args:
+        lq_folder (str | :obj:`Path`): Path to a lq folder.
+        gt_folder (str | :obj:`Path`): Path to a gt folder.
+        pipeline (list[dict | callable]): A sequence of data transformations.
+        scale (int): Upsampling scale ratio.
+        test_mode (bool): Store `True` when building test dataset.
+            Default: `True`.
+    """
+
+    def __init__(self, lq_folder, gt_folder, pipeline, scale, test_mode=True):
+        super().__init__(pipeline, scale, test_mode)
+
+        self.lq_folder = str(lq_folder)
+        self.gt_folder = str(gt_folder)
+        self.data_infos = self.load_annotations()
+
+    def load_annotations(self):
+        """Load annoations for the test dataset.
+
+        Returns:
+            dict: Returned dict for LQ and GT pairs.
+        """
+
+        sequences = sorted(glob.glob(osp.join(self.lq_folder, '*')))
+
+        data_infos = []
+        for sequence in sequences:
+            sequence_length = len(glob.glob(osp.join(sequence, '*.png')))
+            data_infos.append(
+                dict(
+                    lq_path=self.lq_folder,
+                    gt_path=self.gt_folder,
+                    key=sequence.replace(f'{self.lq_folder}/', ''),
+                    sequence_length=int(sequence_length)))
+
+        return data_infos

mmedit/datasets/sr_vimeo90k_multiple_gt_dataset.py

@@ -0,0 +1,68 @@
+import glob
+import os.path as osp
+
+from .base_sr_dataset import BaseSRDataset
+from .registry import DATASETS
+
+
+@DATASETS.register_module()
+class SRVimeo90KMultipleGTDataset(BaseSRDataset):
+    """Vimeo90K dataset for video super resolution for recurrent networks.
+
+    The dataset loads several LQ (Low-Quality) frames and GT (Ground-Truth)
+    frames. Then it applies specified transforms and finally returns a dict
+    containing paired data and other information.
+
+    It reads Vimeo90K keys from the txt file. Each line contains:
+
+        1. video frame folder
+        2. image shape
+
+    Examples:
+
+    ::
+
+        00001/0266 (256,448,3)
+        00001/0268 (256,448,3)
+
+    Args:
+        lq_folder (str | :obj:`Path`): Path to a lq folder.
+        gt_folder (str | :obj:`Path`): Path to a gt folder.
+        ann_file (str | :obj:`Path`): Path to the annotation file.
+        pipeline (list[dict | callable]): A sequence of data transformations.
+        scale (int): Upsampling scale ratio.
+        test_mode (bool): Store `True` when building test dataset.
+            Default: `False`.
+    """
+
+    def __init__(self, lq_folder, gt_folder, ann_file, pipeline, scale,
+                 test_mode):
+        super().__init__(pipeline, scale, test_mode)
+
+        self.lq_folder = str(lq_folder)
+        self.gt_folder = str(gt_folder)
+        self.ann_file = str(ann_file)
+
+        self.data_infos = self.load_annotations()
+
+    def load_annotations(self):
+        """Load annoations for Vimeo-90K dataset.
+
+        Returns:
+            dict: Returned dict for LQ and GT pairs.
+        """
+        # get keys
+        with open(self.ann_file, 'r') as fin:
+            keys = [line.strip().split(' ')[0] for line in fin]
+
+        data_infos = []
+        for key in keys:
+            lq_paths = sorted(
+                glob.glob(osp.join(self.lq_folder, key, '*.png')))
+            gt_paths = sorted(
+                glob.glob(osp.join(self.gt_folder, key, '*.png')))
+
+            data_infos.append(
+                dict(lq_path=lq_paths, gt_path=gt_paths, key=key))
+
+        return data_infos

tests/test_datasets.py

@@ -6,12 +6,17 @@
 import pytest
 from torch.utils.data import Dataset
 
+# yapf: disable
 from mmedit.datasets import (AdobeComp1kDataset, BaseGenerationDataset,
                              BaseSRDataset, GenerationPairedDataset,
                              GenerationUnpairedDataset, RepeatDataset,
                              SRAnnotationDataset, SRFolderDataset,
                              SRFolderGTDataset, SRLmdbDataset, SRREDSDataset,
-                             SRVid4Dataset, SRVimeo90KDataset)
+                             SRREDSMultipleGTDataset, SRTestMultipleGTDataset,
+                             SRVid4Dataset, SRVimeo90KDataset,
+                             SRVimeo90KMultipleGTDataset)
+
+# yapf: enable
 
 
 def mock_open(*args, **kwargs):
@@ -829,3 +834,141 @@ def test_vid4_dataset():
                 pipeline=[],
                 scale=4,
                 test_mode=False)
+
+
+def test_sr_reds_multiple_gt_dataset():
+    root_path = Path(__file__).parent / 'data'
+
+    # official val partition
+    reds_dataset = SRREDSMultipleGTDataset(
+        lq_folder=root_path,
+        gt_folder=root_path,
+        num_input_frames=15,
+        pipeline=[],
+        scale=4,
+        val_partition='official',
+        test_mode=False)
+
+    assert len(reds_dataset.data_infos) == 240  # 240 training clips
+    assert reds_dataset.data_infos[0] == dict(
+        lq_path=str(root_path),
+        gt_path=str(root_path),
+        key='000',
+        sequence_length=100,
+        num_input_frames=15)
+
+    # REDS4 val partition
+    reds_dataset = SRREDSMultipleGTDataset(
+        lq_folder=root_path,
+        gt_folder=root_path,
+        num_input_frames=20,
+        pipeline=[],
+        scale=4,
+        val_partition='REDS4',
+        test_mode=False)
+
+    assert len(reds_dataset.data_infos) == 266  # 266 training clips
+    assert reds_dataset.data_infos[0] == dict(
+        lq_path=str(root_path),
+        gt_path=str(root_path),
+        key='001',
+        sequence_length=100,
+        num_input_frames=20)  # 000 is been removed
+
+    with pytest.raises(ValueError):
+        # wrong val_partitaion
+        reds_dataset = SRREDSMultipleGTDataset(
+            lq_folder=root_path,
+            gt_folder=root_path,
+            num_input_frames=5,
+            pipeline=[],
+            scale=4,
+            val_partition='wrong_val_partition',
+            test_mode=False)
+
+    # test mode
+    # official val partition
+    reds_dataset = SRREDSMultipleGTDataset(
+        lq_folder=root_path,
+        gt_folder=root_path,
+        num_input_frames=5,
+        pipeline=[],
+        scale=4,
+        val_partition='official',
+        test_mode=True)
+
+    assert len(reds_dataset.data_infos) == 30  # 30 test clips
+    assert reds_dataset.data_infos[0] == dict(
+        lq_path=str(root_path),
+        gt_path=str(root_path),
+        key='240',
+        sequence_length=100,
+        num_input_frames=5)
+
+    # REDS4 val partition
+    reds_dataset = SRREDSMultipleGTDataset(
+        lq_folder=root_path,
+        gt_folder=root_path,
+        num_input_frames=5,
+        pipeline=[],
+        scale=4,
+        val_partition='REDS4',
+        test_mode=True)
+
+    assert len(reds_dataset.data_infos) == 4  # 4 test clips
+    assert reds_dataset.data_infos[1] == dict(
+        lq_path=str(root_path),
+        gt_path=str(root_path),
+        key='011',
+        sequence_length=100,
+        num_input_frames=5)
+
+
+def test_sr_vimeo90k_mutiple_gt_dataset():
+    root_path = Path(__file__).parent / 'data/vimeo90k'
+
+    txt_content = ('00001/0266 (256,448,3)\n')
+    mocked_open_function = mock_open(read_data=txt_content)
+    lq_paths = [
+        str(root_path / '00001' / '0266' / f'im{v}.png') for v in range(1, 8)
+    ]
+    gt_paths = [
+        str(root_path / '00001' / '0266' / f'im{v}.png') for v in range(1, 8)
+    ]
+
+    with patch('builtins.open', mocked_open_function):
+        vimeo90k_dataset = SRVimeo90KMultipleGTDataset(
+            lq_folder=root_path,
+            gt_folder=root_path,
+            ann_file='fake_ann_file',
+            pipeline=[],
+            scale=4,
+            test_mode=False)
+
+        assert vimeo90k_dataset.data_infos == [
+            dict(lq_path=lq_paths, gt_path=gt_paths, key='00001/0266')
+        ]
+
+
+def test_sr_test_multiple_gt_dataset():
+    root_path = Path(__file__).parent / 'data/test_multiple_gt'
+
+    test_dataset = SRTestMultipleGTDataset(
+        lq_folder=root_path,
+        gt_folder=root_path,
+        pipeline=[],
+        scale=4,
+        test_mode=True)
+
+    assert test_dataset.data_infos == [
+        dict(
+            lq_path=str(root_path),
+            gt_path=str(root_path),
+            key='sequence_1',
+            sequence_length=2),
+        dict(
+            lq_path=str(root_path),
+            gt_path=str(root_path),
+            key='sequence_2',
+            sequence_length=1)
+    ]