230 lines
8.3 KiB
Python
230 lines
8.3 KiB
Python
import json
|
|
import logging
|
|
import os
|
|
import random
|
|
import typing
|
|
import zipfile
|
|
|
|
import PIL.Image as Image
|
|
import tqdm
|
|
from colorama import Fore, Style
|
|
|
|
from data.image_train_item import ImageCaption, ImageTrainItem
|
|
|
|
class DataResolver:
|
|
def __init__(self, aspects: list[typing.Tuple[int, int]], flip_p=0.0, seed=555):
|
|
self.aspects = aspects
|
|
self.flip_p = flip_p
|
|
|
|
def image_train_items(self, data_root: str) -> list[ImageTrainItem]:
|
|
"""
|
|
Get the list of `ImageTrainItem` for the given data root.
|
|
|
|
:param data_root: The data root, a directory, a file, etc..
|
|
:return: The list of `ImageTrainItem`.
|
|
"""
|
|
raise NotImplementedError()
|
|
|
|
def image_train_item(self, image_path: str, caption: ImageCaption, multiplier: float=1) -> ImageTrainItem:
|
|
return ImageTrainItem(
|
|
image=None,
|
|
caption=caption,
|
|
aspects=self.aspects,
|
|
pathname=image_path,
|
|
flip_p=self.flip_p,
|
|
multiplier=multiplier
|
|
)
|
|
|
|
class JSONResolver(DataResolver):
|
|
def image_train_items(self, json_path: str) -> list[ImageTrainItem]:
|
|
items = []
|
|
with open(json_path, encoding='utf-8', mode='r') as f:
|
|
json_data = json.load(f)
|
|
|
|
for data in json_data:
|
|
caption = JSONResolver.image_caption(data)
|
|
if caption:
|
|
image_value = JSONResolver.get_image_value(data)
|
|
item = self.image_train_item(image_value, caption)
|
|
if item:
|
|
items.append(item)
|
|
|
|
return items
|
|
|
|
@staticmethod
|
|
def get_image_value(json_data: dict) -> typing.Optional[str]:
|
|
"""
|
|
Get the image from the json data if possible.
|
|
|
|
:param json_data: The json data, a dict.
|
|
:return: The image, or None if not found.
|
|
"""
|
|
image_value = json_data.get("image", None)
|
|
if isinstance(image_value, str):
|
|
image_value = image_value.strip()
|
|
if os.path.exists(image_value):
|
|
return image_value
|
|
|
|
@staticmethod
|
|
def get_caption_value(json_data: dict) -> typing.Optional[str]:
|
|
"""
|
|
Get the caption from the json data if possible.
|
|
|
|
:param json_data: The json data, a dict.
|
|
:return: The caption, or None if not found.
|
|
"""
|
|
caption_value = json_data.get("caption", None)
|
|
if isinstance(caption_value, str):
|
|
return caption_value.strip()
|
|
|
|
@staticmethod
|
|
def image_caption(json_data: dict) -> typing.Optional[ImageCaption]:
|
|
"""
|
|
Get the caption from the json data if possible.
|
|
|
|
:param json_data: The json data, a dict.
|
|
:return: The `ImageCaption`, or None if not found.
|
|
"""
|
|
image_value = JSONResolver.get_image_value(json_data)
|
|
caption_value = JSONResolver.get_caption_value(json_data)
|
|
if image_value:
|
|
if caption_value:
|
|
return ImageCaption.resolve(caption_value)
|
|
return ImageCaption.from_file(image_value)
|
|
|
|
|
|
class DirectoryResolver(DataResolver):
|
|
def image_train_items(self, data_root: str) -> list[ImageTrainItem]:
|
|
"""
|
|
Create `ImageTrainItem` objects with metadata for hydration later.
|
|
Unzips all zip files in `data_root` and then recursively searches the
|
|
`data_root` for images and captions.
|
|
|
|
:param data_root: The root directory to recurse through
|
|
"""
|
|
DirectoryResolver.unzip_all(data_root)
|
|
image_paths = list(DirectoryResolver.recurse_data_root(data_root))
|
|
items = []
|
|
multipliers = {}
|
|
skip_folders = []
|
|
randomizer = random.Random(self.seed)
|
|
|
|
for pathname in tqdm.tqdm(image_paths):
|
|
current_dir = os.path.dirname(pathname)
|
|
|
|
if current_dir not in multipliers:
|
|
multiply_txt_path = os.path.join(current_dir, "multiply.txt")
|
|
if os.path.exists(multiply_txt_path):
|
|
try:
|
|
with open(multiply_txt_path, 'r') as f:
|
|
val = float(f.read().strip())
|
|
multipliers[current_dir] = val
|
|
logging.info(f" * DLMA multiply.txt in {current_dir} set to {val}")
|
|
except Exception as e:
|
|
logging.warning(f" * {Fore.LIGHTYELLOW_EX}Error trying to read multiply.txt for {current_dir}: {Style.RESET_ALL}{e}")
|
|
skip_folders.append(current_dir)
|
|
multipliers[current_dir] = 1.0
|
|
else:
|
|
skip_folders.append(current_dir)
|
|
multipliers[current_dir] = 1.0
|
|
|
|
caption = ImageCaption.resolve(pathname)
|
|
item = self.image_train_item(pathname, caption, multiplier=multipliers[current_dir])
|
|
|
|
cur_file_multiplier = multipliers[current_dir]
|
|
|
|
while cur_file_multiplier >= 1.0:
|
|
items.append(item)
|
|
cur_file_multiplier -= 1
|
|
|
|
if cur_file_multiplier > 0:
|
|
if randomizer.random() < cur_file_multiplier:
|
|
items.append(item)
|
|
|
|
if item:
|
|
items.append(item)
|
|
return items
|
|
|
|
@staticmethod
|
|
def unzip_all(path):
|
|
try:
|
|
for root, dirs, files in os.walk(path):
|
|
for file in files:
|
|
if file.endswith('.zip'):
|
|
logging.info(f"Unzipping {file}")
|
|
with zipfile.ZipFile(path, 'r') as zip_ref:
|
|
zip_ref.extractall(path)
|
|
except Exception as e:
|
|
logging.error(f"Error unzipping files {e}")
|
|
|
|
@staticmethod
|
|
def recurse_data_root(recurse_root):
|
|
for f in os.listdir(recurse_root):
|
|
current = os.path.join(recurse_root, f)
|
|
|
|
if os.path.isfile(current):
|
|
ext = os.path.splitext(f)[1].lower()
|
|
if ext in ['.jpg', '.jpeg', '.png', '.bmp', '.webp', '.jfif']:
|
|
yield current
|
|
|
|
sub_dirs = []
|
|
|
|
for d in os.listdir(recurse_root):
|
|
current = os.path.join(recurse_root, d)
|
|
if os.path.isdir(current):
|
|
sub_dirs.append(current)
|
|
|
|
for dir in sub_dirs:
|
|
DirectoryResolver.__recurse_data_root(dir)
|
|
|
|
def strategy(data_root: str):
|
|
if os.path.isfile(data_root) and data_root.endswith('.json'):
|
|
return JSONResolver
|
|
|
|
if os.path.isdir(data_root):
|
|
return DirectoryResolver
|
|
|
|
raise ValueError(f"data_root '{data_root}' is not a valid directory or JSON file.")
|
|
|
|
|
|
def resolve_root(path: str, aspects: list[float], flip_p: float = 0.0, seed=555) -> list[ImageTrainItem]:
|
|
"""
|
|
:param data_root: Directory or JSON file.
|
|
:param aspects: The list of aspect ratios to use
|
|
:param flip_p: The probability of flipping the image
|
|
"""
|
|
if os.path.isfile(path) and path.endswith('.json'):
|
|
resolver = JSONResolver(aspects, flip_p, seed)
|
|
|
|
if os.path.isdir(path):
|
|
resolver = DirectoryResolver(aspects, flip_p, seed)
|
|
|
|
if not resolver:
|
|
raise ValueError(f"data_root '{path}' is not a valid directory or JSON file.")
|
|
|
|
items = resolver.image_train_items(path)
|
|
return items
|
|
|
|
def resolve(value: typing.Union[dict, str], aspects: list[float], flip_p: float=0.0, seed=555) -> list[ImageTrainItem]:
|
|
"""
|
|
Resolve the training data from the value.
|
|
:param value: The value to resolve, either a dict or a string.
|
|
:param aspects: The list of aspect ratios to use
|
|
:param flip_p: The probability of flipping the image
|
|
"""
|
|
if isinstance(value, str):
|
|
return resolve_root(value, aspects, flip_p)
|
|
|
|
if isinstance(value, dict):
|
|
resolver = value.get('resolver', None)
|
|
match resolver:
|
|
case 'directory' | 'json':
|
|
path = value.get('path', None)
|
|
return resolve_root(path, aspects, flip_p, seed)
|
|
case 'multi':
|
|
items = []
|
|
for resolver in value.get('resolvers', []):
|
|
items.extend(resolve(resolver, aspects, flip_p, seed))
|
|
return items
|
|
case _:
|
|
raise ValueError(f"Cannot resolve training data for resolver value '{resolver}'") |