mmagic.datasets.grow_scale_image_dataset

Module Contents

Classes

GrowScaleImgDataset

Grow Scale Unconditional Image Dataset.

class mmagic.datasets.grow_scale_image_dataset.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)

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:

1. 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’;

2. 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.

参数

• 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()

Load annotations.

update_annotations(curr_scale)

Update annotations.

参数

curr_scale (int) – Current image scale.

返回

Whether to update.

返回类型

bool

concat_imgs_list_to(num)

Concat image list to specified length.

参数

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

prepare_train_data(idx)

Prepare training data.

参数

idx (int) – Index of current batch.

返回

Prepared training data batch.

返回类型

dict

prepare_test_data(idx)

Prepare testing data.

参数

idx (int) – Index of current batch.

返回

Prepared training data batch.

返回类型

dict

__getitem__(idx)

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.

参数

idx (int) – The index of self.data_list.

返回

The idx-th image and data information of dataset after self.pipeline.

返回类型

dict

__repr__()

Print self.transforms in sequence.

返回

Formatted string.

返回类型

str