Creates a grid of images with 'nrows' number of image rows, images must all be the same shape.

simple_image_grid.py

from typing import List
import numpy as np
def image_grid(images: List[np.ndarray], nrows: int):
    """
    Creates a grid of images with 'nrows' number of image rows, images must all be the same shape.
    """
    assert len(images) > nrows, "Must have more images than rows"
    ncols = int(np.ceil(len(images) / nrows))
    images_size = images[0].shape[:2]
    grid_img = np.zeros((nrows * images_size[0], ncols * images_size[1], 3), dtype=np.uint8)
    for i, img in enumerate(images):
        r = int(i / ncols)
        c = i % ncols
        grid_img[r * images_size[0]:(r + 1) * images_size[0], c * images_size[1]:(c + 1) * images_size[1]] = img
    return grid_img