Image merge and compress

sample.py

from PIL import Image
import numpy as np
from os import path, getcwd, listdir, makedirs


def merge_and_compress(image_paths, output_path, output_compression="png"):

    # Load images from file
    images_pil = []
    cols = 0
    rows = 0
    for img_path in image_paths:
        img_pil = Image.open(img_path)
        if img_pil.mode != "RGB":
            img_pil = img_pil.convert('RGB')
        # Get total number of rows needed    
        rows += img_pil.size[1]
        # Get max width needed
        if cols < img_pil.size[0]:
            cols = img_pil.size[0]
        images_pil.append(img_pil)

    channels = len(images_pil[0].getbands())
    datatype = np.asarray(images_pil[0]).dtype

    out_array = np.zeros((rows, cols, channels), dtype=datatype)

    # Place image into output array
    next_image_start_row = 0
    for img in images_pil:
        img = np.array(img)
        out_array[next_image_start_row:next_image_start_row + img.shape[0], :img.shape[1]] = img
        next_image_start_row += img.shape[0]

    # Convert to PIL
    merged_img = Image.fromarray(out_array)
    # Save image in given format after compressing
    merged_img.save(path.join(output_path, "merged_image"), output_compression, quality=30, optimize=True)


if __name__ == "__main__":
    input_dir = path.join(getcwd(), "data")
    output_dir = path.join(getcwd(), "output")
    if not path.exists(output_dir):
        print("Creating output folder")
        makedirs(output_dir)

    img_paths = [path.join(input_dir, f) for f in listdir(input_dir)]
    print(img_paths)
    merge_and_compress(img_paths, output_dir, "JPEG")