求输入输出均为图像的Keras神经网络示例

嘿，这里有个完全符合你需求的Keras实现示例——彩色图像转灰度图像的模型，附带完整的训练和推理流程，你可以直接参考上手～

这个模型以3通道彩色图像为输入，输出1通道灰度图像，采用编码器-解码器结构，能很好地保留图像的空间特征，适合像素级的图像转换任务。

1. 模型结构定义 #

import tensorflow as tf
from tensorflow.keras import layers, Model

def build_color_to_gray_model(input_shape=(256, 256, 3)):
    # 输入层：接收256x256的彩色图像
    inputs = layers.Input(shape=input_shape)
    
    # 编码器：逐步提取图像特征并缩小尺寸
    x = layers.Conv2D(64, (3,3), activation='relu', padding='same')(inputs)
    x = layers.MaxPooling2D((2,2), padding='same')(x)
    x = layers.Conv2D(128, (3,3), activation='relu', padding='same')(x)
    x = layers.MaxPooling2D((2,2), padding='same')(x)
    
    # 解码器：还原图像尺寸并重建灰度图
    x = layers.Conv2DTranspose(128, (3,3), strides=(2,2), activation='relu', padding='same')(x)
    x = layers.Conv2DTranspose(64, (3,3), strides=(2,2), activation='relu', padding='same')(x)
    
    # 输出层：1通道灰度图，用sigmoid确保输出值在0-1之间（对应归一化后的像素值）
    outputs = layers.Conv2D(1, (3,3), activation='sigmoid', padding='same')(x)
    
    model = Model(inputs=inputs, outputs=outputs)
    return model

# 初始化并查看模型结构
model = build_color_to_gray_model()
model.summary()

2. 数据准备 #

假设你的数据集按以下结构存放：

• ./color_images/：存放所有彩色输入图像
• ./gray_images/：存放对应的灰度标签图像（文件名与彩色图一致）

我们用自定义生成器来加载并预处理数据：

import os
import numpy as np
from tensorflow.keras.preprocessing.image import load_img, img_to_array

def custom_data_generator(color_dir, gray_dir, batch_size=32):
    color_files = [f for f in os.listdir(color_dir) if f.endswith(('.jpg', '.png'))]
    while True:
        np.random.shuffle(color_files)
        # 按批次加载图像
        for i in range(0, len(color_files), batch_size):
            batch_color, batch_gray = [], []
            for filename in color_files[i:i+batch_size]:
                # 加载并归一化彩色图像
                color_img = load_img(os.path.join(color_dir, filename), target_size=(256,256))
                color_arr = img_to_array(color_img) / 255.0
                batch_color.append(color_arr)
                
                # 加载并归一化灰度图像（标签）
                gray_img = load_img(os.path.join(gray_dir, filename), target_size=(256,256), color_mode='grayscale')
                gray_arr = img_to_array(gray_img) / 255.0
                batch_gray.append(gray_arr)
                
            yield np.array(batch_color), np.array(batch_gray)

# 初始化训练生成器
train_gen = custom_data_generator('./color_images', './gray_images', batch_size=32)

3. 模型训练与推理 #

# 编译模型：用MSE损失（像素级回归任务常用），Adam优化器
model.compile(optimizer='adam', loss='mse', metrics=['mae'])

# 开始训练
train_steps = len(os.listdir('./color_images')) // 32
model.fit(
    train_gen,
    steps_per_epoch=train_steps,
    epochs=50
)

# 单张图像推理示例
def predict_gray_image(model, color_image_path):
    # 加载并预处理输入图像
    color_img = load_img(color_image_path, target_size=(256,256))
    color_arr = img_to_array(color_img) / 255.0
    color_arr = np.expand_dims(color_arr, axis=0)  # 添加batch维度
    
    # 预测并转换为可显示的灰度图
    gray_pred = model.predict(color_arr)[0]
    gray_pred = (gray_pred * 255).astype(np.uint8)
    
    # 保存或显示结果
    from PIL import Image
    result_img = Image.fromarray(gray_pred.squeeze(), mode='L')  # 去掉多余的通道维度
    result_img.save('predicted_gray.jpg')
    result_img.show()

# 测试推理
predict_gray_image(model, './test_color.jpg')

额外：边缘图转全彩图像的模型思路 #

如果需要实现边缘图转全彩图像，只需要做以下调整：

1. 输入层改为(256,256,1)（单通道边缘图）
2. 输出层改为Conv2D(3, ...)（3通道彩色图）
3. 训练时将标签替换为对应的全彩图像
   模型结构可以沿用上面的编码器-解码器，或者改用更强大的U-Net结构来提升效果。

内容的提问来源于stack exchange，提问作者Peter Veselinović