StableDiffusion/KohyaSS
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Update tool

Browse Source
This commit is contained in:
bmaltais 2023-02-12 07:02:05 -05:00
parent a49fb9cb8c
commit 261b6790ee
2 changed files with 108 additions and 21 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
README.md
View File

@ -144,10 +144,10 @@ Then redo the installation instruction within the kohya_ss venv.
## Change history
* 2023/02/11 (v20.7.2):
- ``lora_interrogator.py`` is added in ``networks`` folder. See ``python networks\lora_interrogator.py -h`` for usage.
- `lora_interrogator.py` is added in `networks` folder. See `python networks\lora_interrogator.py -h` for usage.
- For LoRAs where the activation word is unknown, this script compares the output of Text Encoder after applying LoRA to that of unapplied to find out which token is affected by LoRA. Hopefully you can figure out the activation word. LoRA trained with captions does not seem to be able to interrogate.
- Batch size can be large (like 64 or 128).
- ``train_textual_inversion.py`` now supports multiple init words.
- `train_textual_inversion.py` now supports multiple init words.
- Following feature is reverted to be the same as before. Sorry for confusion:
> Now the number of data in each batch is limited to the number of actual images (not duplicated). Because a certain bucket may contain smaller number of actual images, so the batch may contain same (duplicated) images.
- Add new tool to sort, group and average crop image in a dataset

125
tools/crop_images_to_n_buckets.py
View File

@ -7,6 +7,7 @@
import os
import cv2
import argparse
import shutil
def aspect_ratio(img_path):
"""Return aspect ratio of an image"""
@ -38,9 +39,22 @@ def average_aspect_ratio(group):
"""Calculate average aspect ratio for a group"""
aspect_ratios = [aspect_ratio for _, aspect_ratio in group]
avg_aspect_ratio = sum(aspect_ratios) / len(aspect_ratios)
print(f"Average aspect ratio for group: {avg_aspect_ratio}")
return avg_aspect_ratio
def center_crop_image(image, target_aspect_ratio):
"""Crop the input image to the target aspect ratio.
The function calculates the crop region for the input image based on its current aspect ratio and the target aspect ratio.
Args:
image: A numpy array representing the input image.
target_aspect_ratio: A float representing the target aspect ratio.
Returns:
A numpy array representing the cropped image.
"""
height, width = image.shape[:2]
current_aspect_ratio = float(width) / float(height)
@ -58,45 +72,111 @@ def center_crop_image(image, target_aspect_ratio):
return cropped_image
def save_cropped_images(group, folder_name, group_number, avg_aspect_ratio):
def copy_related_files(img_path, save_path):
"""
Copy all files in the same directory as the input image that have the same base name as the input image to the
output directory with the corresponding new filename.
:param img_path: Path to the input image.
:param save_path: Path to the output image.
"""
# Get the base filename and directory
img_dir, img_basename = os.path.split(img_path)
img_base, img_ext = os.path.splitext(img_basename)
save_dir, save_basename = os.path.split(save_path)
save_base, save_ext = os.path.splitext(save_basename)
# Create the output directory if it does not exist
if not os.path.exists(save_dir):
os.makedirs(save_dir)
# Loop over all files in the same directory as the input image
try:
for filename in os.listdir(img_dir):
# Skip files with the same name as the input image
if filename == img_basename:
continue
# Check if the file has the same base name as the input image
file_base, file_ext = os.path.splitext(filename)
if file_base == img_base:
# Build the new filename and copy the file
new_filename = os.path.join(save_dir, f"{save_base}{file_ext}")
shutil.copy2(os.path.join(img_dir, filename), new_filename)
except OSError as e:
print(f"Error: {e}") # Handle errors from os.listdir()
def save_resized_cropped_images(group, folder_name, group_number, avg_aspect_ratio, use_original_name=False):
"""Crop and resize all images in the input group to the smallest resolution, and save them to a folder.
Args:
group: A list of tuples, where each tuple contains the path to an image and its aspect ratio.
folder_name: A string representing the name of the folder to save the images to.
group_number: An integer representing the group number.
avg_aspect_ratio: A float representing the average aspect ratio of the images in the group.
use_original_name: A boolean indicating whether to save the images with their original file names.
"""
if not os.path.exists(folder_name):
os.makedirs(folder_name)
# get the smallest size of the images
small_height = 0
small_width = 0
smallest_res = 100000000
for i, image in enumerate(group):
img_path, aspect_ratio = image
smallest_res = float("inf")
for img_path, _ in group:
image = cv2.imread(img_path)
cropped_image = center_crop_image(image, avg_aspect_ratio)
height, width = cropped_image.shape[:2]
if smallest_res > height * width:
small_height = height
small_width = width
smallest_res = height * width
image_res = height * width
if image_res < smallest_res:
smallest_res = image_res
small_height, small_width = height, width
# resize all images to the smallest resolution of the images in the group
for i, image in enumerate(group):
img_path, aspect_ratio = image
for i, (img_path, aspect_ratio) in enumerate(group):
image = cv2.imread(img_path)
cropped_image = center_crop_image(image, avg_aspect_ratio)
resized_image = cv2.resize(cropped_image, (small_width, small_height))
save_path = os.path.join(folder_name, "group_{}_{}.jpg".format(group_number, i))
if use_original_name:
save_name = os.path.basename(img_path)
else:
save_name = f"group_{group_number}_{i}.jpg"
save_path = os.path.join(folder_name, save_name)
cv2.imwrite(save_path, resized_image)
# Copy matching files named the same as img_path to
copy_related_files(img_path, save_path)
print(f"Saved {save_name} to {folder_name}")
def main():
parser = argparse.ArgumentParser(description='Sort images and crop them based on aspect ratio')
parser.add_argument('--path', type=str, help='Path to the directory containing images', required=True)
parser.add_argument('--dst_path', type=str, help='Path to the directory to save the cropped images', required=True)
parser.add_argument('--batch_size', type=int, help='Size of the batches to create', required=True)
parser.add_argument('input_dir', type=str, help='Path to the directory containing images')
parser.add_argument('output_dir', type=str, help='Path to the directory to save the cropped images')
parser.add_argument('batch_size', type=int, help='Size of the batches to create')
parser.add_argument('--use_original_name', action='store_true', help='Whether to use original file names for the saved images')
args = parser.parse_args()
sorted_images = sort_images_by_aspect_ratio(args.path)
print(f"Sorting images by aspect ratio in {args.input_dir}...")
if not os.path.exists(args.input_dir):
print(f"Error: Input directory does not exist: {args.input_dir}")
return
if not os.path.exists(args.output_dir):
try:
os.makedirs(args.output_dir)
except OSError:
print(f"Error: Failed to create output directory: {args.output_dir}")
return
sorted_images = sort_images_by_aspect_ratio(args.input_dir)
total_images = len(sorted_images)
print(f'Total images: {total_images}')
if args.batch_size <= 0:
print("Error: Batch size must be greater than 0")
return
group_size = total_images // args.batch_size
@ -111,11 +191,18 @@ def main():
print('Creating groups...')
groups = create_groups(sorted_images, group_size)
print(f"Created {len(groups)} groups")
print('Saving cropped and resize images...')
for i, group in enumerate(groups):
avg_aspect_ratio = average_aspect_ratio(group)
save_cropped_images(group, args.dst_path, i+1, avg_aspect_ratio)
print(f"Processing group {i+1} with {len(group)} images...")
try:
save_resized_cropped_images(group, args.output_dir, i+1, avg_aspect_ratio, args.use_original_name)
except Exception as e:
print(f"Error: Failed to save images in group {i+1}: {e}")
print('Done')
if __name__ == '__main__':
main()