Converting a 2D segmentation map to n dimensional numpy masks-CodePudding

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.

CodePudding user response：

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