`mmagic.datasets`¶

Package Contents¶

Classes¶

`BasicConditionalDataset`	Custom dataset for conditional GAN. This class is based on the
`BasicFramesDataset`	BasicFramesDataset for open source projects in OpenMMLab/MMagic.
`BasicImageDataset`	BasicImageDataset for open source projects in OpenMMLab/MMagic.
`CIFAR10`	CIFAR10 Dataset.
`AdobeComp1kDataset`	Adobe composition-1k dataset.
`ControlNetDataset`	Demo dataset to test ControlNet. Modified from https://github.com/lllyas
`DreamBoothDataset`	Dataset for DreamBooth.
`GrowScaleImgDataset`	Grow Scale Unconditional Image Dataset.
`ImageNet`	ImageNet Dataset.
`MSCoCoDataset`	MSCoCo 2014 dataset.
`PairedImageDataset`	General paired image folder dataset for image generation.
`SinGANDataset`	SinGAN Dataset.
`TextualInversionDataset`	Dataset for Textual Inversion and ViCo.
`UnpairedImageDataset`	General unpaired image folder dataset for image generation.

class mmagic.datasets.BasicConditionalDataset(ann_file: str = '', metainfo: Optional[dict] = None, data_root: str = '', data_prefix: Union[str, dict] = '', extensions: Sequence[str] = ('.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif'), lazy_init: bool = False, classes: Union[str, Sequence[str], None] = None, **kwargs)[source]¶

Bases: mmengine.dataset.BaseDataset

Custom dataset for conditional GAN. This class is based on the combination of BaseDataset (https://github.com/open- mmlab/mmclassification/blob/main/mmcls/datasets/base_dataset.py) # noqa and CustomDataset (https://github.com/open- mmlab/mmclassification/blob/main/mmcls/datasets/custom.py). # noqa.

The dataset supports two kinds of annotation format.

A annotation file read by line (e.g., txt) is provided, and each line indicates a sample:

The sample files:
```
data_prefix/
├── folder_1
│   ├── xxx.png
│   ├── xxy.png
│   └── ...
└── folder_2
    ├── 123.png
    ├── nsdf3.png
    └── ...
```
The annotation file (the first column is the image path and the second column is the index of category):
```
folder_1/xxx.png 0
folder_1/xxy.png 1
folder_2/123.png 5
folder_2/nsdf3.png 3
...
```
Please specify the name of categories by the argument classes or metainfo.

A dict-based annotation file (e.g., json) is provided, key and value indicate the path and label of the sample:

The sample files:

data_prefix/
├── folder_1
│   ├── xxx.png
│   ├── xxy.png
│   └── ...
└── folder_2
    ├── 123.png
    ├── nsdf3.png
    └── ...

The annotation file (the key is the image path and the value column is the label):

{
    "folder_1/xxx.png": [1, 2, 3, 4],
    "folder_1/xxy.png": [2, 4, 1, 0],
    "folder_2/123.png": [0, 9, 8, 1],
    "folder_2/nsdf3.png", [1, 0, 0, 2],
    ...
}

In this kind of annotation, labels can be any type and not restricted to an index.

The samples are arranged in the specific way:

data_prefix/
├── class_x
│   ├── xxx.png
│   ├── xxy.png
│   └── ...
│       └── xxz.png
└── class_y
    ├── 123.png
    ├── nsdf3.png
    ├── ...
    └── asd932_.png

If the ann_file is specified, the dataset will be generated by the first two ways, otherwise, try the third way.

Parameters

ann_file (str) – Annotation file path. Defaults to ‘’.
metainfo (dict, optional) – Meta information for dataset, such as class information. Defaults to None.
data_root (str) – The root directory for data_prefix and ann_file. Defaults to ‘’.
data_prefix (str | dict) – Prefix for the data. Defaults to ‘’.
extensions (Sequence[str]) – A sequence of allowed extensions. Defaults to (‘.jpg’, ‘.jpeg’, ‘.png’, ‘.ppm’, ‘.bmp’, ‘.pgm’, ‘.tif’).
lazy_init (bool) – Whether to load annotation during instantiation. In some cases, such as visualization, only the meta information of the dataset is needed, which is not necessary to load annotation file. Basedataset can skip load annotations to save time by set lazy_init=False. Defaults to False.
**kwargs – Other keyword arguments in BaseDataset.

property img_prefix¶: The prefix of images.

property CLASSES¶: Return all categories names.

property class_to_idx¶

Map mapping class name to class index.

Returns: mapping from class name to class index.
Return type: dict

_find_samples(file_backend)[source]¶: find samples from data_prefix.

load_data_list()[source]¶: Load image paths and gt_labels.

is_valid_file(filename: str) → bool[source]¶: Check if a file is a valid sample.

get_gt_labels()[source]¶

Get all ground-truth labels (categories).

Returns: categories for all images.
Return type: np.ndarray

get_cat_ids(idx: int) → List[int][source]¶

Get category id by index.

Parameters: idx (int) – Index of data.
Returns: Image category of specified index.
Return type: cat_ids (List[int])

_compat_classes(metainfo, classes)[source]¶: Merge the old style classes arguments to metainfo.

full_init()[source]¶: Load annotation file and set BaseDataset._fully_initialized to True.

__repr__()[source]¶

Print the basic information of the dataset.

Returns: Formatted string.
Return type: str

extra_repr() → List[str][source]¶: The extra repr information of the dataset.

class mmagic.datasets.BasicFramesDataset(ann_file: str = '', metainfo: Optional[dict] = None, data_root: Optional[str] = None, data_prefix: dict = dict(img=''), pipeline: List[Union[dict, Callable]] = [], test_mode: bool = False, filename_tmpl: dict = dict(), search_key: Optional[str] = None, backend_args: Optional[dict] = None, depth: int = 1, num_input_frames: Optional[int] = None, num_output_frames: Optional[int] = None, fixed_seq_len: Optional[int] = None, load_frames_list: dict = dict(), **kwargs)[source]¶

Bases: mmengine.dataset.BaseDataset

BasicFramesDataset for open source projects in OpenMMLab/MMagic.

This dataset is designed for low-level vision tasks with frames, such as video super-resolution and video frame interpolation.

The annotation file is optional.

If use annotation file, the annotation format can be shown as follows.

Case 1 (Vid4):

    calendar 41
    city 34
    foliage 49
    walk 47

Case 2 (REDS):

    000/00000000.png (720, 1280, 3)
    000/00000001.png (720, 1280, 3)

Case 3 (Vimeo90k):

    00001/0266 (256, 448, 3)
    00001/0268 (256, 448, 3)

Parameters

ann_file (str) – Annotation file path. Defaults to ‘’.
metainfo (dict, optional) – Meta information for dataset, such as class information. Defaults to None.
data_root (str, optional) – The root directory for data_prefix and ann_file. Defaults to None.
data_prefix (dict, optional) – Prefix for training data. Defaults to dict(img=’’, gt=’’).
pipeline (list, optional) – Processing pipeline. Defaults to [].
test_mode (bool, optional) – test_mode=True means in test phase. Defaults to False.
filename_tmpl (str) – Template for each filename. Note that the template excludes the file extension. Default: ‘{}’.
search_key (str) – The key used for searching the folder to get data_list. Default: ‘gt’.
backend_args (dict, optional) – Arguments to instantiate the prefix of uri corresponding backend. Defaults to None.
depth (int) – The depth of path. Default: 1
num_input_frames (None | int) – Number of input frames. Default: None.
num_output_frames (None | int) – Number of output frames. Default: None.
fixed_seq_len (None | int) – The fixed sequence length. If None, BasicFramesDataset will obtain the length of each sequence. Default: None.
load_frames_list (dict) – Load frames list for each key. Default: dict().

Examples

Assume the file structure as the following:

mmagic (root) ├── mmagic ├── tools ├── configs ├── data │ ├── Vid4 │ │ ├── BIx4 │ │ │ ├── city │ │ │ │ ├── img1.png │ │ ├── GT │ │ │ ├── city │ │ │ │ ├── img1.png │ │ ├── meta_info_Vid4_GT.txt │ ├── places │ │ ├── sequences | | | ├── 00001 │ │ │ │ ├── 0389 │ │ │ │ │ ├── img1.png │ │ │ │ │ ├── img2.png │ │ │ │ │ ├── img3.png │ │ ├── tri_trainlist.txt

Case 1: Loading Vid4 dataset for training a VSR model.

dataset = BasicFramesDataset(
    ann_file='meta_info_Vid4_GT.txt',
    metainfo=dict(dataset_type='vid4', task_name='vsr'),
    data_root='data/Vid4',
    data_prefix=dict(img='BIx4', gt='GT'),
    pipeline=[],
    depth=2,
    num_input_frames=5)

Case 2: Loading Vimeo90k dataset for training a VFI model.

dataset = BasicFramesDataset(
    ann_file='tri_trainlist.txt',
    metainfo=dict(dataset_type='vimeo90k', task_name='vfi'),
    data_root='data/vimeo-triplet',
    data_prefix=dict(img='sequences', gt='sequences'),
    pipeline=[],
    depth=2,
    load_frames_list=dict(
        img=['img1.png', 'img3.png'], gt=['img2.png']))

See more details in unittest

tests/test_datasets/test_base_frames_dataset.py: TestFramesDatasets().test_version_1_method()

METAINFO¶

load_data_list() → List[dict][source]¶

Load data list from folder or annotation file.

Returns: A list of annotation.
Return type: list[dict]

_get_path_list()[source]¶

Get list of paths from annotation file or folder of dataset.

Returns: A list of paths.
Return type: list[str]

_get_path_list_from_ann()[source]¶

Get list of paths from annotation file.

Returns: A list of paths.
Return type: list[str]

_get_path_list_from_folder(sub_folder=None, need_ext=True, depth=1)[source]¶

Get list of paths from folder.

Parameters

sub_folder (None | str) – The path of sub_folder. Default: None.
need_ext (bool) – Whether need ext. Default: True.
depth (int) – Residual depth of path, recursively called to depth == 1. Default: 1

Returns

A list of paths.

Return type

list[str]

_set_seq_lens()[source]¶: Get sequence lengths.

_get_frames_list(key, folder)[source]¶

Obtain list of frames.

Parameters

key (str) – The key of frames list, e.g. img, gt.
folder (str) – Folder of frames.

Returns

The paths list of frames.

Return type

list[str]

class mmagic.datasets.BasicImageDataset(ann_file: str = '', metainfo: Optional[dict] = None, data_root: Optional[str] = None, data_prefix: dict = dict(img=''), pipeline: List[Union[dict, Callable]] = [], test_mode: bool = False, filename_tmpl: dict = dict(), search_key: Optional[str] = None, backend_args: Optional[dict] = None, img_suffix: Optional[Union[str, Tuple[str]]] = IMG_EXTENSIONS, recursive: bool = False, **kwards)[source]¶

Bases: mmengine.dataset.BaseDataset

BasicImageDataset for open source projects in OpenMMLab/MMagic.

This dataset is designed for low-level vision tasks with image, such as super-resolution and inpainting.

The annotation file is optional.

If use annotation file, the annotation format can be shown as follows.

Case 1 (CelebA-HQ):

    000001.png
    000002.png

Case 2 (DIV2K):

    0001_s001.png (480,480,3)
    0001_s002.png (480,480,3)
    0001_s003.png (480,480,3)
    0002_s001.png (480,480,3)
    0002_s002.png (480,480,3)

Case 3 (Vimeo90k):

    00001/0266 (256, 448, 3)
    00001/0268 (256, 448, 3)

Parameters

ann_file (str) – Annotation file path. Defaults to ‘’.
metainfo (dict, optional) – Meta information for dataset, such as class information. Defaults to None.
data_root (str, optional) – The root directory for data_prefix and ann_file. Defaults to None.
data_prefix (dict, optional) – Prefix for training data. Defaults to dict(img=None, ann=None).
pipeline (list, optional) – Processing pipeline. Defaults to [].
test_mode (bool, optional) – test_mode=True means in test phase. Defaults to False.
filename_tmpl (dict) – Template for each filename. Note that the template excludes the file extension. Default: dict().
search_key (str) – The key used for searching the folder to get data_list. Default: ‘gt’.
backend_args (dict, optional) – Arguments to instantiate the prefix of uri corresponding backend. Defaults to None.
suffix (str or tuple[str], optional) – File suffix that we are interested in. Default: None.
recursive (bool) – If set to True, recursively scan the directory. Default: False.

Note

Assume the file structure as the following:

mmagic (root)
├── mmagic
├── tools
├── configs
├── data
│   ├── DIV2K
│   │   ├── DIV2K_train_HR
│   │   │   ├── image.png
│   │   ├── DIV2K_train_LR_bicubic
│   │   │   ├── X2
│   │   │   ├── X3
│   │   │   ├── X4
│   │   │   │   ├── image_x4.png
│   │   ├── DIV2K_valid_HR
│   │   ├── DIV2K_valid_LR_bicubic
│   │   │   ├── X2
│   │   │   ├── X3
│   │   │   ├── X4
│   ├── places
│   │   ├── test_set
│   │   ├── train_set
|   |   ├── meta
|   |   |    ├── Places365_train.txt
|   |   |    ├── Places365_val.txt

Examples

Case 1: Loading DIV2K dataset for training a SISR model.

dataset = BasicImageDataset(
    ann_file='',
    metainfo=dict(
        dataset_type='div2k',
        task_name='sisr'),
    data_root='data/DIV2K',
    data_prefix=dict(
        gt='DIV2K_train_HR', img='DIV2K_train_LR_bicubic/X4'),
    filename_tmpl=dict(img='{}_x4', gt='{}'),
    pipeline=[])

Case 2: Loading places dataset for training an inpainting model.

dataset = BasicImageDataset(
    ann_file='meta/Places365_train.txt',
    metainfo=dict(
        dataset_type='places365',
        task_name='inpainting'),
    data_root='data/places',
    data_prefix=dict(gt='train_set'),
    pipeline=[])

METAINFO¶

load_data_list() → List[dict][source]¶

Load data list from folder or annotation file.

Returns: A list of annotation.
Return type: list[dict]

_get_path_list()[source]¶

Get list of paths from annotation file or folder of dataset.

Returns: A list of paths.
Return type: list[dict]

_get_path_list_from_ann()[source]¶

Get list of paths from annotation file.

Returns: List of paths.
Return type: List

_get_path_list_from_folder()[source]¶

Get list of paths from folder.

Returns: List of paths.
Return type: List

class mmagic.datasets.CIFAR10(data_prefix: str, test_mode: bool, metainfo: Optional[dict] = None, data_root: str = '', download: bool = True, **kwargs)[source]¶

Bases: mmagic.datasets.basic_conditional_dataset.BasicConditionalDataset

CIFAR10 Dataset.

This implementation is modified from https://github.com/pytorch/vision/blob/master/torchvision/datasets/cifar.py

Parameters

data_prefix (str) – Prefix for data.
test_mode (bool) – test_mode=True means in test phase. It determines to use the training set or test set.
metainfo (dict, optional) – Meta information for dataset, such as categories information. Defaults to None.
data_root (str) – The root directory for data_prefix. Defaults to ‘’.
download (bool) – Whether to download the dataset if not exists. Defaults to True.
**kwargs – Other keyword arguments in BaseDataset.

base_folder = cifar-10-batches-py¶

url = https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz¶

filename = cifar-10-python.tar.gz¶

tgz_md5 = c58f30108f718f92721af3b95e74349a¶

train_list = [['data_batch_1', 'c99cafc152244af753f735de768cd75f'], ['data_batch_2',...¶

test_list = [['test_batch', '40351d587109b95175f43aff81a1287e']]¶

meta¶

METAINFO¶

load_data_list()[source]¶: Load images and ground truth labels.

_load_meta()[source]¶: Load categories information from metafile.

_check_integrity()[source]¶: Check the integrity of data files.

extra_repr() → List[str][source]¶: The extra repr information of the dataset.

class mmagic.datasets.AdobeComp1kDataset(ann_file: Optional[str] = '', metainfo: Union[collections.abc.Mapping, mmengine.config.Config, None] = None, data_root: Optional[str] = '', data_prefix: dict = dict(img_path=''), filter_cfg: Optional[dict] = None, indices: Optional[Union[int, Sequence[int]]] = None, serialize_data: bool = True, pipeline: List[Union[dict, Callable]] = [], test_mode: bool = False, lazy_init: bool = False, max_refetch: int = 1000)[source]¶

Bases: mmengine.dataset.BaseDataset

Adobe composition-1k dataset.

The dataset loads (alpha, fg, bg) data and apply specified transforms to the data. You could specify whether composite merged image online or load composited merged image in pipeline.

Example for online comp-1k dataset:

[
    {
        "alpha": 'alpha/000.png',
        "fg": 'fg/000.png',
        "bg": 'bg/000.png'
    },
    {
        "alpha": 'alpha/001.png',
        "fg": 'fg/001.png',
        "bg": 'bg/001.png'
    },
]

Example for offline comp-1k dataset:

[
    {
        "alpha": 'alpha/000.png',
        "merged": 'merged/000.png',
        "fg": 'fg/000.png',
        "bg": 'bg/000.png'
    },
    {
        "alpha": 'alpha/001.png',
        "merged": 'merged/001.png',
        "fg": 'fg/001.png',
        "bg": 'bg/001.png'
    },
]

Parameters

ann_file (str) – Annotation file path. Defaults to ‘’.
data_root (str, optional) – The root directory for data_prefix and ann_file. Defaults to None.
pipeline (list, optional) – Processing pipeline. Defaults to [].
test_mode (bool, optional) – test_mode=True means in test phase. Defaults to False.
**kwargs – Other arguments passed to mmengine.dataset.BaseDataset.

Examples

See unit-tests TODO: Move some codes in unittest here

METAINFO¶

load_data_list() → List[dict][source]¶

Load annotations from an annotation file named as self.ann_file

In order to be compatible to both new and old annotation format, we copy implementations from mmengine and do some modifications.

Returns: A list of annotation.
Return type: list[dict]

parse_data_info(raw_data_info: dict) → Union[dict, List[dict]][source]¶: Join data_root to each path in data_info.

class mmagic.datasets.ControlNetDataset(ann_file: str = 'prompt.json', data_root: str = './data/fill50k', control_key='source', image_key='target', pipeline: List[Union[dict, Callable]] = [])[source]¶

Bases: mmengine.dataset.BaseDataset

Demo dataset to test ControlNet. Modified from https://github.com/lllyas viel/ControlNet/blob/16ea3b5379c1e78a4bc8e3fc9cae8d65c42511b1/tutorial_data set.py # noqa.

You can download the demo data from https://huggingface.co/lllyasviel/ControlNet/blob/main/training/fill50k.zip # noqa and then unzip the file to the data folder.

Parameters

ann_file (str) – Path to the annotation file. Defaults to ‘prompt.json’ as ControlNet’s default.
data_root (str) – Path to the data root. Defaults to ‘./data/fill50k’.
pipeline (list[dict | callable]) – A sequence of data transforms.

load_data_list() → List[dict][source]¶

Load annotations from an annotation file named as self.ann_file

Returns: A list of annotation.
Return type: list[dict]

class mmagic.datasets.DreamBoothDataset(data_root: str, concept_dir: str, prompt: str, pipeline: List[Union[dict, Callable]] = [])[source]¶

Bases: mmengine.dataset.BaseDataset

Dataset for DreamBooth.

Parameters

data_root (str) – Path to the data root.
concept_dir (str) – Path to the concept images.
prompt (str) – Prompt of the concept.
pipeline (list[dict | callable]) – A sequence of data transforms.

load_data_list() → list[source]¶: Load data list from concept_dir and class_dir.

class mmagic.datasets.GrowScaleImgDataset(data_roots: dict, pipeline, len_per_stage=int(1000000.0), gpu_samples_per_scale=None, gpu_samples_base=32, io_backend: Optional[str] = None, file_lists: Optional[Union[str, dict]] = None, test_mode=False)[source]¶

Bases: mmengine.dataset.BaseDataset

Grow Scale Unconditional Image Dataset.

This dataset is similar with UnconditionalImageDataset, but offer more dynamic functionalities for the supporting complex algorithms, like PGGAN.

Highlight functionalities:

Support growing scale dataset. The motivation is to decrease data pre-processing load in CPU. In this dataset, you can provide imgs_roots like:
{'64': 'path_to_64x64_imgs', '512': 'path_to_512x512_imgs'}
Then, in training scales lower than 64x64, this dataset will set self.imgs_root as ‘path_to_64x64_imgs’;
Offer samples_per_gpu according to different scales. In this dataset, self.samples_per_gpu will help runner to know the updated batch size.

Basically, This dataset contains raw images for training unconditional GANs. Given a root dir, we will recursively find all images in this root. The transformation on data is defined by the pipeline.

Parameters

imgs_root (str) – Root path for unconditional images.
pipeline (list[dict | callable]) – A sequence of data transforms.
len_per_stage (int, optional) – The length of dataset for each scale. This args change the length dataset by concatenating or extracting subset. If given a value less than 0., the original length will be kept. Defaults to 1e6.
gpu_samples_per_scale (dict | None, optional) – Dict contains samples_per_gpu for each scale. For example, {'32': 4} will set the scale of 32 with samples_per_gpu=4, despite other scale with samples_per_gpu=self.gpu_samples_base.
gpu_samples_base (int, optional) – Set default samples_per_gpu for each scale. Defaults to 32.
io_backend (str, optional) – The storage backend type. Options are “disk”, “ceph”, “memcached”, “lmdb”, “http” and “petrel”. Default: None.
test_mode (bool, optional) – If True, the dataset will work in test mode. Otherwise, in train mode. Default to False.

_VALID_IMG_SUFFIX = ('.jpg', '.png', '.jpeg', '.JPEG')¶

load_data_list()[source]¶: Load annotations.

update_annotations(curr_scale)[source]¶

Update annotations.

Parameters: curr_scale (int) – Current image scale.
Returns: Whether to update.
Return type: bool

concat_imgs_list_to(num)[source]¶

Concat image list to specified length.

Parameters: num (int) – The length of the concatenated image list.

prepare_train_data(idx)[source]¶

Prepare training data.

Parameters: idx (int) – Index of current batch.
Returns: Prepared training data batch.
Return type: dict

prepare_test_data(idx)[source]¶

Prepare testing data.

Parameters: idx (int) – Index of current batch.
Returns: Prepared training data batch.
Return type: dict

__getitem__(idx)[source]¶

Get the idx-th image and data information of dataset after self.pipeline, and full_init will be called if the dataset has not been fully initialized.

During training phase, if self.pipeline get None, self._rand_another will be called until a valid image is fetched or

the maximum limit of refetch is reached.

Parameters: idx (int) – The index of self.data_list.
Returns: The idx-th image and data information of dataset after self.pipeline.
Return type: dict

__repr__()[source]¶

Print self.transforms in sequence.

Returns: Formatted string.
Return type: str

class mmagic.datasets.ImageNet(ann_file: str = '', metainfo: Optional[dict] = None, data_root: str = '', data_prefix: Union[str, dict] = '', **kwargs)[source]¶

Bases: mmagic.datasets.basic_conditional_dataset.BasicConditionalDataset

ImageNet Dataset.

The dataset supports two kinds of annotation format. More details can be found in CustomDataset.

Parameters

ann_file (str) – Annotation file path. Defaults to ‘’.
metainfo (dict, optional) – Meta information for dataset, such as class information. Defaults to None.
data_root (str) – The root directory for data_prefix and ann_file. Defaults to ‘’.
data_prefix (str | dict) – Prefix for training data. Defaults to ‘’.
**kwargs – Other keyword arguments in CustomDataset and BaseDataset.

IMG_EXTENSIONS = ('.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif')¶

METAINFO¶

class mmagic.datasets.MSCoCoDataset(ann_file: str = '', metainfo: Optional[dict] = None, data_root: str = '', drop_caption_rate=0.0, phase='train', year=2014, data_prefix: Union[str, dict] = '', extensions: Sequence[str] = ('.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif'), lazy_init: bool = False, classes: Union[str, Sequence[str], None] = None, caption_style: str = '', **kwargs)[source]¶

Bases: mmagic.datasets.basic_conditional_dataset.BasicConditionalDataset

MSCoCo 2014 dataset.

Parameters

ann_file (str) – Annotation file path. Defaults to ‘’.
metainfo (dict, optional) – Meta information for dataset, such as class information. Defaults to None.
data_root (str) – The root directory for data_prefix and ann_file. Defaults to ‘’.
drop_caption_rate (float, optional) – Rate of dropping caption, used for training. Defaults to 0.0.
phase (str, optional) – Subdataset used for certain phase, can be set to train, test and val. Defaults to ‘train’.
year (int, optional) – Version of CoCo dataset, can be set to 2014 and 2017. Defaults to 2014.
data_prefix (str | dict) – Prefix for the data. Defaults to ‘’.
extensions (Sequence[str]) – A sequence of allowed extensions. Defaults to (‘.jpg’, ‘.jpeg’, ‘.png’, ‘.ppm’, ‘.bmp’, ‘.pgm’, ‘.tif’).
lazy_init (bool) – Whether to load annotation during instantiation. In some cases, such as visualization, only the meta information of the dataset is needed, which is not necessary to load annotation file. Basedataset can skip load annotations to save time by set lazy_init=False. Defaults to False.
caption_style (str) – If you want to add a style description for each caption, you can set caption_style to your style prompt. For example, ‘realistic style’. Defaults to empty str.
**kwargs – Other keyword arguments in BaseDataset.

METAINFO¶

load_data_list()[source]¶: Load image paths and gt_labels.

class mmagic.datasets.PairedImageDataset(data_root, pipeline, io_backend: Optional[str] = None, test_mode=False, test_dir='test')[source]¶

Bases: mmengine.dataset.BaseDataset

General paired image folder dataset for image generation.

It assumes that the training directory is ‘/path/to/data/train’. During test time, the directory is ‘/path/to/data/test’. ‘/path/to/data’ can be initialized by args ‘dataroot’. Each sample contains a pair of images concatenated in the w dimension (A|B).

Parameters

dataroot (str | Path) – Path to the folder root of paired images.
pipeline (List[dict | callable]) – A sequence of data transformations.
test_mode (bool) – Store True when building test dataset. Default: False.
test_dir (str) – Subfolder of dataroot which contain test images. Default: ‘test’.

load_data_list()[source]¶

Load paired image paths.

Returns: List that contains paired image paths.
Return type: list[dict]

scan_folder(path)[source]¶

Obtain image path list (including sub-folders) from a given folder.

Parameters: path (str | Path) – Folder path.
Returns: Image list obtained from the given folder.
Return type: list[str]

class mmagic.datasets.SinGANDataset(data_root, min_size, max_size, scale_factor_init, pipeline, num_samples=- 1)[source]¶

Bases: mmengine.dataset.BaseDataset

SinGAN Dataset.

In this dataset, we create an image pyramid and save it in the cache.

Parameters

img_path (str) – Path to the single image file.
min_size (int) – Min size of the image pyramid. Here, the number will be set to the min(H, W).
max_size (int) – Max size of the image pyramid. Here, the number will be set to the max(H, W).
scale_factor_init (float) – Rescale factor. Note that the actual factor we use may be a little bit different from this value.
num_samples (int, optional) – The number of samples (length) in this dataset. Defaults to -1.

full_init()[source]¶: Skip the full init process for SinGANDataset.

load_data_list(min_size, max_size, scale_factor_init)[source]¶

Load annotations for SinGAN Dataset.

Parameters

min_size (int) – The minimum size for the image pyramid.
max_size (int) – The maximum size for the image pyramid.
scale_factor_init (float) – The initial scale factor.

__getitem__(index)[source]¶

Get :attr:self.data_dict. For SinGAN, we use single image with different resolution to train the model.

Parameters: idx (int) – This will be ignored in SinGANDataset.
Returns: Dict contains input image in different resolution. self.pipeline.
Return type: dict

__len__()[source]¶

Get the length of filtered dataset and automatically call full_init if the dataset has not been fully init.

Returns: The length of filtered dataset.
Return type: int

class mmagic.datasets.TextualInversionDataset(data_root: str, concept_dir: str, placeholder: str, template: str, with_image_reference: bool = False, pipeline: List[Union[dict, Callable]] = [])[source]¶

Bases: mmengine.dataset.BaseDataset

Dataset for Textual Inversion and ViCo.

Parameters

data_root (str) – Path to the data root.
concept_dir (str) – Path to the concept images.
placeholder (str) – A string to denote the concept.
template (list[str]) – A list of strings like ‘A photo of {}’.
with_image_reference (bool) – Is used for vico training.
pipeline (list[dict | callable]) – A sequence of data transforms.

load_data_list() → list[source]¶: Load data list from concept_dir and class_dir.

prepare_data(idx)[source]¶

Get data processed by self.pipeline.

Parameters: idx (int) – The index of data_info.
Returns: Depends on self.pipeline.
Return type: Any

class mmagic.datasets.UnpairedImageDataset(data_root, pipeline, io_backend: Optional[str] = None, test_mode=False, domain_a='A', domain_b='B')[source]¶

Bases: mmengine.dataset.BaseDataset

General unpaired image folder dataset for image generation.

It assumes that the training directory of images from domain A is ‘/path/to/data/trainA’, and that from domain B is ‘/path/to/data/trainB’, respectively. ‘/path/to/data’ can be initialized by args ‘dataroot’. During test time, the directory is ‘/path/to/data/testA’ and ‘/path/to/data/testB’, respectively.

Parameters

dataroot (str | Path) – Path to the folder root of unpaired images.
pipeline (List[dict | callable]) – A sequence of data transformations.
io_backend (str, optional) – The storage backend type. Options are “disk”, “ceph”, “memcached”, “lmdb”, “http” and “petrel”. Default: None.
test_mode (bool) – Store True when building test dataset. Default: False.
domain_a (str, optional) – Domain of images in trainA / testA. Defaults to ‘A’.
domain_b (str, optional) – Domain of images in trainB / testB. Defaults to ‘B’.

load_data_list()[source]¶

Load the data list.

Returns: The data info list of source and target domain.
Return type: list

_load_domain_data_list(dataroot)[source]¶

Load unpaired image paths of one domain.

Parameters: dataroot (str) – Path to the folder root for unpaired images of one domain.
Returns: List that contains unpaired image paths of one domain.
Return type: list[dict]

get_data_info(idx) → dict[source]¶

Get annotation by index and automatically call full_init if the dataset has not been fully initialized.

Parameters: idx (int) – The index of data.
Returns: The idx-th annotation of the dataset.
Return type: dict

__len__()[source]¶: The length of the dataset.

scan_folder(path)[source]¶

Obtain image path list (including sub-folders) from a given folder.

Parameters: path (str | Path) – Folder path.
Returns: Image list obtained from the given folder.
Return type: list[str]

mmagic.datasets¶

Package Contents¶

Classes¶

`mmagic.datasets`¶