Image Folder -> Thumbnail Grid

make_image_grid.py

import argparse
import os
import random

import cv2
import numpy as np

###############################################################################

IMG_EXTENSIONS = ['.jpg', '.JPG', '.jpeg', '.JPEG', '.png', '.PNG']


def is_image_file(filename):
    return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)


def make_dataset(dir):
    images = []
    assert os.path.isdir(dir), '%s is not a valid directory' % dir

    for root, _, fnames in sorted(os.walk(dir)):
        for fname in fnames:
            if is_image_file(fname):
                path = os.path.join(root, fname)
                images.append(path)

    return images

###############################################################################


def get_args():
    parser = argparse.ArgumentParser()

    parser.add_argument('--img_dir', required=True)
    parser.add_argument('--thumbnail_height', type=int, default=128)
    parser.add_argument('--thumbnail_width', type=int, default=128)
    parser.add_argument('--out_height', type=int, default=1024)
    parser.add_argument('--out_width', type=int, default=2048)
    parser.add_argument('--out_img_path', type=str, default='grid.jpg')

    args = parser.parse_args()
    return args


if __name__ == '__main__':
    args = get_args()
    print(args)

    img_paths = make_dataset(args.img_dir)
    random.shuffle(img_paths)

    grid_img = np.zeros((args.out_height, args.out_width, 3), dtype=np.uint8)

    num_rows = args.out_height//args.thumbnail_height
    num_cols = args.out_width//args.thumbnail_width
    total = len(img_paths)

    for row in range(num_rows):
        for col in range(num_cols):
            img_idx = min(row * num_cols + col, total)
            img = cv2.imread(img_paths[img_idx])
            img_thumbnail = cv2.resize(
                img, (args.thumbnail_width, args.thumbnail_height))
            grid_img[row*args.thumbnail_height:(row+1)*args.thumbnail_height,
                     col*args.thumbnail_width:(col+1)*args.thumbnail_width] = \
                img_thumbnail

    cv2.imwrite(args.out_img_path, grid_img)