Source code for mmedit.datasets.sr_folder_dataset

# Copyright (c) OpenMMLab. All rights reserved.
import os.path as osp

from .base_sr_dataset import BaseSRDataset
from .registry import DATASETS

[docs]@DATASETS.register_module() class SRFolderDataset(BaseSRDataset): """General paired image folder dataset for image restoration. The dataset loads lq (Low Quality) and gt (Ground-Truth) image pairs, applies specified transforms and finally returns a dict containing paired data and other information. This is the "folder mode", which needs to specify the lq folder path and gt folder path, each folder containing the corresponding images. Image lists will be generated automatically. You can also specify the filename template to match the lq and gt pairs. For example, we have two folders with the following structures: :: data_root ├── lq │ ├── 0001_x4.png │ ├── 0002_x4.png ├── gt │ ├── 0001.png │ ├── 0002.png then, you need to set: .. code-block:: python lq_folder = data_root/lq gt_folder = data_root/gt filename_tmpl = '{}_x4' 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: `False`. filename_tmpl (str): Template for each filename. Note that the template excludes the file extension. Default: '{}'. """ def __init__(self, lq_folder, gt_folder, pipeline, scale, test_mode=False, filename_tmpl='{}'): super().__init__(pipeline, scale, test_mode) self.lq_folder = str(lq_folder) self.gt_folder = str(gt_folder) self.filename_tmpl = filename_tmpl self.data_infos = self.load_annotations()
[docs] def load_annotations(self): """Load annoations for SR dataset. It loads the LQ and GT image path from folders. Returns: list[dict]: A list of dicts for paired paths of LQ and GT. """ data_infos = [] lq_paths = self.scan_folder(self.lq_folder) gt_paths = self.scan_folder(self.gt_folder) assert len(lq_paths) == len(gt_paths), ( f'gt and lq datasets have different number of images: ' f'{len(lq_paths)}, {len(gt_paths)}.') for gt_path in gt_paths: basename, ext = osp.splitext(osp.basename(gt_path)) lq_path = osp.join(self.lq_folder, (f'{self.filename_tmpl.format(basename)}' f'{ext}')) assert lq_path in lq_paths, f'{lq_path} is not in lq_paths.' data_infos.append(dict(lq_path=lq_path, gt_path=gt_path)) return data_infos
Read the Docs v: v0.13.0
On Read the Docs
Project Home

Free document hosting provided by Read the Docs.