github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 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,10 +441,16 @@ class AiToolkitDataset(LatentCachingMixin, CLIPCachingMixin, BucketsMixin, Capti
if not os.path.isdir(self.dataset_path): if not os.path.isdir(self.dataset_path):
dataset_folder = os.path.dirname(dataset_folder) dataset_folder = os.path.dirname(dataset_folder)
dataset_size_file = os.path.join(dataset_folder, '.aitk_size.json') dataset_size_file = os.path.join(dataset_folder, '.aitk_size.json')
dataloader_version = "0.1.1"
if os.path.exists(dataset_size_file): if os.path.exists(dataset_size_file):
try: try:
with open(dataset_size_file, 'r') as f: with open(dataset_size_file, 'r') as f:
self.size_database = json.load(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: except Exception as e:
print(f"Error loading size database: {dataset_size_file}") print(f"Error loading size database: {dataset_size_file}")
print(e) print(e)
@@ -452,6 +458,8 @@ class AiToolkitDataset(LatentCachingMixin, CLIPCachingMixin, BucketsMixin, Capti
else: else:
self.size_database = {} self.size_database = {}
self.size_database["__version__"] = dataloader_version
bad_count = 0 bad_count = 0
for file in tqdm(file_list): for file in tqdm(file_list):
try: try:
@@ -461,6 +469,7 @@ class AiToolkitDataset(LatentCachingMixin, CLIPCachingMixin, BucketsMixin, Capti
dataset_config=dataset_config, dataset_config=dataset_config,
dataloader_transforms=self.transform, dataloader_transforms=self.transform,
size_database=self.size_database, size_database=self.size_database,
dataset_root=dataset_folder,
) )
self.file_list.append(file_item) self.file_list.append(file_item)
except Exception as e: 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.path = kwargs.get('path', '')
self.dataset_config: 'DatasetConfig' = kwargs.get('dataset_config', None) self.dataset_config: 'DatasetConfig' = kwargs.get('dataset_config', None)
size_database = kwargs.get('size_database', {}) size_database = kwargs.get('size_database', {})
filename = os.path.basename(self.path) dataset_root = kwargs.get('dataset_root', None)
if filename in size_database: if dataset_root is not None:
w, h = size_database[filename] # remove dataset root from path
file_key = self.path.replace(dataset_root, '')
else: 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 # process width and height
try: # try:
w, h = image_utils.get_image_size(self.path) # w, h = image_utils.get_image_size(self.path)
except image_utils.UnknownImageFormat: # except image_utils.UnknownImageFormat:
print_once(f'Warning: Some images in the dataset cannot be fast read. ' + \ # print_once(f'Warning: Some images in the dataset cannot be fast read. ' + \
f'This process is faster for png, jpeg') # f'This process is faster for png, jpeg')
img = exif_transpose(Image.open(self.path)) img = exif_transpose(Image.open(self.path))
w, h = img.size w, h = img.size
size_database[filename] = (w, h) size_database[file_key] = (w, h)
self.width: int = w self.width: int = w
self.height: int = h self.height: int = h
self.dataloader_transforms = kwargs.get('dataloader_transforms', None) self.dataloader_transforms = kwargs.get('dataloader_transforms', None)