ostris/ai-toolkit
1
0
Fork 0
mirror of https://github.com/ostris/ai-toolkit.git synced 2026-02-27 15:53:58 +00:00
Code Issues Packages Projects Releases Wiki Activity

Added caching to image sizes so we dont do it every time.

Browse Source
This commit is contained in:
Jaret Burkett
2024-07-15 19:07:41 -06:00
parent e4558dff4b
commit 58dffd43a8
7 changed files with 90 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
toolkit/data_transfer_object/data_loader.py
View File

@@ -1,3 +1,4 @@
import os
import weakref
from _weakref import ReferenceType
from typing import TYPE_CHECKING, List, Union
@@ -40,16 +41,22 @@ class FileItemDTO(
ArgBreakMixin,
):
def __init__(self, *args, **kwargs):
self.path = kwargs.get('path', None)
self.path = kwargs.get('path', '')
self.dataset_config: 'DatasetConfig' = kwargs.get('dataset_config', None)
# 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))
h, w = img.size
size_database = kwargs.get('size_database', {})
filename = os.path.basename(self.path)
if filename in size_database:
w, h = size_database[filename]
else:
# 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))
h, w = img.size
size_database[filename] = (w, h)
self.width: int = w
self.height: int = h
self.dataloader_transforms = kwargs.get('dataloader_transforms', None)