How can I use tf.string_split() in tensorflow?

Question

I want to get the extension of image files to invoke different image decoder, and I found there's a function called tf.string_split in tensorflow r0.11.

filename_queue = tf.train.string_input_producer(filenames, shuffle=shuffle)
reader = tf.WholeFileReader()
img_src, img_bytes = reader.read(filename_queue)
split_result = tf.string_split(img_src, '.')

But when I run it, I get this error：

ValueError: Shape must be rank 1 but is rank 0 for 'StringSplit' (op: 'StringSplit') with input shapes: [], [].

I think it may caused by the shape inference of img_src. I try to use img_src.set_shape([1,]) to fix it, but it seems not work, I get this error:

ValueError: Shapes () and (1,) are not compatible

Also, I can't get the shape of img_src using

tf.Print(split_result, [tf.shape(img_src)],'img_src shape=')

The result is img_src shape=[]. But if I use the following code:

tf.Print(split_result, [img_src],'img_src=')

The result is img_src=test_img/test1.png. Am I doing something wrong?

Answer 1

Just pack img_src into a tensor.

split_result = tf.string_split([img_src], '.')