I am writing a python program to preprocess images to be used as labels for a semantic segmentation task. The original images have three channels, where the vector of three values representing each pixel represent a class label for that pixel. For example, a pixel of [0,0,0] could be class 1, [0,0,255] could be class 2, and so on.
I need to convert these images into a single channel image, with pixel values starting from 0 and increasing serially to represent each class. Essentially, I need to convert [0,0,0] in the old image to 0 in the new image, [0,0,255] to 1, and so on for all the classes.
The images are fairly high resolution, with more than 2000 pixels width and height. I need to do this for hundreds of images. The current approach I have involves iterating over each pixel and replacing the 3-dimensional value with the corresponding scalar value.
filename="file.png"
label_list = [[0,0,0], [0,0,255]] # for example. there are more classes like this
image = imread(filename)
new_image = np.empty((image.shape[0], image.shape[1]))
for i in range(image.shape[0]):
for j in range(image.shape[1]):
for k, label in enumerate(label_list):
if np.array_equal(image[i][j], label):
new_image[i][j] = k
break
imsave("newname.png", new_image)
The problem is that the above program is very inefficient, and takes a few minutes to run for each image. This is too much to handle all my images, and hence I need to improve it.
Firstly, I think it might be possible to remove the innermost loop by converting label_list
to a numpy array and using np.where
. However, I am not sure how to do a np.where
to find a 1-dimensional array inside a two-dimensional array, and whether it would improve anything.
From this thread, I tried to define a function and apply it directly on the image. However, I need to map every 3-dimensional label to a scalar. A dictionary cannot contain a list as a key. Would there be a better way to do this, and would it help?
Is there a way to improve (by a lot) the efficiency, or is there a better way, to do what the above program does?
Thank you.