ostris/ai-toolkit
1
0
Fork 0
mirror of https://github.com/ostris/ai-toolkit.git synced 2026-04-30 19:21:39 +00:00
Code Issues Packages Projects Releases Wiki Activity

Reworked bucket loader to scale buckets to pixels amounts not just minimum size. Makes the network more consistant

Browse Source
This commit is contained in:
Jaret Burkett
2023-08-30 14:52:12 -06:00
parent d401348c2e
commit 33267e117c
6 changed files with 137 additions and 69 deletions
Download Patch File Download Diff File Expand all files Collapse all files

85
toolkit/dataloader_mixins.py
View File

@@ -1,3 +1,4 @@
import math
import os
import random
from typing import TYPE_CHECKING, List, Dict, Union
@@ -94,56 +95,56 @@ class BucketsMixin:
bucket_tolerance = config.bucket_tolerance
file_list: List['FileItemDTO'] = self.file_list
# make sure out resolution is divisible by bucket_tolerance
if resolution % bucket_tolerance != 0:
# reduce it to the nearest divisible number
resolution = resolution - (resolution % bucket_tolerance)
total_pixels = resolution * resolution
# for file_item in enumerate(file_list):
for idx, file_item in enumerate(file_list):
width = file_item.crop_width
height = file_item.crop_height
# determine new size, smallest dimension should be equal to resolution
# the other dimension should be the same ratio it is now (bigger)
new_width = resolution
new_height = resolution
if width > height:
# scale width to match new resolution,
new_width = int(width * (resolution / height))
file_item.crop_width = new_width
file_item.scale_to_width = new_width
file_item.crop_height = resolution
file_item.scale_to_height = resolution
# make sure new_width is divisible by bucket_tolerance
# determine new resolution to have the same number of pixels
current_pixels = width * height
if current_pixels == total_pixels:
# no change
continue
aspect_ratio = width / height
new_height = int(math.sqrt(total_pixels / aspect_ratio))
new_width = int(aspect_ratio * new_height)
# increase smallest one to be divisible by bucket_tolerance and increase the other to match
if new_width < new_height:
# increase width
if new_width % bucket_tolerance != 0:
# reduce it to the nearest divisible number
reduction = new_width % bucket_tolerance
file_item.crop_width = new_width - reduction
new_width = file_item.crop_width
# adjust the new x position so we evenly crop
file_item.crop_x = int(file_item.crop_x + (reduction / 2))
elif height > width:
# scale height to match new resolution
new_height = int(height * (resolution / width))
file_item.crop_height = new_height
file_item.scale_to_height = new_height
file_item.scale_to_width = resolution
file_item.crop_width = resolution
# make sure new_height is divisible by bucket_tolerance
if new_height % bucket_tolerance != 0:
# reduce it to the nearest divisible number
reduction = new_height % bucket_tolerance
file_item.crop_height = new_height - reduction
new_height = file_item.crop_height
# adjust the new x position so we evenly crop
file_item.crop_y = int(file_item.crop_y + (reduction / 2))
crop_amount = new_width % bucket_tolerance
new_width = new_width + (bucket_tolerance - crop_amount)
new_height = int(new_width / aspect_ratio)
else:
# square image
file_item.crop_height = resolution
file_item.scale_to_height = resolution
file_item.scale_to_width = resolution
file_item.crop_width = resolution
# increase height
if new_height % bucket_tolerance != 0:
crop_amount = new_height % bucket_tolerance
new_height = new_height + (bucket_tolerance - crop_amount)
new_width = int(aspect_ratio * new_height)
# Ensure that the total number of pixels remains the same.
# assert new_width * new_height == total_pixels
file_item.scale_to_width = new_width
file_item.scale_to_height = new_height
file_item.crop_width = new_width
file_item.crop_height = new_height
# make sure it is divisible by bucket_tolerance, decrease if not
if new_width % bucket_tolerance != 0:
crop_amount = new_width % bucket_tolerance
file_item.crop_width = new_width - crop_amount
else:
file_item.crop_width = new_width
if new_height % bucket_tolerance != 0:
crop_amount = new_height % bucket_tolerance
file_item.crop_height = new_height - crop_amount
else:
file_item.crop_height = new_height
# check if bucket exists, if not, create it
bucket_key = f'{new_width}x{new_height}'