imggen_reproduce.py

import os
import sys
import shutil
import argparse

import numpy as np
from scipy.misc import imsave
from keras.preprocessing.image import ImageDataGenerator


def pop_imgs(dirname):
    """Put 128 images in dirname"""
    for i in range(128):
        fname = "{}.jpg".format(i)
        img = np.random.randint(0, 256, size=(32, 32, 3), dtype=np.uint8)
        imsave(os.path.join(dirname, fname), img)


def main(dirname):
    if os.path.isdir(dirname):
        shutil.rmtree(dirname)
    os.makedirs(dirname)
    try:
        pop_imgs(dirname)
        gen = ImageDataGenerator()
        batches = gen.flow_from_directory(dirname, class_mode=None, shuffle=False)
        batch = next(batches)
    finally:
        shutil.rmtree(dirname)
    return 0


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("img_test_dir",
                        help="directory to use as a scratchpad. "
                             "will be deleted and used and deleted again.")
    args = parser.parse_args()
    sys.exit(main(args.img_test_dir))