mirror of https://github.com/bmaltais/kohya_ss
Add Group Images tool and GUI
parent
61f52f0eb5
commit
6bb8ec7641
|
|
@ -345,6 +345,8 @@ This will store a backup file with your current locally installed pip packages a
|
|||
|
||||
## Change History
|
||||
|
||||
* 2023/05/28 (v21.5.14)
|
||||
- Add Create Groupo tool and GUI
|
||||
* 2023/05/24 (v21.5.13)
|
||||
- Upgrade gradio release to fix issue with UI refresh on config load.
|
||||
- [D-Adaptation v3.0](https://github.com/facebookresearch/dadaptation) is now supported. [PR #530](https://github.com/kohya-ss/sd-scripts/pull/530) Thanks to sdbds!
|
||||
|
|
|
|||
|
|
@ -95,7 +95,7 @@ def gradio_group_images_gui_tab(headless=False):
|
|||
info='Do not copy other files in the input folder to the output folder',
|
||||
)
|
||||
|
||||
group_images_button = gr.Button('Caption images')
|
||||
group_images_button = gr.Button('Group images')
|
||||
|
||||
group_images_button.click(
|
||||
group_images,
|
||||
|
|
|
|||
|
|
@ -1,18 +1,18 @@
|
|||
import argparse
|
||||
import shutil
|
||||
from PIL import Image
|
||||
from PIL import Image, ImageOps
|
||||
import os
|
||||
import numpy as np
|
||||
# import math
|
||||
|
||||
class ImageProcessor:
|
||||
|
||||
def __init__(self, input_folder, output_folder, group_size, include_subfolders, do_not_copy_other_files):
|
||||
def __init__(self, input_folder, output_folder, group_size, include_subfolders, do_not_copy_other_files, pad):
|
||||
self.input_folder = input_folder
|
||||
self.output_folder = output_folder
|
||||
self.group_size = group_size
|
||||
self.include_subfolders = include_subfolders
|
||||
self.do_not_copy_other_files = do_not_copy_other_files
|
||||
self.pad = pad
|
||||
self.image_extensions = ('.png', '.jpg', '.jpeg', '.gif', '.webp')
|
||||
|
||||
def get_image_paths(self):
|
||||
|
|
@ -100,6 +100,42 @@ class ImageProcessor:
|
|||
for i, group in enumerate(groups):
|
||||
print(f"Processing group {i+1} with {len(group)} images...")
|
||||
self.process_group(group, i)
|
||||
|
||||
def process_group(self, group, group_index):
|
||||
if len(group) > 0:
|
||||
aspect_ratios = self.get_aspect_ratios(group)
|
||||
avg_aspect_ratio = np.mean(aspect_ratios)
|
||||
if self.pad:
|
||||
padded_images = self.pad_images(group, avg_aspect_ratio)
|
||||
self.resize_and_save_images(padded_images, group_index)
|
||||
else:
|
||||
cropped_images = self.crop_images(group, avg_aspect_ratio)
|
||||
self.resize_and_save_images(cropped_images, group_index)
|
||||
if not self.do_not_copy_other_files:
|
||||
self.copy_other_files(group, group_index)
|
||||
|
||||
def pad_images(self, group, avg_aspect_ratio):
|
||||
padded_images = []
|
||||
for j, path in enumerate(group):
|
||||
with Image.open(path) as img:
|
||||
print(f" Processing image {j+1}: {path}")
|
||||
img = self.pad_image(img, avg_aspect_ratio)
|
||||
padded_images.append(img)
|
||||
return padded_images
|
||||
|
||||
def pad_image(self, img, avg_aspect_ratio):
|
||||
img_aspect_ratio = img.width / img.height
|
||||
if img_aspect_ratio < avg_aspect_ratio:
|
||||
# Too tall, increase width
|
||||
new_width = avg_aspect_ratio * img.height
|
||||
pad_width = int((new_width - img.width) / 2)
|
||||
img = ImageOps.expand(img, border=(pad_width, 0), fill='black')
|
||||
else:
|
||||
# Too wide, increase height
|
||||
new_height = img.width / avg_aspect_ratio
|
||||
pad_height = int((new_height - img.height) / 2)
|
||||
img = ImageOps.expand(img, border=(0, pad_height), fill='black')
|
||||
return img
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description='Process groups of images.')
|
||||
|
|
@ -108,10 +144,11 @@ def main():
|
|||
parser.add_argument('group_size', type=int, help='Number of images in each group')
|
||||
parser.add_argument('--include_subfolders', action='store_true', help='Include subfolders in search for images')
|
||||
parser.add_argument('--do_not_copy_other_files', '--no_copy', dest='do_not_copy_other_files', action='store_true', help='Do not copy other files with the same name as images')
|
||||
parser.add_argument('--pad', action='store_true', help='Pad images instead of cropping them')
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
processor = ImageProcessor(args.input_folder, args.output_folder, args.group_size, args.include_subfolders, args.do_not_copy_other_files)
|
||||
processor = ImageProcessor(args.input_folder, args.output_folder, args.group_size, args.include_subfolders, args.do_not_copy_other_files, args.pad)
|
||||
processor.process_images()
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
|
|
|||
|
|
@ -0,0 +1,128 @@
|
|||
import argparse
|
||||
from PIL import Image
|
||||
import os
|
||||
import numpy as np
|
||||
import itertools
|
||||
|
||||
class ImageProcessor:
|
||||
|
||||
def __init__(self, input_folder, min_group, max_group, include_subfolders, pad):
|
||||
self.input_folder = input_folder
|
||||
self.min_group = min_group
|
||||
self.max_group = max_group
|
||||
self.include_subfolders = include_subfolders
|
||||
self.pad = pad
|
||||
self.image_extensions = ('.png', '.jpg', '.jpeg', '.gif', '.webp')
|
||||
self.losses = [] # List to store loss values for each image
|
||||
|
||||
def get_image_paths(self):
|
||||
images = []
|
||||
if self.include_subfolders:
|
||||
for dirpath, dirnames, filenames in os.walk(self.input_folder):
|
||||
for filename in filenames:
|
||||
if filename.endswith(self.image_extensions):
|
||||
images.append(os.path.join(dirpath, filename))
|
||||
else:
|
||||
images = [os.path.join(self.input_folder, f) for f in os.listdir(self.input_folder) if f.endswith(self.image_extensions)]
|
||||
return images
|
||||
|
||||
def group_images(self, images, group_size):
|
||||
sorted_images = sorted(images, key=lambda path: Image.open(path).size[0] / Image.open(path).size[1])
|
||||
groups = [sorted_images[i:i+group_size] for i in range(0, len(sorted_images), group_size)]
|
||||
return groups
|
||||
|
||||
def process_group(self, group):
|
||||
if len(group) > 0:
|
||||
aspect_ratios = self.get_aspect_ratios(group)
|
||||
avg_aspect_ratio = np.mean(aspect_ratios)
|
||||
self.calculate_losses(group, avg_aspect_ratio)
|
||||
|
||||
def get_aspect_ratios(self, group):
|
||||
aspect_ratios = []
|
||||
for path in group:
|
||||
with Image.open(path) as img:
|
||||
width, height = img.size
|
||||
aspect_ratios.append(width / height)
|
||||
return aspect_ratios
|
||||
|
||||
def calculate_losses(self, group, avg_aspect_ratio):
|
||||
for j, path in enumerate(group):
|
||||
with Image.open(path) as img:
|
||||
loss = self.calculate_loss(img, avg_aspect_ratio)
|
||||
self.losses.append((path, loss)) # Add (path, loss) tuple to the list
|
||||
|
||||
def calculate_loss(self, img, avg_aspect_ratio):
|
||||
img_aspect_ratio = img.width / img.height
|
||||
if img_aspect_ratio > avg_aspect_ratio:
|
||||
# Too wide, reduce width
|
||||
new_width = avg_aspect_ratio * img.height
|
||||
loss = abs(img.width - new_width) / img.width # Calculate loss value
|
||||
else:
|
||||
# Too tall, reduce height
|
||||
new_height = img.width / avg_aspect_ratio
|
||||
loss = abs(img.height - new_height) / img.height # Calculate loss value
|
||||
return loss
|
||||
|
||||
def monte_carlo_optimization(self, groups):
|
||||
best_groups = groups.copy()
|
||||
best_loss = np.inf
|
||||
best_removed_images = []
|
||||
|
||||
for group in groups:
|
||||
num_images = len(group)
|
||||
all_combinations = []
|
||||
# Generate all possible combinations of images to remove
|
||||
for r in range(1, num_images + 1):
|
||||
combinations = list(itertools.combinations(group, r))
|
||||
all_combinations.extend(combinations)
|
||||
|
||||
for combination in all_combinations:
|
||||
self.losses = [] # Reset losses for each combination
|
||||
remaining_images = list(set(group) - set(combination))
|
||||
self.process_group(remaining_images)
|
||||
avg_loss = np.mean(self.losses)
|
||||
|
||||
if avg_loss < best_loss:
|
||||
best_loss = avg_loss
|
||||
best_groups[best_groups.index(group)] = remaining_images
|
||||
best_removed_images = combination
|
||||
|
||||
return best_groups, best_loss, best_removed_images
|
||||
|
||||
def process_images(self):
|
||||
images = self.get_image_paths()
|
||||
num_images = len(images)
|
||||
results = []
|
||||
|
||||
for group_size in range(self.min_group, self.max_group + 1):
|
||||
groups = self.group_images(images, group_size)
|
||||
optimized_groups, avg_loss, removed_images = self.monte_carlo_optimization(groups)
|
||||
num_remaining = num_images % group_size
|
||||
|
||||
results.append((group_size, avg_loss, num_remaining, optimized_groups, removed_images))
|
||||
|
||||
# Sort results based on average crop loss in ascending order
|
||||
sorted_results = sorted(results, key=lambda x: x[1])
|
||||
|
||||
for group_size, avg_loss, num_remaining, optimized_groups, removed_images in sorted_results:
|
||||
print(f"Group size: {group_size}, Average crop loss: {avg_loss}, Number of images remaining: {num_remaining}")
|
||||
print(f"Optimized Groups: {optimized_groups}")
|
||||
print(f"Removed Images: {removed_images}")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description='Process groups of images.')
|
||||
parser.add_argument('input_folder', type=str, help='Input folder containing images')
|
||||
parser.add_argument('min_group', type=int, help='Minimum group size')
|
||||
parser.add_argument('max_group', type=int, help='Maximum group size')
|
||||
parser.add_argument('--include_subfolders', action='store_true', help='Include subfolders in search for images')
|
||||
parser.add_argument('--pad', action='store_true', help='Pad images instead of cropping them')
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
processor = ImageProcessor(args.input_folder, args.min_group, args.max_group, args.include_subfolders, args.pad)
|
||||
processor.process_images()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Loading…
Reference in New Issue