How to create the categorical mask for images specifically for Tensor? Or port the NumPy function correctly to Dataset.map function

General Discussion
, ,
#1

Hi, I’m Maifee Ul Asad. And this is my first question to this community. If I make any mistake, please correct me and any help is really appreciated.

Now as for the problem:

I’m trying to move from NumPy array as my dataset to tensorflow.Dataset.

Now, I’ve created a pipeline to train the model for classification problems. At some point, I just normalize all the images using map function:

dataset['train'] = dataset['train'].map(pre_pr, num_parallel_calls=tf.data.experimental.AUTOTUNE)

And the function description looks like this:

@tf.function
def normalize(input_image: tf.Tensor, input_mask: tf.Tensor) -> tuple:
    input_image = tf.cast(input_image, tf.float32) / 255.0
    input_mask= tf.cast(input_mask, tf.float32) / 255.0
    return input_image, input_mask

@tf.function
def pre_pr(datapoint: dict) -> tuple:
    input_image = tf.image.resize(datapoint['image'], (IMG_SIZE, IMG_SIZE))
    input_mask = tf.image.resize(datapoint['y_mask'], (IMG_SIZE, IMG_SIZE))
    return normalize(input_image, input_mask)

This works pretty fine.

After all the mapping it converts a <MapDataset shapes: {image: (None, None, 3), y_mask: (None, None, 3)}, types: {image: tf.uint8, y_mask: tf.uint8}> to <PrefetchDataset shapes: ((None, 128, 128, 3), (None, 128, 128, 3)), types: (tf.float32, tf.float32)>

But the issue is in any segmentation problem, when I’m trying to create the categorical mask it causes issues.

I have got a function that creates the segmentation mask from images, with a given color palette, like this:

def flat_labels(label):
    label_seg = np.zeros(label.shape, dtype=np.float32)
    for i in range(len(encodings)):
        label_seg[np.all(label == encodings[i], axis=-1)] = i
    return label_seg[:, :, 0]

Here encoding is a list/array, like [255,255,0],[0,0,255]...

When I try to call it like:

input_mask = tf.numpy_function(flat_labels, [input_mask], tf.uint8)

It changes the datatype to something else:

<PrefetchDataset shapes: ((None, 128, 128, 3), <unknown>), types: (tf.float32, tf.uint8)>

I really don’t want to change the rest of the code, I want to just call this flat_labels NumPy-based function to a tensor. Or some way to create a function to create a categorical mask for Tensor.

Thanks in advance.

#3

Okay, resolved it using like this:

@tf.function
def flat_labels(label):
    map = []
    for colour in encodings:
        class_map = tf.reduce_all(tf.equal(label, colour), axis=-1)
        map.append(class_map)
    map = tf.stack(map, axis=-1)
    map = tf.cast(map, tf.float32)
    
    return map

Similar to NumPy function.

Thanks.

1 Like