v0.0.2
This commit is contained in:
@@ -6,5 +6,6 @@ from scepter.modules.data.dataset.dataset import (Image2ImageDataset,
|
||||
ImageClassifyPublicDataset,
|
||||
ImageTextPairDataset,
|
||||
Text2ImageDataset)
|
||||
from scepter.modules.data.dataset.ms_dataset import ImageTextPairMSDataset
|
||||
from scepter.modules.data.dataset.ms_dataset import (
|
||||
ImageTextPairFolderDataset, ImageTextPairMSDataset)
|
||||
from scepter.modules.data.dataset.registry import DATASETS
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
|
||||
import numbers
|
||||
import os
|
||||
import sys
|
||||
from collections.abc import Iterable
|
||||
|
||||
@@ -235,6 +236,7 @@ class Text2ImageDataset(BaseDataset):
|
||||
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)
|
||||
@@ -257,11 +259,14 @@ class Text2ImageDataset(BaseDataset):
|
||||
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)
|
||||
elif key != 'meta':
|
||||
item[key] = value
|
||||
else:
|
||||
continue
|
||||
|
||||
self.items.append(item)
|
||||
if use_num > 0:
|
||||
self.items = self.items[:use_num]
|
||||
|
||||
@@ -9,6 +9,7 @@ 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 FS
|
||||
|
||||
|
||||
@DATASETS.register_class()
|
||||
@@ -105,14 +106,17 @@ class ImageTextPairMSDataset(BaseDataset):
|
||||
raise (
|
||||
'Your must set MS_DATASET_NAME as modelscope dataset or your local dataset orignized '
|
||||
'as modelscope dataset.')
|
||||
if FS.exists(ms_dataset_name):
|
||||
ms_dataset_name = FS.get_dir_to_local_dir(ms_dataset_name)
|
||||
ms_remap_path = ms_dataset_name
|
||||
try:
|
||||
self.data = MsDataset.load(str(ms_dataset_name),
|
||||
namespace=ms_dataset_namespace,
|
||||
subset_name=ms_dataset_subname,
|
||||
split=ms_dataset_split)
|
||||
except Exception as e:
|
||||
except Exception:
|
||||
self.logger.info(
|
||||
f"Load Modelscope dataset failed with {e}, retry with download_mode='force_redownload'."
|
||||
"Load Modelscope dataset failed, retry with download_mode='force_redownload'."
|
||||
)
|
||||
try:
|
||||
self.data = MsDataset.load(
|
||||
@@ -177,3 +181,119 @@ class ImageTextPairMSDataset(BaseDataset):
|
||||
__class__.__name__,
|
||||
ImageTextPairMSDataset.para_dict,
|
||||
set_name=True)
|
||||
|
||||
|
||||
@DATASETS.register_class()
|
||||
class ImageTextPairFolderDataset(BaseDataset):
|
||||
para_dict = {
|
||||
'DATA_FOLDER': {
|
||||
'value': '',
|
||||
'description': 'Dataset folder.'
|
||||
},
|
||||
'TRIGGER_WORDS': {
|
||||
'value':
|
||||
'',
|
||||
'description':
|
||||
'The words used to describe the common features of your data, especially when you customize a '
|
||||
'tuner. Use these words you can get what you want.'
|
||||
},
|
||||
'REPLACE_STYLE': {
|
||||
'value':
|
||||
False,
|
||||
'description':
|
||||
'Whether use the MS_DATASET_SUBNAME to replace the word in your description, default is False.'
|
||||
},
|
||||
'HIGHLIGHT_KEYWORDS': {
|
||||
'value':
|
||||
'',
|
||||
'description':
|
||||
'The keywords you want to highlight in prompt, which will be replace by <HIGHLIGHT_KEYWORDS>.'
|
||||
},
|
||||
'KEYWORDS_SIGN': {
|
||||
'value':
|
||||
'',
|
||||
'description':
|
||||
'The keywords sign you want to add, which is like <{HIGHLIGHT_KEYWORDS}{KEYWORDS_SIGN}>'
|
||||
},
|
||||
'OUTPUT_SIZE': {
|
||||
'value':
|
||||
None,
|
||||
'description':
|
||||
'If you use the FlexibleResize transforms, this filed will output the image_size as [h, w],'
|
||||
'which will be used to set the output size of images used to train the model.'
|
||||
},
|
||||
}
|
||||
|
||||
def __init__(self, cfg, logger=None):
|
||||
super().__init__(cfg=cfg, logger=logger)
|
||||
data_folder = cfg.get('DATA_FOLDER', None)
|
||||
self.replace_style = cfg.get('REPLACE_STYLE', False)
|
||||
self.trigger_words = cfg.get('TRIGGER_WORDS', '')
|
||||
self.replace_keywords = cfg.get('HIGHLIGHT_KEYWORDS', '')
|
||||
self.keywords_sign = cfg.get('KEYWORDS_SIGN', '')
|
||||
self.output_size = cfg.get('OUTPUT_SIZE', None)
|
||||
if self.output_size is not None:
|
||||
if isinstance(self.output_size, numbers.Number):
|
||||
self.output_size = [self.output_size, self.output_size]
|
||||
# Use modelscope dataset
|
||||
if not data_folder or not FS.exists(data_folder):
|
||||
raise ('Your must set datafolder for local dataset.')
|
||||
data_folder = FS.get_dir_to_local_dir(data_folder)
|
||||
all_lines = open(os.path.join(data_folder, 'train.csv'),
|
||||
'r').read().split('\n')
|
||||
assert all_lines[0] == 'Target:FILE,Prompt'
|
||||
self.data = []
|
||||
for line in all_lines[1:]:
|
||||
line = line.strip()
|
||||
if line == '':
|
||||
continue
|
||||
self.data.append({
|
||||
'Target:FILE':
|
||||
os.path.join(data_folder,
|
||||
line.split(',', 1)[0]),
|
||||
'Prompt':
|
||||
line.split(',', 1)[1]
|
||||
})
|
||||
self.real_number = len(self.data)
|
||||
|
||||
def __len__(self):
|
||||
if self.mode == 'train':
|
||||
return sys.maxsize
|
||||
else:
|
||||
return len(self.data)
|
||||
|
||||
def _get(self, index: int):
|
||||
current_data = self.data[index % len(self.data)]
|
||||
# print(current_data.keys())
|
||||
image_path = current_data['Target:FILE']
|
||||
prompt = current_data['Prompt']
|
||||
style = current_data['Style'] if 'Style' in current_data else ''
|
||||
# print(prompt, style)
|
||||
if self.replace_style and not style == '':
|
||||
prompt = prompt.replace(style, f'<{self.keywords_sign}>')
|
||||
elif not self.replace_keywords.strip() == '':
|
||||
prompt = prompt.replace(
|
||||
self.replace_keywords,
|
||||
'<' + self.replace_keywords + f'{self.keywords_sign}>')
|
||||
if not self.trigger_words == '':
|
||||
prompt = self.trigger_words.strip() + ' ' + prompt
|
||||
if we.debug:
|
||||
print(prompt, self.replace_keywords.strip())
|
||||
ret_item = {
|
||||
'meta': {
|
||||
'img_path': image_path,
|
||||
'data_key': style,
|
||||
'data_num': self.real_number
|
||||
},
|
||||
'prompt': prompt
|
||||
}
|
||||
if self.output_size is not None:
|
||||
ret_item['meta']['image_size'] = self.output_size
|
||||
return ret_item
|
||||
|
||||
@staticmethod
|
||||
def get_config_template():
|
||||
return dict_to_yaml('DATASet',
|
||||
__class__.__name__,
|
||||
ImageTextPairMSDataset.para_dict,
|
||||
set_name=True)
|
||||
|
||||
@@ -0,0 +1,25 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||||
|
||||
from PIL import Image
|
||||
from torch.utils.data.dataloader import default_collate
|
||||
|
||||
|
||||
def pil_collate_fn(self, batch):
|
||||
batch_data = {}
|
||||
for items in batch:
|
||||
for key, item in items.items():
|
||||
if isinstance(item, Image.Image):
|
||||
if key not in batch_data:
|
||||
batch_data[key] = []
|
||||
batch_data[key].append(item)
|
||||
else:
|
||||
if key not in batch_data:
|
||||
batch_data[key] = []
|
||||
batch_data[key].append(item)
|
||||
|
||||
for key, item in batch_data.items():
|
||||
if not all(isinstance(x, Image.Image) for x in item):
|
||||
batch_data[key] = default_collate(item)
|
||||
|
||||
return batch_data
|
||||
Reference in New Issue
Block a user