Создание записей TFRecord для задач сегментации

05.01.202105.01.2021 Александр Кручинин Tensorflow, TFRecord, Сегментация

Для эффективной работы с данными, в том числе и на TPU, TensorFlow предоставляет формат TFRecord. Формат TFRecord это простой формат для хранения последовательности двоичных записей. Подробнее об этом вы можете почитать здесь:

https://www.tensorflow.org/tutorials/load_data/tfrecord?hl=ru

Здесь же будет показано, как использовать TFRecord для задач сегментации, например сетей UNet. Пример был сделал с использованием Googler Colab. Ссылка на GitHub с кодом представлена внизу.

Первоначально создадим тестовые данные, которые будут добавляться в tfrec. Для этого воспользуемся OpenCV и генератором случайных окружностей

import numpy as np 
import os
import cv2
from PIL import Image
import matplotlib.pyplot as plt
import random
# Создать данные для помещения в  записи
count_records = 100
images = [ ]
random.seed(0)
for i in range(count_records):
  img = np.zeros((512,512,3),dtype=np.uint8)
  circles = random.randint(0,7) # от 0 до 7 окружностей
  for j in range(circles):
    x = random.randint(0,511)
    y = random.randint(0,511)
    r = random.randint(50,200)
    img = cv2.circle(img, (x,y), r, 255, -1)
  images.append(img)

Выведем часть полученных изображений.

# Вывод изображений
columns, rows = 4,4
fig=plt.figure(figsize=(columns*4, rows*4))
for i in range(rows):
  for j in range(columns):
    idx = j+i*columns
    fig.add_subplot(rows, columns, j+i*columns+1)
    plt.axis('off')
    plt.imshow(Image.fromarray(images[idx]))            
plt.show()

Создадим маски путём перевода изображений в градации серого, а затем бинаризуя их.

# Создать маски для изображений
masks=[ ]
for img in images:
  im = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
  ret,im = cv2.threshold(im,1,1,cv2.THRESH_BINARY)
  masks.append(im)
# Вывод масок
columns, rows = 4,4
fig=plt.figure(figsize=(columns*4, rows*4))
for i in range(rows):
  for j in range(columns):
    idx = j+i*columns
    fig.add_subplot(rows, columns, j+i*columns+1)
    plt.axis('off')
    plt.imshow(Image.fromarray(masks[idx]))            
plt.show()
 

Теперь необходимо сохранить все эти изображения в tfrec. Для этого подключаем TensorFlow

import tensorflow as tf

Используем рекомендуемые функции для перевода в byte. Каждая функция принимает на вход скалярное значение и возвращает tf.train.Feature содержащий один из трех вышеприведенных list типов:

def _bytes_feature(value):
  """Returns a bytes_list from a string / byte."""
  if isinstance(value, type(tf.constant(0))):
    value = value.numpy() # BytesList won't unpack a string from an EagerTensor.
  return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))
def _float_feature(value):
  """Returns a float_list from a float / double."""
  return tf.train.Feature(float_list=tf.train.FloatList(value=[value]))
def _int64_feature(value):
  """Returns an int64_list from a bool / enum / int / uint."""
  return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))

Объединяющая функция, чтобы создать один элемент записи из изображения и маски

def serialize_example(image, mask):
  """
  Creates a tf.train.Example message ready to be written to a file.
  """
  # Create a dictionary mapping the feature name to the tf.train.Example-compatible
  # data type.
  feature = {
      'image': _bytes_feature(image),
      'mask': _bytes_feature(mask),
  }
  # Create a Features message using tf.train.Example.
  example_proto = tf.train.Example(features=tf.train.Features(feature=feature))
  return example_proto.SerializeToString()

Создание записи

# Создать tfrec (помните, что оптимальный размер не больше 200Mb)
filename = 'test.tfrec'
with tf.io.TFRecordWriter(filename) as writer:
  for i in range(count_records):
    example = serialize_example(images[i].tobytes(),masks[i].tobytes())
    writer.write(example)  

Видно, что файл размером около 100Mb создан. Далее читаем эти данные.

TRAINING_FILENAMES =[filename]
AUTO = -1
BATCH_SIZE = 4
image_feature_description = {
    'image': tf.io.FixedLenFeature([, tf.string),
    'mask': tf.io.FixedLenFeature([, tf.string)
}
def _parse_image_function(example_proto):
    single_example = tf.io.parse_single_example(example_proto, image_feature_description)
    image = tf.reshape( tf.io.decode_raw(single_example['image'],out_type=np.dtype('uint8')), (512, 512, 3))
    mask =  tf.reshape(tf.io.decode_raw(single_example['mask'],out_type='bool'),(512, 512,1))
    return image,mask# tf.cast(image, tf.float32), tf.cast(mask, tf.float32) # cast as float32 required for TPU
def load_dataset(filenames, ordered=False):
    ignore_order = tf.data.Options()
    if not ordered:
        ignore_order.experimental_deterministic = False
    dataset = tf.data.TFRecordDataset(filenames, num_parallel_reads=AUTO)
    dataset = dataset.with_options(ignore_order)
    dataset = dataset.map(_parse_image_function, num_parallel_calls=AUTO)
    return dataset
def get_training_dataset():
    dataset = load_dataset(TRAINING_FILENAMES)
    dataset = dataset.repeat()
    dataset = dataset.shuffle(128)
    dataset = dataset.batch(BATCH_SIZE,drop_remainder=True)
    dataset = dataset.prefetch(AUTO)
    return dataset
 
for img, mask in get_training_dataset().take(1):
    pass

Смотрим маску и изображение:

plt.imshow(mask[1].numpy().squeeze())

<matplotlib.image.AxesImage at 0x7fc1046ef6a0>

plt.imshow(img[1].numpy().squeeze())
 
<matplotlib.image.AxesImage at 0x7fc104590be0>

На этом все. Для удобства можете загрузить отсюда код в Google Colab и попробовать

https://github.com/recogru/papers/blob/master/TFRecord_sample.ipynb

Вам также может понравиться

Изменение шага learning rate в TensorFlow Keras

Распознавание с помощью Tensorflow на С-API и OpenCV

Основы работы с TensorFlow и Keras на Python. Часть 3