Python利用mask提取保存图片前景

考虑这么一个情形。二分割任务，输入了一张图像(比方说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通道表透明度的概念。