How can I convert an image from pixels to one-hot encodings?

Question

I have a PNG image I am loading within Tensorflow using:

image = tf.io.decode_png(tf.io.read_file(path), channels=3)

The image contains pixels that match a lookup like this:

image_colors = [
    (0, 0, 0),        # black
    (0.5, 0.5, 0.5),  # grey
    (1, 0.5, 1),      # pink
]

How can I convert it so that the output has the pixels mapped into one-hot encodings where the hot component would be the matching color?

Answer 1

Let me assume for convenience that all values in image_colors are in [0, 255]:

image_colors = [
    (0, 0, 0),  # black
    (127, 127, 127),  # grey
    (255, 127, 255),  # pink
]

My approach maps pixels into one-hot values as follows:

# Create a "color reference" tensor from image_colors
color_reference = tf.cast(tf.constant(image_colors), dtype=tf.uint8)

# Load the image and obtain tensor with one-hot values
image = tf.io.decode_png(tf.io.read_file(path), channels=3)
comp = tf.equal(image[..., None, :], color_reference)
one_hot = tf.cast(tf.reduce_all(comp, axis=-1), dtype=tf.float32)

Note that you can easily add new colors to image_colors without changing the TF implementation. Also, this assumes that all pixels in the image are in image_colors. If that is not the case, one could define a range for each color and then use other tf.math operations (e.g. tf.greater and tf.less) instead of tf.equal.

Answer 2

There might be better approaches than this.

def map_colors(pixel):
    if pixel[0] < 10 and pixel[1] < 10 and pixel[2] < 10:   ## Black
        return 0
    elif pixel[0] > 245 and pixel[1] > 245 and pixel[2] > 245:  ## White
        return 1
    else:
        return 11


image = tf.io.decode_png(tf.io.read_file(path), channels=3)
img_shape = image.shape
# Arrange the pixels in RGB format from 3D array to 2D array.
image = tf.reshape(image, [-1, 3])

colors = tf.map_fn(lambda x: map_colors(x), image)
one_hot = tf.one_hot(colors, depth=12)
print(one_hot)

# If necessary
image_one_hot = tf.reshape(one_hot, [img_shape[0], img_shape[1], 12])

Ensure that in map_colors, you put all your 12 colors with the range of RGB color values they can accept. Make sure that all the combinations are covered or else, add an extra class of None of the above and name it 12.