295 lines
8.9 KiB
Python
295 lines
8.9 KiB
Python
# -*- coding: utf-8 -*-
|
|
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
|
|
import numbers
|
|
import os
|
|
import sys
|
|
from collections.abc import Iterable
|
|
|
|
import numpy as np
|
|
import torchvision
|
|
|
|
from scepter.modules.data.dataset.base_dataset import BaseDataset
|
|
from scepter.modules.data.dataset.registry import DATASETS
|
|
from scepter.modules.utils.config import dict_to_yaml
|
|
from scepter.modules.utils.distribute import we
|
|
from scepter.modules.utils.file_system import DATA_FS as FS
|
|
|
|
|
|
@DATASETS.register_class()
|
|
class ImageClassifyPublicDataset(BaseDataset):
|
|
"""
|
|
Dataset for image classification wrapper
|
|
|
|
Args:
|
|
json_path (str): json file which contains all instances, should be a list of dict
|
|
which contains img_path and gt_label
|
|
image_dir (str or None): image directory, if None, img_path in json_path will be considered as absolute path
|
|
classes (list[str] or None): image class description
|
|
"""
|
|
para_dict = {
|
|
'DATASET': {
|
|
'value': 'cifar10',
|
|
'description': 'the public dataset name'
|
|
},
|
|
'DATA_ROOT': {
|
|
'value': '',
|
|
'description': 'the download data save path'
|
|
}
|
|
}
|
|
|
|
para_dict.update(BaseDataset.para_dict)
|
|
|
|
def __init__(self, cfg, logger=None):
|
|
|
|
super(ImageClassifyPublicDataset, self).__init__(cfg, logger=logger)
|
|
|
|
self.dataset_name = cfg.DATASET
|
|
self.data_root = cfg.DATA_ROOT
|
|
self.phase = cfg.MODE
|
|
if self.dataset_name == 'cifar10':
|
|
self.dataset = torchvision.datasets.CIFAR10(
|
|
root=self.data_root,
|
|
train=self.phase == 'train',
|
|
download=True)
|
|
|
|
def __len__(self) -> int:
|
|
return len(self.dataset)
|
|
|
|
def _get(self, index: int):
|
|
img, target = self.dataset.__getitem__(index)
|
|
ret = {
|
|
'meta': {},
|
|
'label': np.asarray(target, dtype=np.int64),
|
|
'img': img
|
|
}
|
|
return ret
|
|
|
|
def worker_init_fn(self, worker_id, num_workers=1):
|
|
super(ImageClassifyPublicDataset,
|
|
self).worker_init_fn(worker_id, num_workers=num_workers)
|
|
|
|
@staticmethod
|
|
def get_config_template():
|
|
'''
|
|
{ "ENV" :
|
|
{ "description" : "",
|
|
"A" : {
|
|
"value": 1.0,
|
|
"description": ""
|
|
}
|
|
}
|
|
}
|
|
:return:
|
|
'''
|
|
return dict_to_yaml('modename_DATA',
|
|
__class__.__name__,
|
|
ImageClassifyPublicDataset.para_dict,
|
|
set_name=True)
|
|
|
|
|
|
@DATASETS.register_class()
|
|
class ImageTextPairDataset(BaseDataset):
|
|
"""
|
|
Dataset for diffusion model training
|
|
"""
|
|
para_dict = {
|
|
'P_ZERO': {
|
|
'value': 0.0,
|
|
'description': '',
|
|
},
|
|
'NEGTIVE_PROMPT': {
|
|
'value': '',
|
|
'description': 'The default negtive prompt',
|
|
}
|
|
}
|
|
para_dict.update(BaseDataset.para_dict)
|
|
|
|
def __init__(self, cfg, logger=None):
|
|
super(ImageTextPairDataset, self).__init__(cfg, logger=logger)
|
|
self.p_zero = cfg.get('P_ZERO', 0.0)
|
|
self._default_item = {
|
|
'meta': {},
|
|
'prompt':
|
|
'Plants in the Water, Nature, Lake, Horizontal, Reflection, Photography, Backgrounds, Swamp, No People'
|
|
}
|
|
|
|
def _get(self, index):
|
|
meta = dict()
|
|
# the last item is field_keys
|
|
for key, value in zip(index[-1], index[:-1]):
|
|
if key in ['oss_key', 'path', 'img_path', 'target_img_path']:
|
|
meta['img_path'] = value
|
|
elif key in ['prompt', 'caption', 'text']:
|
|
meta['ori_prompt'] = value
|
|
elif key in ['width', 'height']:
|
|
meta[key] = int(value)
|
|
else:
|
|
meta[key] = value
|
|
|
|
prompt = meta.get('prompt_prefix', '') + meta.get('ori_prompt', '')
|
|
if self.mode == 'train' and np.random.uniform() < self.p_zero:
|
|
prompt = ''
|
|
|
|
item = {
|
|
'meta': meta,
|
|
'prompt': prompt,
|
|
}
|
|
return item
|
|
|
|
def __getitem__(self, index):
|
|
item = self._get(index)
|
|
item = self.pipeline(item)
|
|
return item
|
|
|
|
def __len__(self) -> int:
|
|
return sys.maxsize
|
|
|
|
@staticmethod
|
|
def get_config_template():
|
|
return dict_to_yaml('DATASETS',
|
|
__class__.__name__,
|
|
ImageTextPairDataset.para_dict,
|
|
set_name=True)
|
|
|
|
|
|
@DATASETS.register_class()
|
|
class Image2ImageDataset(BaseDataset):
|
|
"""
|
|
Dataset for diffusion model training
|
|
"""
|
|
para_dict = {}
|
|
para_dict.update(BaseDataset.para_dict)
|
|
|
|
def __init__(self, cfg, logger=None):
|
|
super(Image2ImageDataset, self).__init__(cfg, logger=logger)
|
|
self._default_item = {
|
|
'meta': {},
|
|
'prompt':
|
|
'Plants in the Water, Nature, Lake, Horizontal, Reflection, Photography, Backgrounds, Swamp, No People'
|
|
}
|
|
|
|
def _get(self, index):
|
|
meta = dict()
|
|
# the last item is field_keys
|
|
for key, value in zip(index[-1], index[:-1]):
|
|
if key in ['oss_key', 'path', 'img_path']:
|
|
meta['img_path'] = value
|
|
elif key in ['prompt', 'caption', 'text']:
|
|
meta['ori_prompt'] = value
|
|
elif key in ['width', 'height']:
|
|
meta[key] = int(value)
|
|
else:
|
|
meta[key] = value
|
|
item = {'meta': meta}
|
|
return item
|
|
|
|
def __getitem__(self, index):
|
|
item = self._get(index)
|
|
item = self.pipeline(item)
|
|
return item
|
|
|
|
def __len__(self) -> int:
|
|
return sys.maxsize
|
|
|
|
@staticmethod
|
|
def get_config_template():
|
|
return dict_to_yaml('DATASETS',
|
|
__class__.__name__,
|
|
Image2ImageDataset.para_dict,
|
|
set_name=True)
|
|
|
|
|
|
@DATASETS.register_class()
|
|
class Text2ImageDataset(BaseDataset):
|
|
para_dict = {
|
|
'PROMPT_FILE': {
|
|
'value': '',
|
|
'description': ''
|
|
},
|
|
'FIELDS': {
|
|
'value': '',
|
|
'description': ''
|
|
},
|
|
'DELIMITER': {
|
|
'value': ',',
|
|
'description': ''
|
|
},
|
|
'PROMPT_PREFIX': {
|
|
'value': '',
|
|
'description': ''
|
|
},
|
|
'IMAGE_SIZE': {
|
|
'value': 512,
|
|
'description': ''
|
|
},
|
|
'USE_NUM': {
|
|
'value': -1,
|
|
'description': ''
|
|
},
|
|
}
|
|
para_dict.update(BaseDataset.para_dict)
|
|
|
|
def __init__(self, cfg, logger=None):
|
|
super(Text2ImageDataset, self).__init__(cfg, logger=logger)
|
|
|
|
delimiter = cfg.get('DELIMITER', ',')
|
|
fields = cfg.get('FIELDS', ['row_key', 'prompt'])
|
|
prompt_prefix = cfg.get('PROMPT_PREFIX', '')
|
|
path_prefix = cfg.get('PATH_PREFIX', '')
|
|
use_num = cfg.get('USE_NUM', -1)
|
|
|
|
image_size = cfg.get('IMAGE_SIZE', 1024)
|
|
if isinstance(image_size, numbers.Number):
|
|
image_size = [image_size, image_size]
|
|
assert isinstance(image_size, Iterable) and len(image_size) == 2
|
|
|
|
if cfg.PROMPT_FILE is not None and cfg.PROMPT_FILE != '':
|
|
prompt_file = cfg.PROMPT_FILE
|
|
with FS.get_object(prompt_file) as local_data:
|
|
rows = [
|
|
i.split(delimiter,
|
|
len(fields) - 1)
|
|
for i in local_data.decode('utf-8').strip().split('\n')
|
|
]
|
|
else:
|
|
rows = [
|
|
i.split(delimiter,
|
|
len(fields) - 1) for i in cfg.PROMPT_DATA
|
|
]
|
|
|
|
self.items = list()
|
|
for i, row in enumerate(rows):
|
|
item = {'index': i, 'meta': {'image_size': image_size}}
|
|
for key, value in zip(fields, row):
|
|
if key in ['prompt', 'caption', 'text']:
|
|
item['ori_prompt'] = value
|
|
item['prompt'] = prompt_prefix + value
|
|
elif key in ['oss_key', 'path', 'img_path', 'target_img_path']:
|
|
item['meta']['img_path'] = os.path.join(path_prefix, value)
|
|
elif key in ['width', 'height']:
|
|
item['meta'][key] = int(value)
|
|
else:
|
|
item['meta'][key] = value
|
|
|
|
self.items.append(item)
|
|
if use_num > 0:
|
|
self.items = self.items[:use_num]
|
|
if we.rank == 0:
|
|
logger.info(f'eval prompt num: {len(self.items)}')
|
|
logger.info('eval prompts: {}'.format(
|
|
[k['prompt'] for k in self.items]))
|
|
|
|
def _get(self, index: int):
|
|
return self.items[index]
|
|
|
|
def __len__(self) -> int:
|
|
return len(self.items)
|
|
|
|
@staticmethod
|
|
def get_config_template():
|
|
return dict_to_yaml('DATASETS',
|
|
__class__.__name__,
|
|
Text2ImageDataset.para_dict,
|
|
set_name=True)
|