%matplotlib inline
import matplotlib.pyplot as plt
import numpy as np
import imageio
import tensorflow as tf

import os
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"]="7"

def batch(file_list, batch_size=16):
    fileQ = tf.train.string_input_producer(file_list, shuffle=False)
    reader = tf.WholeFileReader()

    filename, data = reader.read(fileQ)
    image = tf.image.decode_jpeg(data, channels=3)

    img = imageio.imread(file_list[0])
    w, h, c = img.shape
    shape = [w, h, 3]
    
    image.set_shape(shape)
    min_after_dequeue = 10000
    capacity = min_after_dequeue + 3 * batch_size

    queue = tf.train.shuffle_batch(
        [image], batch_size=batch_size,
        num_threads=4, capacity=capacity,
        min_after_dequeue=min_after_dequeue, name='synthetic_inputs')
    
    crop_queue = tf.image.crop_to_bounding_box(queue, 50, 25, 128, 128)
    resized = tf.image.resize_nearest_neighbor(crop_queue, [64, 64])

    return tf.to_float(resized)

import os
dir_ = '/root/CelebA/images'
jpgs = os.listdir(dir_)
filelist = [os.path.join(dir_, jpg) for jpg in jpgs]
plt.imshow(imageio.imread(filelist[0]))

data_loader = batch(filelist)

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    
    coord = tf.train.Coordinator()
    threads = tf.train.start_queue_runners(sess=sess, coord=coord)

    im = sess.run(data_loader)
    
    coord.request_stop()
    coord.join(threads)

im = im.astype(np.uint8)
plt.imshow(im[8, :])