ostris/ai-toolkit
1
0
Fork 0
mirror of https://github.com/ostris/ai-toolkit.git synced 2026-03-02 17:19:48 +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

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)