From 261b6790ee1e92d6d84220ef2d990542acc3b8aa Mon Sep 17 00:00:00 2001 From: bmaltais Date: Sun, 12 Feb 2023 07:02:05 -0500 Subject: [PATCH] Update tool --- README.md | 4 +- tools/crop_images_to_n_buckets.py | 125 +++++++++++++++++++++++++----- 2 files changed, 108 insertions(+), 21 deletions(-) diff --git a/README.md b/README.md index 0b0d45e..d5cb7b8 100644 --- a/README.md +++ b/README.md @@ -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 diff --git a/tools/crop_images_to_n_buckets.py b/tools/crop_images_to_n_buckets.py index 3b0d68a..dff7825 100644 --- a/tools/crop_images_to_n_buckets.py +++ b/tools/crop_images_to_n_buckets.py @@ -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() \ No newline at end of file