Make a grid of images that are stored individually in a numpy array

make_grid.py

def make_grid(arr: np.ndarray, nrows: int, pad_size: int = 2, pad_value: int = 0):
    """
    Expected array layout is (N, C, H, W)
    """
    arr = np.pad(arr,
                 ((0,0),
                  (0,0),
                  (pad_size // 2, pad_size // 2),
                  (pad_size // 2, pad_size // 2)),
                 constant_values=(pad_value,))

    N, C, H, W = arr.shape
    if N % nrows != 0:
        additional_arrays = nrows - N % nrows
        arr = np.concatenate([arr, (pad_value * np.ones((additional_arrays, C, H, W))).astype(np.uint8)], axis=0)

    ncols = arr.shape[0] // nrows
    return (arr
        .transpose(0, 2, 3, 1)
        .reshape(nrows, ncols, H, W, C)
        .transpose(0, 2, 1, 3, 4)
        .reshape(H * nrows, W * ncols, C))