Storing multiple values in a tfrecord feature

Solution 1:

You achieve this in 3 simple steps, although it is hard to say what you actually intend to do without further details:

Create and parse data:

import tensorflow as tf
import pandas as pd
import tabulate
import numpy as np

d = {'image_id': ['0002cc93b.jpg', '0007a71bf.jpg', '000a4bcdd.jpg'], 
     'class_1_rle': ['', '18661 28 18863 82...', '131973 1 132228 4...'], 
     'class_2_rle': ['29102 12 29346 24...', '', ''], 
     'class_3_rle': ['', '', '229501 11 229741 33...']}

df = pd.DataFrame(data=d)
default_value = '1 0'
df = df.replace(r'^\s*$', default_value, regex=True)
print(df.to_markdown())

image_ids = np.asarray(df.pop('image_id'))
rle_classes = df.to_numpy()
image_ids_shape = image_ids.shape
rle_classes_shape = rle_classes.shape

image_ids = np.vectorize(lambda x: x.encode('utf-8'))(image_ids).ravel()
rle_classes = np.vectorize(lambda x: x.encode('utf-8'))(rle_classes).ravel()

|    | image_id      | class_1_rle          | class_2_rle          | class_3_rle            |
|---:|:--------------|:---------------------|:---------------------|:-----------------------|
|  0 | 0002cc93b.jpg | 1 0                  | 29102 12 29346 24... | 1 0                    |
|  1 | 0007a71bf.jpg | 18661 28 18863 82... | 1 0                  | 1 0                    |
|  2 | 000a4bcdd.jpg | 131973 1 132228 4... | 1 0                  | 229501 11 229741 33... |

Create tfrecord:

def bytes_feature(value):
  return tf.train.Feature(bytes_list=tf.train.BytesList(value = value))

def create_example(image_ids, rle_classes):

  feature = {'img_id': bytes_feature(image_ids),
             'rle': bytes_feature(rle_classes)}
  example = tf.train.Example(features = tf.train.Features(feature = feature))
  return example

test_writer = tf.io.TFRecordWriter('data.tfrecords')

example = create_example(image_ids, rle_classes)
test_writer.write(example.SerializeToString())
test_writer.close() 

Read tfrecord:

def parse_tfrecord(example):
  feature = {'img_id': tf.io.FixedLenFeature([image_ids_shape[0]], tf.string),
             'rle': tf.io.FixedLenFeature([rle_classes_shape[0], rle_classes_shape[1]], tf.string)}
  parsed_example = tf.io.parse_single_example(example, feature)
  return parsed_example

serialised_example = tf.data.TFRecordDataset('data.tfrecords')
parsed_example_dataset = serialised_example.map(parse_tfrecord)
parsed_example_dataset = parsed_example_dataset.flat_map(tf.data.Dataset.from_tensor_slices)
for features in parsed_example_dataset:
  print(features['img_id'], features['rle'])

tf.Tensor(b'0002cc93b.jpg', shape=(), dtype=string) tf.Tensor([b'1 0' b'29102 12 29346 24...' b'1 0'], shape=(3,), dtype=string)
tf.Tensor(b'0007a71bf.jpg', shape=(), dtype=string) tf.Tensor([b'18661 28 18863 82...' b'1 0' b'1 0'], shape=(3,), dtype=string)
tf.Tensor(b'000a4bcdd.jpg', shape=(), dtype=string) tf.Tensor([b'131973 1 132228 4...' b'1 0' b'229501 11 229741 33...'], shape=(3,), dtype=string)