语义分割任务中如何筛选含目标特征的非全黑掩码？

默认你的数据集采用图像、掩码分文件夹存储，两者文件名一一对应的常规组织格式，若命名规则有差异按需调整匹配逻辑即可。

核心实现逻辑 #

• 遍历掩码文件夹下的所有掩码文件，提取文件名用于匹配对应原图
• 加载单张掩码后，直接通过数组判断是否存在待分割目标的像素值（你的场景为像素值=1），无需逐像素遍历，效率更高
• 将符合条件（含目标）的掩码和对应原图，分别拷贝到新的独立文件夹，避免修改原始数据集
• 可选：将筛选后的文件名列表导出为csv文件，后续训练直接读取列表即可，无需重复执行筛选逻辑

Python 实现示例 #

依赖第三方库：numpy、Pillow，可通过pip install numpy pillow安装

import os
import numpy as np
from PIL import Image
import shutil

# 路径配置，按需修改为你的本地路径
MASK_FOLDER = "./data/masks"
IMAGE_FOLDER = "./data/images"
SAVE_MASK_FOLDER = "./data/filtered_masks"
SAVE_IMAGE_FOLDER = "./data/filtered_images"

# 自动创建输出文件夹
os.makedirs(SAVE_MASK_FOLDER, exist_ok=True)
os.makedirs(SAVE_IMAGE_FOLDER, exist_ok=True)

# 遍历所有掩码
for mask_file in os.listdir(MASK_FOLDER):
    # 跳过非图片格式文件
    if not mask_file.lower().endswith(("png", "jpg", "jpeg")):
        continue
    # 加载掩码转为numpy数组
    mask_arr = np.array(Image.open(os.path.join(MASK_FOLDER, mask_file)))
    # 判断是否存在目标像素（像素值为1）
    if np.any(mask_arr == 1):
        # 匹配对应原图，若命名规则不同可修改此处匹配逻辑
        # 示例为掩码和原图同名，仅后缀不同，按需调整
        img_file = mask_file.rsplit(".", 1)[0] + ".jpg"
        # 拷贝符合要求的文件到目标文件夹
        shutil.copy2(
            os.path.join(MASK_FOLDER, mask_file),
            os.path.join(SAVE_MASK_FOLDER, mask_file)
        )
        shutil.copy2(
            os.path.join(IMAGE_FOLDER, img_file),
            os.path.join(SAVE_IMAGE_FOLDER, img_file)
        )

优化建议 #

• 若数据集规模极大，可提前读取所有文件名，用多进程并行加载判断掩码，提升筛选速度
• 若你的掩码存在多类目标，修改np.any(mask_arr == 1)中的判断条件即可，比如筛选同时包含类别1和类别2的样本可以写np.any(mask_arr ==1) and np.any(mask_arr ==2)
• 筛选后仍存在类别不平衡问题的话，可以搭配加权损失、过采样少量类等方法进一步优化训练效果

内容的提问来源于stack exchange，提问作者PietroB.A.