tensorflow: Combining to_categorical() with Datasets and map()

Viewed 2882

In short, I am struggling in trying to convert an image per-pixel category mask in a tf.data.Dataset from an integer-class encoding to a one-hot encoding.

Consider the image segmentation tensorflow tutorial example here: https://www.tensorflow.org/tutorials/images/segmentation.

The input is an image and the output is a per-pixel integer-labeled category mask. In their example, the mask has a category value at each pixel represented by an integer: {0, 1, or 2}.

The train and test variables are of type tf.data.Dataset and each sample is an (image,mask) tuple.

This form of mask/output is consistent with the sparse_categorical_crossentropy loss function in the tutorial. However, I would like to be able to use other loss functions that require a one-hot encoding instead.

I have been attempting to convert the datasets via the tf.keras.utils.to_categorical() function using a map() call, ie.:

    def mask_to_categorical(image, mask):
        mask = tf.keras.utils.to_categorical(mask,3)
        return image, mask
    
    train = train.map(mask_to_categorical)

However, this fails with an error such as:

    {...}/tensorflow_core/python/keras/utils/np_utils.py:40 to_categorical
        y = np.array(y, dtype='int')

    TypeError: __array__() takes 1 positional argument but 2 were given

Note:

My searching thus far has pointed towards eager/non-eager issues as one possible cause. For what it's worth, I verified that I am running in eager mode via:

    >>> print('tf.executing_eagerly() = ', tf.executing_eagerly())
    
    tf.executing_eagerly() =  True

Any suggestions? Thank you!

1 Answers

Try modifying your function for the one-hot encoding like this:

def mask_to_categorical(image, mask):
    mask = tf.one_hot(tf.cast(mask, tf.int32), 3)
    mask = tf.cast(mask, tf.float32)
    return image, mask
Related