ostris/ai-toolkit
1
0
Fork 0
mirror of https://github.com/ostris/ai-toolkit.git synced 2026-03-09 12:39:49 +00:00
Code Issues Packages Projects Releases Wiki Activity

Style and content loss working

Browse Source
This commit is contained in:
Jaret Burkett
2023-07-18 07:47:01 -06:00
parent 439310e4dc
commit 94d52572d4
3 changed files with 281 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
toolkit/data_loader.py
View File

@@ -51,15 +51,20 @@ class ImageDataset(Dataset):
# Downscale the source image first
img = img.resize((int(img.size[0] * self.scale), int(img.size[1] * self.scale)), Image.BICUBIC)
min_img_size = min(img.size)
if self.random_crop:
if self.random_scale:
scale_size = random.randint(int(img.size[0] * self.scale), self.resolution)
if self.random_scale and min_img_size > self.resolution:
if min_img_size < self.resolution:
print(
f"Unexpected values: min_img_size={min_img_size}, self.resolution={self.resolution}, image file={file}")
scale_size = self.resolution
else:
scale_size = random.randint(self.resolution, int(min_img_size))
img = img.resize((scale_size, scale_size), Image.BICUBIC)
img = transforms.RandomCrop(self.resolution)(img)
else:
min_dimension = min(img.size)
img = transforms.CenterCrop(min_dimension)(img)
img = transforms.CenterCrop(min_img_size)(img)
img = img.resize((self.resolution, self.resolution), Image.BICUBIC)
img = self.transform(img)