34、基于自定义自编码器和预训练模型的图像上色项目

基于自定义自编码器和预训练模型的图像上色项目

1. 项目概述

本项目主要围绕图像上色任务展开，包含两个部分：一是构建自定义自编码器模型进行图像上色；二是使用预训练的 VGG16 模型作为编码器，实现图像特征提取和上色，以节省训练时间并提升特征提取效果。

2. 自定义自编码器模型

2.1 数据加载

可以从 Kaggle 下载数据集，也可以从指定链接下载：

# 从 Kaggle 下载
#!pip install -q kaggle
#!mkdir ~/.kaggle
#!touch ~/.kaggle/kaggle.json
#api_token = {"username":"Your UserName", "key":"Your key"}
#import json
#with open('/root/.kaggle/kaggle.json', 'w') as file:
#    json.dump(api_token, file)
#!chmod 600 ~/.kaggle/kaggle.json
#!kaggle datasets download -d thedownhill/art-images-drawings-painting-sculpture-engraving

# 从指定链接下载
!wget --no-check-certificate -r 'https://drive.google.com/uc?export=download&id=1CKs7s_MZMuZFBXDchcL_AgmCxgPBTJXK' -O art-images-drawings-painting-sculpture-engraving.zip
!unzip art-images-drawings-painting-sculpture-engraving.zip

下载并解压后，定义相关变量：

import numpy as np
import os
from skimage.io import imread
from tqdm import tqdm

IMG_WIDTH = 256
IMG_HEIGHT = 256
TRAIN_PATH = '/content/dataset/dataset_updated/training_set/painting/'
train_ids = next(os.walk(TRAIN_PATH))[2]

2.2 数据预处理

检查并移除无法读取的图像，将训练图像调整为 256x256 大小：

missing_count = 0
for n, id_ in tqdm(enumerate(train_ids), total=len(train_ids)):
    path = TRAIN_PATH + id_ + ''
    try:
        img = imread(path)
    except:
        missing_count += 1
print("\n\nTotal missing: " + str(missing_count))

X_train = np.zeros((len(train_ids) - missing_count, IMG_HEIGHT, IMG_WIDTH, 3), dtype=np.uint8)
missing_images = 0
for n, id_ in tqdm(enumerate(train_ids), total=len(train_ids)):
    path = TRAIN_PATH + id_ + ''
    try:
        img = imread(path)
        img = resize(img, (IMG_HEIGHT, IMG_WIDTH), mode='constant', preserve_range=True)
        X_train[n - missing_images] = img
    except:
        missing_images += 1
X_train = X_train.astype('float32') / 255.

创建训练集和测试集：

from sklearn.model_selection import train_test_split
x_train, x_test = train_test_split(X_train, test_size=20)

2.3 训练数据准备

使用 ImageDataGenerator 对图像进行增强，并编写函数创建训练批次：

from tensorflow.keras.preprocessing.image import ImageDataGenerator
from skimage.color import rgb2gray, rgb2lab

datagen = ImageDataGenerator(
    shear_range=0.2,
    zoom_range=0.2,
    rotation_range=20,
    horizontal_flip=True)

def create_training_batches(dataset=X_train, batch_size = 20):
    for batch in datagen.flow(dataset, batch_size=batch_size):
        X_batch = rgb2gray(batch)
        lab_batch = rgb2lab(batch)
        X_batch = lab_batch[:,:,:,0]
        X_batch = X_batch.reshape(X_batch.shape+(1,))
        Y_batch = lab_batch[:,:,:,1:] / 128
        yield X_batch, Y_batch

2.4 模型定义

构建自编码器模型，包括编码器和解码器：

from tensorflow.keras.models import Model
from tensorflow.keras.layers import Input, Conv2D, UpSampling2D

# 编码器
inputs1 = Input(shape=(IMG_WIDTH, IMG_HEIGHT, 1,))
encoder_output = Conv2D(64, (3,3), activation='relu', padding='same', strides=2)(inputs1)
encoder_output = Conv2D(128, (3,3), activation='relu', padding='same')(encoder_output)
encoder_output = Conv2D(128, (3,3), activation='relu', padding='same', strides=2)(encoder_output)
encoder_output = Conv2D(256, (3,3), activation='relu', padding='same')(encoder_output)
encoder_output = Conv2D(256, (3,3), activation='relu', padding='same', strides=2)(encoder_output)
encoder_output = Conv2D(512, (3,3), activation='relu', padding='same')(encoder_output)
encoder_output = Conv2D(512, (3,3), activation='relu', padding='same')(encoder_output)
encoder_output = Conv2D(256, (3,3), activation='relu', padding='same')(encoder_output)

# 解码器
decoder_output = Conv2D(128, (3,3), activation='relu', padding='same')(encoder_output)
decoder_output = UpSampling2D((2, 2))(decoder_output)
decoder_output = Conv2D(64, (3,3), activation='relu', padding='same')(decoder_output)
decoder_output = Conv2D(64, (3,3), activation='relu', padding='same')(decoder_output)
decoder_output = UpSampling2D((2, 2))(decoder_output)
decoder_output = Conv2D(32, (3,3), activation='relu', padding='same')(decoder_output)
decoder_output = Conv2D(2, (3, 3), activation='tanh', padding='same')(decoder_output)
decoder_output = UpSampling2D((2, 2))(decoder_output)

model = Model(inputs=inputs1, outputs=decoder_output)
model.compile(loss='mse', optimizer='adam', metrics=['accuracy'])
print(model.summary())

2.5 模型训练

BATCH_SIZE = 20
model.fit_generator(create_training_batches(X_train,BATCH_SIZE),
                    epochs= 100,
                    verbose=1,
                    steps_per_epoch=X_train.shape[0]/BATCH_SIZE)

2.6 模型测试

import matplotlib.pyplot as plt
from skimage.color import lab2rgb

test_image = rgb2lab(x_test)[:,:,:,0]
test_image = test_image.reshape(test_image.shape+(1,))
output = model.predict(test_image)
output = output * 128

generated_images = np.zeros((len(output),256, 256, 3))
for i in range(len(output)):
    cur = np.zeros((256, 256, 3))
    cur[:,:,0] = test_image[i][:,:,0]
    cur[:,:,1:] = output[i]
    generated_images[i] = lab2rgb(cur)

plt.figure(figsize=(20, 6))
for i in range(10):
    # 灰度图
    plt.subplot(3, 10, i + 1)
    plt.imshow(rgb2gray(x_test)[i].reshape(256, 256))
    plt.gray()
    plt.axis('off')
    # 上色图
    plt.subplot(3, 10, i + 1 +10)
    plt.imshow(generated_images[i].reshape(256, 256,3))
    plt.axis('off')
    # 原图
    plt.subplot(3, 10, i + 1 + 20)
    plt.imshow(x_test[i].reshape(256, 256,3))
    plt.axis('off')
plt.tight_layout()
plt.show()

2.7 对未知图像进行推理

!wget https://raw.githubusercontent.com/Apress/artificial-neural-networks-with-tensorflow-2/main/ch14/mountain.jpg
from skimage.io import imread
from skimage.transform import resize

img = imread("mountain.jpg")
img = resize(img, (IMG_HEIGHT, IMG_WIDTH), mode='constant', preserve_range=True)
img = img.astype('float32') / 255.
test_image = rgb2lab(img)[:,:,0]
test_image = test_image.reshape((1,)+test_image.shape+(1,))
output = model.predict(test_image)
output = output * 128
plt.imshow(img)
plt.axis('off')

2.8 自定义自编码器流程

graph LR
    A[数据加载] --> B[数据预处理]
    B --> C[训练数据准备]
    C --> D[模型定义]
    D --> E[模型训练]
    E --> F[模型测试]
    F --> G[未知图像推理]

3. 使用预训练的 VGG16 模型作为编码器

3.1 项目描述

使用与上一个项目相同的图像数据集，仅更改模型定义和推理部分。由于 VGG16 是在 224x224 大小的图像上训练的，因此将图像大小调整为 224x224：

IMG_WIDTH = 224
IMG_HEIGHT = 224

3.2 模型定义

使用 VGG16 的前 18 层作为编码器：

from tensorflow.keras.applications.vgg16 import VGG16
from tensorflow.keras.models import Sequential

vggmodel = VGG16()
newmodel = Sequential()
for i, layer in enumerate(vggmodel.layers):
    if i < 19:
        newmodel.add(layer)
newmodel.summary()
for layer in newmodel.layers:
    layer.trainable = False

3.3 预训练模型使用流程

graph LR
    A[数据加载与预处理] --> B[调整图像大小]
    B --> C[模型定义（使用 VGG16 前 18 层）]
    C --> D[模型训练与推理]

通过以上步骤，我们可以完成基于自定义自编码器和预训练 VGG16 模型的图像上色任务。使用预训练模型可以显著节省训练时间，并提升特征提取效果。

4. 对比分析

4.1 自定义自编码器与预训练 VGG16 模型的特点

模型类型	优点	缺点
自定义自编码器	- 可根据具体任务灵活设计网络结构 - 对特定数据集有更好的适应性	- 训练时间长 - 特征提取能力可能不如预训练模型
预训练 VGG16 模型	- 训练速度快 - 特征提取能力强	- 网络结构固定，灵活性较差 - 可能需要根据具体任务进行微调

4.2 性能对比

• 训练时间 ：自定义自编码器在 GPU 上每个 epoch 训练时间略超过一分钟，而使用预训练的 VGG16 模型每个 epoch 训练时间可降至约一秒。
• 特征提取效果 ：预训练的 VGG16 模型经过大规模数据集的训练，能够提取更丰富、更具代表性的图像特征，在图像上色任务中通常能得到更好的结果。

4.3 适用场景

• 自定义自编码器 ：适用于数据集具有独特特征，需要专门设计网络结构的场景；或者对模型灵活性要求较高的情况。
• 预训练 VGG16 模型 ：适用于数据集规模较小，希望借助预训练模型的强大特征提取能力，快速获得较好结果的场景。

5. 注意事项

5.1 数据处理

• 图像大小 ：在使用不同模型时，要确保图像大小与模型要求一致。如使用 VGG16 模型时，需将图像调整为 224x224 大小；使用自定义自编码器时，图像大小为 256x256。
• 数据增强 ：合理使用数据增强技术可以提高模型的泛化能力，但在测试阶段不要对测试数据进行增强操作。

5.2 模型训练

• 损失函数和优化器 ：选择合适的损失函数和优化器对模型性能有重要影响。在本项目中，使用均方误差（MSE）作为损失函数，Adam 作为优化器。
• 训练轮数 ：训练轮数过多可能导致过拟合，过少则可能导致欠拟合。需要根据具体情况进行调整。

5.3 模型保存与加载

在实际应用中，为了方便后续使用，可以将训练好的模型保存下来：

# 保存模型
model.save('my_model.h5')

# 加载模型
from tensorflow.keras.models import load_model
loaded_model = load_model('my_model.h5')

6. 总结

6.1 主要内容回顾

• 我们首先介绍了自定义自编码器的构建过程，包括数据加载、预处理、模型定义、训练、测试和对未知图像的推理。
• 然后使用预训练的 VGG16 模型作为编码器，通过调整图像大小和重新定义模型，实现了图像上色任务。
• 最后对两种方法进行了对比分析，并给出了使用过程中的注意事项。

6.2 未来展望

• 可以尝试使用其他预训练模型，如 ResNet、Inception 等，进一步提升图像上色的效果。
• 结合生成对抗网络（GAN），可以生成更逼真、更自然的上色图像。
• 探索将图像上色技术应用于更多领域，如历史图像修复、动漫上色等。

6.3 整体流程总结

graph LR
    A[选择模型类型] --> B{自定义自编码器}
    A --> C{预训练 VGG16 模型}
    B --> D[数据加载与处理（256x256）]
    C --> E[数据加载与处理（224x224）]
    D --> F[自定义模型构建与训练]
    E --> G[使用 VGG16 前 18 层构建模型并训练]
    F --> H[模型评估与应用]
    G --> H

通过本项目的实践，我们可以看到图像上色是一个有趣且具有挑战性的任务，不同的模型和方法各有优劣。合理选择和应用这些技术，可以为我们带来更好的图像上色效果。