带有ListDirectory的Tensorflow数据集API

时间:2017-12-19 10:42:28

标签: python tensorflow tensorflow-datasets

我想使用Tensorflow数据集API为每个文件夹创建一个批处理(每个文件夹包含图像)。我有以下简单的代码片段:

import tensorflow as tf
import os
import pdb

def parse_file(filename):
    image_string = tf.read_file(filename)
    image_decoded = tf.image.decode_png(image_string)
    image_resized = tf.image.resize_images(image_decoded, [48, 48])
    return image_resized #, label

def parse_dir(frame_dir):
    filenames = tf.gfile.ListDirectory(frame_dir)
    batch = tf.constant(5)
    batch = tf.map_fn(parse_file, filenames)
    return batch

directory = "../Detections/NAC20171125"
# filenames = tf.constant([os.path.join(directory, f) for f in os.listdir(directory)])
frames = [os.path.join(directory, str(f)) for f in range(10)]


dataset = tf.data.Dataset.from_tensor_slices((frames))
dataset = dataset.map(parse_dir)

dataset = dataset.batch(256)
iterator = dataset.make_initializable_iterator()
next_element = iterator.get_next()


with tf.Session() as sess:
    sess.run(iterator.initializer)
    while True:
        try:
            batch = sess.run(next_element)
            print(batch.shape)
        except tf.errors.OutOfRangeError:
            break

但是,tf.gfile.ListDirectory(在parse_dir中)需要一个普通的字符串而不是Tensor。所以现在错误是

TypeError: Expected binary or unicode string, got <tf.Tensor 'arg0:0' shape=() dtype=string>

有一种简单的方法可以解决这个问题吗?

1 个答案:

答案 0 :(得分:3)

这里的问题是tf.gfile.ListDirectory()是一个需要Python字符串的Python函数,而frame_dir的{​​{1}}参数是parse_dir()。因此,您需要等效的TensorFlow操作来列出目录中的文件,而tf.data.Dataset.list_files()(基于tf.matching_files())可能是最接近的等价物。

tf.Tensor