Converting a 2D segmentation map to n dimensional numpy masks

Question

I have a segmentation map with 10 classes (A numpy array of size (m,n,1) which every element is a number from 1~10 specifying a class that the pixel belongs to). I want to convert it to an array of size (m,n,10) where each channel is mask for elements of that specific class. I can do it using a for loop like this:

for i in range(10):
    mask[:,:,i] = (seg_map==i)[:,:,0]

but I need a faster way to do this. The for loop takes too much time. Is there any built in function that can outperform the for loop.

Thanks in advance.

Answer 1

One approach:

import numpy as np

np.random.seed(42)

# toy data
data = np.random.randint(0, 10, 20).reshape((5, 4, 1))

# https://stackoverflow.com/a/37323404/4001592
n_values = 10
values = data.flatten()
encoded = np.eye(n_values)[data.ravel()].reshape((5, 4, 10))

match = np.allclose(data.reshape(5, 4), encoded.argmax(-1))
print(match)

One way to verify that the output is correct is to verify that the one-hot encoded value matches back with the index, as below:

match = np.allclose(data.reshape(5, 4), encoded.argmax(-1))
print(match)

Output

True