ostris/ai-toolkit
1
0
Fork 0
mirror of https://github.com/ostris/ai-toolkit.git synced 2026-05-01 03:31:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

Bugfixes and cleanup

Browse Source
This commit is contained in:
Jaret Burkett
2024-08-01 11:45:12 -06:00
parent 47744373f2
commit 03613c523f
4 changed files with 39 additions and 127 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
toolkit/dataloader_mixins.py
View File

@@ -760,11 +760,15 @@ class ClipImageFileItemDTOMixin:
# do a flip
img = img.transpose(Image.FLIP_TOP_BOTTOM)
# image must be square. If it is not, we will resize/squish it so it is, that way we don't crop out data
if img.width != img.height:
# resize to the smallest dimension
min_size = min(img.width, img.height)
img = img.resize((min_size, min_size), Image.BICUBIC)
if self.dataset_config.square_crop:
# center crop to a square
img = transforms.CenterCrop(min_size)(img)
else:
# image must be square. If it is not, we will resize/squish it so it is, that way we don't crop out data
# resize to the smallest dimension
img = img.resize((min_size, min_size), Image.BICUBIC)
if self.has_clip_augmentations:
self.clip_image_tensor = self.augment_clip_image(img, transform=None)