Search code examples
pythonmatplotlibmaskopencv

How to save the 2d array of values generated by semantic-segmenetaion as an image?

When an image is given as input to a semantic-segmentation model its output is a 2D array of values. Each value in that array represents an object that model thinks is present at that position in the original image. The array looks something like this:

print(image_mask)

array([[2, 2, 2, ..., 7, 7, 7],
       [2, 2, 2, ..., 7, 7, 7],
       [2, 2, 2, ..., 7, 7, 7],
       ...,
       [0, 0, 0, ..., 0, 0, 0],
       [0, 0, 0, ..., 0, 0, 0],
       [0, 0, 0, ..., 0, 0, 0]])

When one plots this as an image using matplotlib it adds false colors and makes an image:

plt.imshow(image_mask)

enter image description here

Plotting this array on top of an image results in a mask like affect:

plt.subplot(1,2,1)
plt.imshow(image, 'gray', interpolation='none')
plt.subplot(1,2,2)
plt.imshow(image, 'gray', interpolation='none')
plt.imshow(image_mask, 'jet', interpolation='none', alpha=0.7)
plt.show()

enter image description here

Now my aim is to use the mask created by the model to extract the people in the image. Initially I had thought that image_mask was a 3-channel RGB image and tried to make the yellow in the image to be white and the background to be black, something like this:

image_mask[np.where((image_mask==[253,231,36]).all())] = [255,255,255] # MAKE YELLOW WHITE
image_mask[np.where((image_mask==[68,1,84]).all())] = [0,0,0] # MAKE BACKGROUND BLACK
result = cv2.bitwise_and(image,image,mask = image_mask) # DO BITWISE AND WITH THE ORIGINAL TO GET THE PEOPLE

But I soon realized this won't be possible as image_mask is not an image. What should I do now? Is there a way to convert this array to an image? Or does cv2 offer some method to help in this scenario?

Solution

  • Got it!

    I used matrix multiplication along with np.expand_dims:

    image_mask_copy = image_mask.copy()
np.place(image_mask_copy,image_mask_copy!=15,[0]) # REMOVE ALL EXCEPT PEOPLE, == 0
np.place(image_mask_copy,image_mask_copy==15,[255]) # MAKE PEOPLE == 255
plt.imshow(image_mask_copy)

    enter image description here

    new_image_mask = np.expand_dims(image_mask_copy,-1)*np.ones((1,1,3))
plt.imshow(new_image_mask)

    enter image description here

    Now it is an image.