Create an image grid in Python

image_grid.py

import argparse
import math
from glob import glob

from PIL import Image


class MissingArgument(Exception):
    pass


def calculate_grid_dimensions(n):
    sqrt_n = math.sqrt(n)
    rows = int(sqrt_n)
    columns = math.ceil(n / rows)
    return rows, columns


def create_image_grid(folder_path, wh, r, c, output_file):
    images = []
    files = sorted(glob(f'{folder_path}/*'))
    wh = wh.split(',')
    image_width = int(wh[0])
    image_height = int(wh[1])
    if not r and not c:
        r, c = calculate_grid_dimensions(len(files))
    elif not r or not c:
        raise MissingArgument('If you pass --rows or --columns, '
                              'you need to pass a value for both.')

    output_width = c * image_width
    output_height = r * image_height
    output_image = Image.new('RGB', (output_width, output_height))

    for file in files:
        try:
            img = Image.open(file)
            img = img.resize((image_width, image_height), Image.ANTIALIAS)
            images.append(img)
        except Exception as e:
            print(f'Error processing {file}: {e}')

    for i in range(r):
        for j in range(c):
            index = i * c + j
            if index < len(images):
                output_image.paste(images[index],
                                   (j * image_width, i * image_height))

    output_image.save(output_file)
    output_image.show()


def opts() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument('-p',
                        '--folder-path',
                        help='Path to the folder.',
                        required=True)
    parser.add_argument(
        '-wh',
        '--width-height',
        help='Width and height of each square image, separated by a comma.',
        default='300,300')
    parser.add_argument('-r', '--rows', help='Number of rows.', type=int)
    parser.add_argument('-c', '--columns', help='Number of columns.', type=int)
    parser.add_argument('-o',
                        '--output-file',
                        help='Output file name or path.',
                        default='image_grid.png')
    return parser.parse_args()


def main():
    args = opts()
    create_image_grid(args.folder_path, args.width_height, args.rows,
                      args.columns, args.output_file)


if __name__ == "__main__":
    main()