考虑这么一个情形。二分割任务,输入了一张图像(比方说3×256×256),网络给了个单通道的预测mask(1×256×256),现在要利用这个mask把前景提出来。一个例子如下:
图像:
mask:
至于原理的话是很简单的,把这图像和mask给concat在一起,变为四通道图像,即B, G, R, alpha,然后再转换回三通道即可。代码如下:
import os
import cv2
import numpy as np
from PIL import Image
def gen_foreground(img_path, mask_path, res_path):
img = cv2.imread(img_path)
mask = cv2.imread(mask_path, 0)
height, width, channel = img.shape
b, g, r = cv2.split(img)
res = np.zeros((4, height, width), dtype=img.dtype)
res[0][0:height, 0:width] = b
res[1][0:height, 0:width] = g
res[2][0:height, 0:width] = r
res[3][0:height, 0:width] = mask
cv2.imwrite(res_path, cv2.merge(res))
if __name__ == "__main__":
gen_foreground("img.png", "mask.png", "res.png")
生成结果:
注意事项:保存的结果图像应该是png格式,而不能是jpg等格式,因为只有png才存在alpha通道表透明度的概念。