How to read a datapoint with multiple lables from a tfrecord file

Question

I am writing data with multiple labels for each image, in this case bounding box and classification labels and am using the following function to write the data to a tfrecord :

   def tfr_write_sr(data_split_name,save_dir, label_array, data_array):

       filename = os.path.join(save_dir, data_split_name + '.tfrecords')
       writer = tf.python_io.TFRecordWriter(filename)
       for index in range(data_array.shape[0]):

       image = data_array[index].tostring()
       example = tf.train.Example(features=tf.train.Features(
        feature={
            'height': tf.train.Feature(
                int64_list=tf.train.Int64List(
                    value=[data_array.shape[1]])),
            'width': tf.train.Feature(
                int64_list=tf.train.Int64List(
                    value=[data_array.shape[2]])),
            'depth': tf.train.Feature(
                int64_list=tf.train.Int64List(
                    value=[data_array.shape[3]])),
            'shape_type': tf.train.Feature(
                    int64_list=tf.train.Int64List(
                        value=[int(label_array[index][3])])),
            'bbtl_x': tf.train.Feature(
                    int64_list=tf.train.Int64List(
                        value=[int(label_array[index][1][0])])),
            'bbtl_y': tf.train.Feature(
                    int64_list=tf.train.Int64List(
                        value=[int(label_array[index][1][1])])),
            'bbbr_x': tf.train.Feature(
                    int64_list=tf.train.Int64List(
                        value=[int(label_array[index][0][0])])),
            'bbbr_y': tf.train.Feature(
                    int64_list=tf.train.Int64List(
                        value=[int(label_array[index][0][1])])),                
            'image_raw': tf.train.Feature(
                bytes_list=tf.train.BytesList(
                    value=[image]))}))
         writer.write(example.SerializeToString())
       writer.close() 

I have verified that the record is being written correctly but all the examples I have seen before only read one label per image, how do I read multiple labels ?

Answer 1

First we read in our tfrecord and get its features :

  reader = tf.TFRecordReader()
  _ , serialized_example = reader.read(filename_queue)

   features = tf.parse_single_example(serialized_example, 
        features={
            'image_raw': tf.FixedLenFeature([],tf.string),
            'shape_type' : tf.FixedLenFeature([], tf.int64),
            'bbtl_x' : tf.FixedLenFeature([], tf.int64),
            'bbtl_y' : tf.FixedLenFeature([], tf.int64),
            'bbbr_x' : tf.FixedLenFeature([], tf.int64),
            'bbbr_y' : tf.FixedLenFeature([], tf.int64)
    })

Now we have our features we can use tf.stack() to build a tensor for our multilables and add it to our graph :

     label  = tf.stack([features['shape_type'],
                        features['bbtl_x'],
                        features['bbtl_y'],
                        features['bbbr_x'],
                        features['bbbr_y'] ], axis=0 )


      image = tf.decode_raw(features['image_raw'], tf.uint8)

      images_batch, labels_batch = tf.train.shuffle_batch([image,label],
                                                 batch_size=128,
                                                 capacity=2000,
                                                 min_after_dequeue=1000)