在 tfrecord 功能中存储多个值

Question

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

我正在尝试使用上表创建tfrecords。我需要以rle 的形式将rle（运行长度编码）功能组合在一起。例如。最终tfrecord 中的功能看起来像

img_id: b'0002cc93b.jpg'
rle: [b'1 0'  b'29102 12 29346 24...'  b'1 0']

img_id: b'000a4bcdd.jpg'
rle: [b'131973 1 132228 4...'  b'1 0'  b'229501 11 229741 33...']

rle 特征应包含对应图像的所有 3 个掩码的rles 作为字符串，空的rle 应编码为'1 0'

我尝试使用列表。但它给出了以下错误

TypeError: ['29102 12 29346 24 29602 24 29858 24 30114 24 30370 24 30626 24 30882 24 31139 23 31395 23 31651 23 has type list, but expected one of: bytes

Answer 1

您只需 3 个简单的步骤即可实现这一目标，但如果没有更多详细信息，很难说出您的实际意图：

创建和解析数据：

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... |

创建 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() 

读取 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)

Answer 2

我找到了适合我情况的整体solution。

在一个特征中存储多个值的具体solution。

用pandas替换df中的空rles为'1 0'

一个从df中获取图像对应的rles的函数。

def rle_class_1(image_id):
    temp_df = df['class_1_rle'][df['image_id'] == image_id]
    for rle in temp_df:
        rle_tensor = tf.constant(rle)
        return rle_tensor.numpy() 

class_2 和 class_3 的类似函数。

创建 tfrecord

paths_dict = dict(zip(file_ids, file_paths))

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

def _int_feature(list_of_ints):
    return tf.train.Feature(int64_list=tf.train.Int64List(value=list_of_ints))

def image_bits_from_id(image_id):
    image = Image.open(paths_dict[image_id])
    image = tf.constant(image)
    image_bits = tf.image.encode_jpeg(image, optimize_size=True, chroma_downsampling=False)
    image_bits = image_bits.numpy()
    return image_bits

def create_tfrec_example(image_id):

    image = image_bits_from_id(image_id) 
    rle_1 = rle_class_1(image_id)
    rle_2 = rle_class_2(image_id)
    rle_3 = rle_class_3(image_id)
    
    feature = {
        'image': _bytestring_feature([image]),
        'img_id': _bytestring_feature([image_id.encode()]),
        'rle': _bytestring_feature([rle_1, rle_2, rle_3])
        }

    tfrec_example = tf.train.Example(features=tf.train.Features(feature=feature))
    return tfrec_example

解析和查看记录

def parse_tfrecord_fn(example):
    features = {
        'image': tf.io.FixedLenFeature([], tf.string), 
        'img_id': tf.io.FixedLenFeature([], tf.string)
        }
    features['rle'] = tf.io.FixedLenFeature([3], tf.string)

    example = tf.io.parse_single_example(example, features)
    example["image"] = tf.io.decode_jpeg(example["image"], channels=3)
    return example

raw_dataset = tf.data.TFRecordDataset(TFREC[0])   # TFREC is a shard
parsed_dataset = raw_dataset.map(parse_tfrecord_fn)

for features in parsed_dataset.take(5):
    for key in features.keys():
        if key != "image":
            print(f"{key}: {features[key]}")

    print(f"Image shape: {features['image'].shape}")
    plt.figure(figsize=(7, 7))
    plt.imshow(features["image"].numpy())
    plt.show()

输出

img_id: b'0002cc93b.jpg'
rle: [b'29102 12 29346 24 29602 24 29858 24... '
 b'1 0' b'1 0']
Image plot