Search code examples
pythontensorflowimage-processingcomputer-visiontensorflow-datasets

Create Tensorflow Dataset with dataframe of images and labels


I want to create a dataset with tensorflow and feed this with images as array (dtype=unit8) and labels as string. The images and the according labels are stored in a dataframe and the columns named as Image as Array and Labels.

Image as Array (type = array) Labels (type = string)
img_1 'ok'
img_2 'not ok'
img_3 'ok'
img_4 'ok'

My challenge: I don't know how to feed the Dataset out of a dataframe, the most tutorials prefer the way to load the data from a directory.

Thank you in forward and I hope you can help me to load the images in the dataset.


Solution

  • You can actually pass a dataframe directly to tf.data.Dataset.from_tensor_slices:

    import tensorflow as tf
    import numpy as np
    import pandas as pd
    
    
    df = pd.DataFrame(data={'images': [np.random.random((64, 64, 3)) for _ in range(100)],
                            'labels': ['ok', 'not ok']*50})
    
    dataset = tf.data.Dataset.from_tensor_slices((list(df['images'].values), df['labels'].values)).batch(2)
    
    for x, y in dataset.take(1):
      print(x.shape, y)
    # (2, 64, 64, 3) tf.Tensor([b'ok' b'not ok'], shape=(2,), dtype=string)