collage.py

import argparse
import itertools
import logging
from pathlib import Path
from PIL import Image
import sys


def collage_size(*, image_count, x_count, source_image_size):
    if image_count <= x_count:
        collage_width = source_image_size[0] * image_count
    else:
        collage_width = source_image_size[0] * x_count

    collage_height = source_image_size[1] * (1 + (image_count - 1) // x_count)

    return collage_width, collage_height


def create_collage_from_dir(source_dir: Path, x_count=10):
    supported_extensions = {".png", ".jpg", ".bmp", ".jpeg"}
    image_paths = [path for path in sorted(source_dir.iterdir()) if path.is_file() and path.suffix in supported_extensions]

    images = [Image.open(path) for path in image_paths]
    if not images:
        logging.warning(f"Expected images in {source_dir}, but found nothing")
        return

    # Check that all images are the same size
    images.sort(key=lambda img: img.size)
    image_sizes = {size: list(imgs) for size, imgs in itertools.groupby(images, lambda img: img.size)}
    if len(image_sizes) > 1:
        logging.error(f"Found multiple image sizes: {list(image_sizes)}")
        for size, imgs in image_sizes.items():
            logging.error(f"{size}: {[img.filename for img in imgs]}")
        sys.exit(1)

    source_image_size = images[0].size
    size = collage_size(
        image_count=len(images),
        x_count=x_count,
        source_image_size=source_image_size,
    )

    collage = Image.new("RGBA", size)
    for idx, image in enumerate(images):
        position = (
            source_image_size[0] * (idx % x_count),
            source_image_size[1] * (idx // x_count)
        )
        collage.paste(image, position)

    image_path = source_dir.parent.joinpath(source_dir.name).with_suffix(".png")
    collage.save(image_path)
    logging.info(f"Created image at {image_path}")


def main():
    parser = argparse.ArgumentParser(
        description="Run without arguments to create grids of images from sibling directories"
    )
    parser.parse_args()

    source_directories = [path for path in Path.cwd().iterdir() if path.is_dir()]
    if not source_directories:
        logging.error("No sibling directories found")
        sys.exit(1)

    for d in source_directories:
        logging.info(f"Creating collage for {d}")
        create_collage_from_dir(d)


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    main()