如何在Python中提取掩码黑色标记区域并生成白色背景目标图像？

我有两张图像：一张原始图像，一张对应掩码，其中待提取区域以黑色标记。想要生成白色背景、提取区域保留原始色彩的最终图像（比如只保留彩色衬衫，背景改为白色）。但自己编写的代码只生成了反转颜色的掩码，没能保留彩色衬衫的图像。

原始代码 #

import os
import cv2
import numpy as np

def extract_cloth(input, mask, output):
    if not os.path.exists(mask):
        os.makedirs(mask)

    if not os.path.exists(output):
        os.makedirs(output)

    input_image_filenames = [filename for filename in os.listdir(input) if filename.endswith('.png')]

    for input_filename in input_image_filenames:
        input_path = os.path.join(input, input_filename)
        mask_filename = os.path.splitext(input_filename)[0] + '.png'
        mask_path = os.path.join(mask, mask_filename)

        if os.path.exists(mask_path):
            original_image = cv2.imread(input_path)
            mask_image = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)

            # 反转掩码（黑变白，白变黑）
            inverted_mask = cv2.bitwise_not(mask_image)

            # 用反转掩码提取衣物
            extracted_cloth = cv2.bitwise_and(original_image, original_image, mask=inverted_mask)

            # 创建白色背景
            white_background = 255 * np.ones_like(original_image, dtype=np.uint8)
            
            # 合并提取的衣物和白色背景
            extracted_cloth_on_white = cv2.bitwise_or(extracted_cloth, white_background, mask=inverted_mask)

            output_filename = os.path.splitext(input_filename)[0] + '.png'
            output_path = os.path.join(output, output_filename)
            cv2.imwrite(output_path, extracted_cloth_on_white)
            print(f'已保存: {output_path}')
        else:
            print(f'未找到{input_filename}对应的掩码')

问题原因与修正方案 #

问题出在最后一步的bitwise_or操作错误使用了反转掩码，导致白色背景无法正确填充到非提取区域。正确的合并逻辑应该是：

• 用反转掩码提取出原始图中的待保留区域（黑色提取区反转后为白色，bitwise_and会保留该区域的彩色像素）
• 用原始掩码将白色背景的提取区域置空（即只保留背景的非提取区为白色）
• 将提取的彩色区域和处理后的白色背景叠加

修正后的代码 #

import os
import cv2
import numpy as np

def extract_cloth(input_dir, mask_dir, output_dir):
    # 避免参数名和内置函数input冲突，重命名参数
    if not os.path.exists(mask_dir):
        os.makedirs(mask_dir)

    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    input_image_filenames = [filename for filename in os.listdir(input_dir) if filename.endswith('.png')]

    for input_filename in input_image_filenames:
        input_path = os.path.join(input_dir, input_filename)
        mask_filename = os.path.splitext(input_filename)[0] + '.png'
        mask_path = os.path.join(mask_dir, mask_filename)

        if os.path.exists(mask_path):
            original_image = cv2.imread(input_path)
            mask_image = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)

            # 反转掩码：黑色提取区变为白色，用于保留原始图对应区域
            inverted_mask = cv2.bitwise_not(mask_image)
            # 提取彩色的待保留区域
            extracted_cloth = cv2.bitwise_and(original_image, original_image, mask=inverted_mask)

            # 创建白色背景，并用原始掩码把提取区域（黑色）置为0（空）
            white_background = 255 * np.ones_like(original_image, dtype=np.uint8)
            background_with_hole = cv2.bitwise_and(white_background, white_background, mask=mask_image)

            # 叠加彩色提取区域和带空位的白色背景
            final_image = cv2.add(extracted_cloth, background_with_hole)

            output_filename = os.path.splitext(input_filename)[0] + '.png'
            output_path = os.path.join(output_dir, output_filename)
            cv2.imwrite(output_path, final_image)
            print(f'已保存: {output_path}')
        else:
            print(f'未找到{input_filename}对应的掩码')

关键改动说明 #

1. 重命名参数input为input_dir，避免和Python内置函数冲突
2. 用原始掩码处理白色背景，生成带空位的背景（提取区域为空，非提取区域为白色）
3. 使用cv2.add替代错误的bitwise_or，正确叠加彩色区域和白色背景

内容的提问来源于stack exchange，提问作者Navin Chandra