ostris/ai-toolkit
1
0
Fork 0
mirror of https://github.com/ostris/ai-toolkit.git synced 2026-02-06 21:49:57 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fixed issue with sizes for some images being loaded sideways resulting in squished images.

Browse Source
This commit is contained in:
Jaret Burkett
2024-10-20 11:51:29 -06:00
parent e3ebd73610
commit bedb8197a2
2 changed files with 26 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
toolkit/data_loader.py
View File

@@ -441,16 +441,24 @@ class AiToolkitDataset(LatentCachingMixin, CLIPCachingMixin, BucketsMixin, Capti
if not os.path.isdir(self.dataset_path):
dataset_folder = os.path.dirname(dataset_folder)
dataset_size_file = os.path.join(dataset_folder, '.aitk_size.json')
dataloader_version = "0.1.1"
if os.path.exists(dataset_size_file):
try:
with open(dataset_size_file, 'r') as f:
self.size_database = json.load(f)
if "__version__" not in self.size_database or self.size_database["__version__"] != dataloader_version:
print("Upgrading size database to new version")
# old version, delete and recreate
self.size_database = {}
except Exception as e:
print(f"Error loading size database: {dataset_size_file}")
print(e)
self.size_database = {}
else:
self.size_database = {}
self.size_database["__version__"] = dataloader_version
bad_count = 0
for file in tqdm(file_list):
@@ -461,6 +469,7 @@ class AiToolkitDataset(LatentCachingMixin, CLIPCachingMixin, BucketsMixin, Capti
dataset_config=dataset_config,
dataloader_transforms=self.transform,
size_database=self.size_database,
dataset_root=dataset_folder,
)
self.file_list.append(file_item)
except Exception as e:

28
toolkit/data_transfer_object/data_loader.py
View File

@@ -44,19 +44,25 @@ class FileItemDTO(
self.path = kwargs.get('path', '')
self.dataset_config: 'DatasetConfig' = kwargs.get('dataset_config', None)
size_database = kwargs.get('size_database', {})
filename = os.path.basename(self.path)
if filename in size_database:
w, h = size_database[filename]
dataset_root = kwargs.get('dataset_root', None)
if dataset_root is not None:
# remove dataset root from path
file_key = self.path.replace(dataset_root, '')
else:
file_key = os.path.basename(self.path)
if file_key in size_database:
w, h = size_database[file_key]
else:
# original method is significantly faster, but some images are read sideways. Not sure why. Do slow method for now.
# process width and height
try:
w, h = image_utils.get_image_size(self.path)
except image_utils.UnknownImageFormat:
print_once(f'Warning: Some images in the dataset cannot be fast read. ' + \
f'This process is faster for png, jpeg')
img = exif_transpose(Image.open(self.path))
w, h = img.size
size_database[filename] = (w, h)
# try:
# w, h = image_utils.get_image_size(self.path)
# except image_utils.UnknownImageFormat:
# print_once(f'Warning: Some images in the dataset cannot be fast read. ' + \
# f'This process is faster for png, jpeg')
img = exif_transpose(Image.open(self.path))
w, h = img.size
size_database[file_key] = (w, h)
self.width: int = w
self.height: int = h
self.dataloader_transforms = kwargs.get('dataloader_transforms', None)