This commit is contained in:
zeyinzi.jzyz
2024-01-19 00:44:01 +08:00
parent 88d3322612
commit 47c528360d
149 changed files with 24224 additions and 271 deletions
+2 -1
View File
@@ -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
+6 -1
View File
@@ -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]
+122 -2
View File
@@ -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)
+25
View File
@@ -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