How to apply keras ImageDataGenerator class to TFRecordsDataset to augment it?

Viewed 493

I want to efficiently retrive examples from TFRecords files and to augment them with keras ImageDataGenerator class, but as I understand ImageDataGenerator can only sample numpy arrays, pandas DataFrames or directories (composed by images readable by PIL). I also know that one can transform it into a tf.data.Dataset object with tf.data.Datset.from_generator(), but not of a method for the inverse.

Maybe I should just use ImageDataGenartor.flow_from_directory(), but I think it is slower, although I have not measured it precisely with large datasets. If it is approximately equivalent, I would thank a source stating it. The next code shows an example of how I would like it to be:

from tf.keras.preprocessing.image import ImageDataGenerator
from tf.data import TFRecordDataset;
import tensorflow as tf;
from models import DenseNetBN100;

eps=150; batch_size=32; tot_examples=50000;
imre_shape=(32,32,3); lare_shape=(10,);
train_fname='cifar10_trainRecords';
model = DenseNetBN100(imre_shpae, lare_shape);

def _parse_record(proto, clip=False):
    features = {
      'image':tf.FixedLenFeature([],tf.string),
      'label':tf.FixedLenFeature([],tf.string),
    }    
    example = tf.parse_single_example(proto, features)
    im = tf.decode_raw(example['image'], tf.float32)
    im = tf.reshape(im, imre_shape)
    la = tf.decode_raw(example['label'], tf.int8)
    la = tf.reshape(la, lare_shape);
    la = tf.cast(la, tf.float32);
    return im, la;
dtst = TFRecordDataset(train_fname).map(_parse_record)
dtst=dtst.repeat(eps).shuffle(10000)
dtst=dtst.batch(batch_size)

train_datagen = ImageDataGenerator(
        rotation_range=40,
        width_shift_range=0.2,
        height_shift_range=0.2,
        shear_range=0.2,
        zoom_range=0.2,
        horizontal_flip=True,
        fill_mode='nearest');

train_generator = train_datagen.flow(dtst)
model.fit(train_generator, epochs=eps
    steps_per_epoch=tot_examples//batch_size)

There is still the posibility of getting more efficiency through transforming a ImageDataGenerator into a tf.data.Dataset, thanks to tf.data.Dataset management. But that is still just another theory, and I think it would be better having trustworthy sources than using sparse personal measurements only.

0 Answers
Related