python 使用TensorFlow的多线程图像阅读器示例

Posted 2021-05-09

#!/usr/bin/python
import sys
from time import time
import tensorflow as tf

num_threads = int(sys.argv[1])

db_dir = "/Data3/cityscapes/scp/"

image_scp = db_dir + "train.img.scp"
# label_scp = db_dir + "train.label.scp"

image_paths = [line.rstrip('\n') for line in open(image_scp)]

is_png = image_paths[0][-3:].lower() == "png"
decoder = tf.image.decode_png if is_png else tf.image.decode_jpeg

image_paths = tf.train.string_input_producer(image_paths, shuffle=False)

# http://stackoverflow.com/questions/33648322/tensorflow-image-reading-display
reader = tf.WholeFileReader()

key, value = reader.read(image_paths)

images = decoder(value)

height = 1024
width = 2048
channel = 3

batch_size = 1
min_after_dequeue = 32
capacity = min_after_dequeue + 3 * batch_size
batch = tf.train.shuffle_batch(
    [images], batch_size=batch_size, capacity=capacity,
    shapes=[[height, width, channel]],
    num_threads = num_threads,
    min_after_dequeue=min_after_dequeue)

counter = 0
N_EXAMPLES = 100 * num_threads
with tf.Session() as sess:
    coord = tf.train.Coordinator()
    threads = tf.train.start_queue_runners(sess=sess, coord=coord)

    print "Starting reading {} examples ... hold tight".format(N_EXAMPLES)
    timer = -time()
    while not coord.should_stop() and counter < N_EXAMPLES:
        image = sess.run(batch)
        counter += 1
        # print "#{:04d}: image.shape = {} [{}]".format(counter, image.shape, image.dtype)

    timer += time()
    print "Took {} seconds for {} examples. {} seconds / example".format(timer, N_EXAMPLES, timer / N_EXAMPLES)
    coord.request_stop()

    coord.join(threads)
    sess.close()
print "Done."