JPG 前景与 mask 图合成 PNG 前景

将 jpg 格式的前景和 mask 图融合，输出 png 格式的前景

# -*- coding: utf-8 -*-
# @Time : 2022/6/23 11:12
# @Author : lc
# @File : test.py
# @Software: PyCharm
# @Composite png formate foreground

import os
import cv2
import numpy as np


def comPng(jpg_dir, mask_dir, png_dir):
    file_list = os.listdir(jpg_dir)
    for file in sorted(file_list):
        img_file = os.path.join(jpg_dir, file)
        mask_file = os.path.join(mask_dir, file)
        save_file = os.path.join(png_dir, file)
        if os.path.exists(save_file):
            print("{} passed!".format(file))
            continue
        else:
            os.makedirs(save_file)
        img_list = os.listdir(img_file)
        for item in sorted(img_list):
            img_path = os.path.join(img_file, item)
            pha_path = os.path.join(mask_file, item)
            save_path = os.path.join(save_file, item).replace('jpg', 'png')
            img = cv2.imread(img_path)
            pha = cv2.imread(pha_path, 0)
            height, width, channel = img.shape
            b, g, r = cv2.split(img)
            fgr = np.zeros((4, height, width), dtype=img.dtype)
            fgr[0][0:height, 0:width] = b
            fgr[1][0:height, 0:width] = g
            fgr[2][0:height, 0:width] = r
            fgr[3][0:height, 0:width] = pha
            cv2.imwrite(save_path, cv2.merge(fgr))
        print("{} saved!".format(file))


if __name__ == "__main__":
    jpg_dir = r"./fgr_jpg"
    mask_dir = r"./mask"
    png_dir = r"./fgr_png"
    comPng(jpg_dir, mask_dir, png_dir)