如何基于布尔掩码实现3通道RGB图像的背景去除与分割裁剪

问题说明 #

需要实现基于多通道布尔掩码的图像裁剪逻辑：输入形状为(w, h, c)的RGB图像、形状为(w, h, n)的True/False分割掩码，对每个掩码输出两个结果：

• 裁剪出掩码对应目标的最小外接矩形区域
• 矩形区域内掩码值不为True的背景像素全部置为0，保留不规则形状的分割目标
  原有基础裁剪函数仅能完成矩形裁剪，无法实现背景置零，且存在逻辑冗余：

# crop image based on segmentation mask
def get_segment_crop(img,tol=0, mask=None):
    for i in range(mask.shape[2]):  
        crops = []
        if mask is None:
            mask = img > tol
        for i in range(mask.shape[2]):
            crops.append(img[np.ix_(mask[:,:,i].any(1), mask[:,:,i].any(0))])
        return crops

前期尝试的两种方案均存在问题，可复现代码如下：

import numpy as np
from scipy import misc
import matplotlib.pyplot as plt

file = misc.face()
img = np.array(file)

#create random masks
tf_mask = np.full((img.shape[0], img.shape[1], 2), False)
tf_mask[500:600,500:600,0] = True
tf_mask[500:650, 550:700, 0] = True
tf_mask[100:200, 100:200, 1] = True
tf_mask[200:300, 100:150, 1] = True
# np.where(tf_mask[:,:,0])[0] #sanity check True values were inputed

# 方案1缺陷：直接修改原图像数组，处理后续掩码时原图像已被覆盖，无法获取正确的原图区域
for i in range(tf_mask.shape[2]):
    img[tf_mask[:,:,i]==False] = 0
    plt.imshow(img)
    plt.show()

# 方案2缺陷：逐通道处理后未做维度拼接，无法得到正常的3通道彩色图像
for i in range(tf_mask.shape[2]):
    for j in range(3):
        new_img = np.where(tf_mask[:,:, i]==False, 0, img[:,:,j])
        plt.imshow(new_img)
        plt.show()

实现方案 #

核心逻辑利用numpy的广播机制，无需逐通道循环即可完成3通道图像的掩码置零，所有操作基于裁剪后的局部副本执行，不会破坏原始输入：

import numpy as np

def get_segment_crop(img, tol=0, mask=None):
    if mask is None:
        mask = img > tol
    crops = []
    # 逐个处理每个分割掩码
    for mask_idx in range(mask.shape[2]):
        current_mask = mask[:, :, mask_idx]
        # 计算最小外接矩形的行列边界
        row_coords = np.where(current_mask.any(axis=1))[0]
        col_coords = np.where(current_mask.any(axis=0))[0]
        r_start, r_end = row_coords.min(), row_coords.max() + 1
        c_start, c_end = col_coords.min(), col_coords.max() + 1
        # 裁剪原图和对应掩码到矩形区域
        cropped_img = img[r_start:r_end, c_start:c_end, :].copy()
        cropped_mask = current_mask[r_start:r_end, c_start:c_end]
        # 广播机制自动匹配3通道维度，掩码外区域直接置0
        cropped_img[~cropped_mask] = 0
        crops.append(cropped_img)
    return crops

实现要点 #

• 裁剪时对图像数组调用.copy()生成局部副本，所有修改仅作用于副本，完全不影响原始图像，多掩码处理互不干扰
• numpy广播规则支持形状为(crop_h, crop_w)的2D布尔掩码直接作用于(crop_h, crop_w, 3)的3通道图像，会自动将掩码为False位置的所有通道像素置为0，无需手动遍历通道
• 修复了原函数中重复外层循环、提前返回的逻辑bug，原函数外层多余的for i in range(mask.shape[2])无实际作用，且内层循环第一次执行完就会返回，逻辑不完整

调用示例 #

from scipy import misc
import matplotlib.pyplot as plt

img = np.array(misc.face())
# 构造测试掩码
tf_mask = np.full((img.shape[0], img.shape[1], 2), False)
tf_mask[500:600,500:600,0] = True
tf_mask[500:650, 550:700, 0] = True
tf_mask[100:200, 100:200, 1] = True
tf_mask[200:300, 100:150, 1] = True

# 获取裁剪结果
result_crops = get_segment_crop(img, mask=tf_mask)
# 可视化输出
for idx, crop in enumerate(result_crops):
    plt.subplot(1, 2, idx+1)
    plt.imshow(crop)
plt.show()

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