image concatenator

imgconcat.py

import glob
import numpy as np
from PIL import Image
from einops import rearrange


def img_concatenator(images, row_batch=0):
    # concatenate n images to 1 image
    # rearrange to ceil(n/row_batch) x row_batch
    np_imgs = [np.array(image) for image in images]
    if row_batch == 0:
        np_imgs = rearrange(np_imgs, 'b h w c -> b h w c')
        np_imgs2 = rearrange(np_imgs, 'b2 h w c -> h (b2 w) c')
    else:
        while len(np_imgs) % row_batch != 0:
            np_imgs.append(np.zeros_like(np_imgs[0]))
        np_imgs = rearrange(np_imgs, 'b h w c -> b h w c')
        np_imgs2 = rearrange(np_imgs, '(b1 b2) h w c -> (b1 h) (b2 w) c', b2=row_batch)
    new_im = Image.fromarray(np_imgs2)

    return new_im

if __name__ == "__main__":
    filenames = glob.glob(target_wildcard)
    images = [Image.open(x) for x in filenames]
    new_im = img_concatenator(images, row_batch=4)
    new_im.save(f'{output_name}.png')