使用TensorFlow Keras实现灰度掩码转二进制并存入Numpy数组

针对二值图像分割任务（输入为RGB图像，掩码为灰度格式），以下代码实现了将灰度掩码转换为二进制格式，并将输入图像与处理后的掩码统一存储到Numpy数组中：

from tensorflow.keras.preprocessing.image import load_img, ImageDataGenerator, img_to_array
from skimage.transform import resize
import os
import numpy as np
from tqdm import tqdm

im_height, im_width = 256, 256
threshold = 150
ids_test = next(os.walk("data/test/image"))[2]  # 获取指定路径下所有图像的文件名列表
print("图像数量 = ", len(ids_test))

# 初始化存储输入图像和掩码的Numpy数组
X_ts = np.zeros((len(ids_test), im_height, im_width, 3), dtype=np.float32)
Y_ts = np.zeros((len(ids_test), im_height, im_width, 1), dtype=np.float32)

# 遍历处理每个图像与掩码
for n, id_ in tqdm(enumerate(ids_test), total=len(ids_test)):
    # 加载并预处理RGB输入图像
    img = load_img(f"data/test/image/{id_}", color_mode="rgb")
    x_img = img_to_array(img)
    x_img = resize(x_img, (im_height, im_width, 3), mode='constant', preserve_range=True)
    # 归一化到[0,1]区间
    X_ts[n] = x_img / 255.0

    # 加载灰度掩码并转换为二进制格式
    mask = img_to_array(load_img(f"data/test/label/{id_}", color_mode="grayscale"))
    # 阈值二值化：大于阈值的像素设为1，其余为0
    binarized = 1.0 * (mask > threshold)
    # 调整掩码尺寸以匹配输入图像
    mask_resized = resize(binarized, (im_height, im_width, 1), mode='constant', preserve_range=True)
    Y_ts[n] = mask_resized

关键步骤说明 #

• 图像尺寸与阈值定义：统一设置输入图像和掩码的目标尺寸为256×256，指定二值化阈值为150（可根据实际掩码的灰度分布调整）
• 文件列表读取：通过os.walk获取测试集所有图像的文件名，确保图像与掩码一一对应
• 数组初始化：提前创建固定形状的Numpy数组，提升数据加载效率
• RGB图像预处理：加载图像后转换为数组，调整尺寸后归一化到[0,1]，符合Keras模型的输入要求
• 掩码二值化：利用阈值判断将灰度掩码转换为仅包含0和1的二进制掩码，再调整尺寸后存入数组

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