Group images by aspect ratio and list images which belong to an aspect ratio with 12 or fewer images. Useful for Kohya SS

ratios.py

from PIL import Image
import os


def get_aspect_ratio(image_path):
    with Image.open(image_path) as img:
        width, height = img.size
        aspect_ratio = round(width / height, 1)
    return aspect_ratio


def get_all_aspect_ratios(directory):
    image_files = [f for f in os.listdir(directory) if
                   f.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp', '.tiff'))]

    aspect_ratios = {}

    for image_file in image_files:
        image_path = os.path.join(directory, image_file)
        aspect_ratio = get_aspect_ratio(image_path)
        aspect_ratios[image_file] = aspect_ratio

    return aspect_ratios


if __name__ == "__main__":
    directory_path = "./"
    bucket_count_threshold = 12

    aspect_ratios = get_all_aspect_ratios(directory_path)
    dir(aspect_ratios)

    counts = {}

    small_ratios = []

    for image_file, aspect_ratio in aspect_ratios.items():
        ratio_format = f"{aspect_ratio:.1f}"
        if ratio_format not in counts:
            counts[ratio_format] = 0
        counts[ratio_format] += 1
    for ratio, count in counts.items():
        # print(f"{ratio}: {count}")
        if count < bucket_count_threshold:
            # print("ratio has too few images")
            if ratio not in small_ratios:
                small_ratios.append(ratio)

    small_files = []

    for ratio in small_ratios:
        for image_file, aspect_ratio in aspect_ratios.items():
            ratio_format = f"{aspect_ratio:.1f}"
            if ratio_format == ratio:
                small_files.append(f"{image_file}: {ratio_format}")
    print(f"The following files are in buckets {bucket_count_threshold} or smaller:")
    for small_file in small_files:
        print(small_file)