如何将语义分割生成的二维值数组保存为图像?

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

当图像作为语义分割模型的输入时,其输出是一个二维值数组。该数组中的每个值代表模型认为存在于原始图像中该位置的对象。该数组看起来像这样:

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]])

当使用 matplotlib 将其绘制为图像时,它会添加假颜色并制作图像:

plt.imshow(image_mask)

将此数组绘制在图像之上会产生类似 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()

现在我的目标是使用模型创建的蒙版来提取图像中的人物。一开始我以为image_mask是一个3通道的RGB图像,想把图像中的黄色变成白色,背景变成黑色,像这样:

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

但我很快意识到这是不可能的,因为 image_mask 不是图像。我现在该怎么办?有没有办法将这个数组转换为图像?或者 cv2 是否提供了一些方法来帮助解决这种情况?

知道了!

我使用了矩阵乘法和 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)


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

现在是图片了。