【python实现】剔除数据集中不含掩码的部分

摘要：数据集中部分图片不含待分割的目标，如果不想用这部分来训练，可以通过一段代码剔除。

from PIL import Image
import cv2
import os

def extract_mask_and_original(mask_path, original_path, output_folder):
    # 读取黑白掩码图片和同名原图图片
    mask_image = Image.open(mask_path).convert("L")
    original_image = Image.open(original_path)

    has_white_pixel = any(pixel > 0 for pixel in mask_image.getdata())

    if has_white_pixel:
        # 提取同名的原图图片
        original_filename = os.path.basename(mask_path).replace(".png", ".png")
        original_path = os.path.join(input_image_folder, original_filename)

        # 保存掩码图片和原图图片
        output_mask_path = os.path.join(output_folder, os.path.basename(mask_path))
        output_original_path = os.path.join(output_folder_image, original_filename)

        mask_image.save(output_mask_path)
        original_image.save(output_original_path)

        print(f"已提取掩码和原图：{output_mask_path}，{output_original_path}")
    else:
        print(f"掩码中不存在白色像素：{mask_path}")

    print("提取完成")


# 设置输入文件夹和输出文件夹
input_image_folder = r"D:\DDR_png512\val\images"
input_mask_folder = r"D:\DDR_png512\val\masks"

output_folder_image = r"D:\DDR_png512tichu\val\images"
output_folder_mask =  r"D:\DDR_png512tichu\val\masks"
# 遍历输入文件夹中的掩码图片
for mask_filename in os.listdir(input_mask_folder):
    if mask_filename.endswith(".png"):
        mask_path = os.path.join(input_mask_folder, mask_filename)
        original_filename = mask_filename.replace(".png", ".png")
        original_path = os.path.join(input_image_folder, original_filename)

        extract_mask_and_original(mask_path, original_path, output_folder_mask)