使用 TensorFlow 进行图像识别

使用 TensorFlow 进行图像识别

图像识别是深度学习的一个关键应用，而 TensorFlow 及其 tf.keras API 为构建和训练图像识别模型提供了强大的工具。本章将指导您如何使用按文件夹组织的自定义数据集创建模型来识别图像。

在本例中，假设您有一个图像数据集，其中每个图像类别都存储在主 dataset_image 目录下的一个单独的子文件夹中。例如： dataset_image/ class_A/ image1.jpg image2.png … class_B/ imageX.jpg imageY.jpeg … 这是组织图像数据集的常见方式。

我们将使用 tf.keras.utils.image_dataset_from_directory 来加载图像，这是一个方便的工具，适用于此类文件夹结构。它会自动从目录名推断类别标签，并创建一个 tf.data.Dataset 对象，这对于训练非常高效。

# 步骤 1: 导入所需的库
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import matplotlib.pyplot as plt # 用于可视化 (可选)
import numpy as np
import os # 如有需要，用于模拟数据集创建

# --- (可选) 用于创建演示用虚拟数据集的辅助函数 ---
def create_dummy_dataset(base_dir="dataset_image_demo", num_classes=2, img_per_class=10, img_size=(64,64)):
    if os.path.exists(base_dir):
        import shutil
        shutil.rmtree(base_dir) # 如果存在则清理
    os.makedirs(base_dir, exist_ok=True)
    for i in range(num_classes):
        class_dir = os.path.join(base_dir, f"class_{i}")
        os.makedirs(class_dir, exist_ok=True)
        for j in range(img_per_class):
            # 创建虚拟图像（随机噪声）
            img_array = np.random.rand(img_size[0], img_size[1], 3) * 255
            img = keras.utils.array_to_img(img_array.astype('uint8'))
            img.save(os.path.join(class_dir, f"img_{j}.png"))
    print(f"已在 {base_dir} 创建虚拟数据集")
    return base_dir

# 如果您没有数据集，则创建一个虚拟数据集
# 替换为您实际的数据集路径
IMAGE_DATASET_PATH = create_dummy_dataset(img_size=(64,64))
# IMAGE_DATASET_PATH = "path/to/your/dataset_image"

# --- 步骤 2: 使用 image_dataset_from_directory 加载数据 ---
IMG_WIDTH = 64
IMG_HEIGHT = 64
BATCH_SIZE = 32

# 创建训练数据集
train_dataset = keras.utils.image_dataset_from_directory(
    IMAGE_DATASET_PATH,
    labels='inferred',       # 从目录名推断标签
    label_mode='categorical', # 对于多类别分类，生成独热编码标签
    image_size=(IMG_HEIGHT, IMG_WIDTH),
    interpolation='nearest', # 用于调整大小
    batch_size=BATCH_SIZE,
    shuffle=True,
    seed=42,
    validation_split=0.2,   # 使用 20% 的数据作为验证集
    subset='training'       # 指定这是训练集
)

# 创建验证数据集
validation_dataset = keras.utils.image_dataset_from_directory(
    IMAGE_DATASET_PATH,
    labels='inferred',
    label_mode='categorical',
    image_size=(IMG_HEIGHT, IMG_WIDTH),
    interpolation='nearest',
    batch_size=BATCH_SIZE,
    shuffle=False, # 验证数据无需打乱
    seed=42,
    validation_split=0.2,
    subset='validation'     # 指定这是验证集
)

class_names = train_dataset.class_names
num_classes = len(class_names)
print(f"找到的类别: {class_names}, 类别数量: {num_classes}")

# --- 步骤 3: 为提升性能配置数据集并标准化数据 ---
AUTOTUNE = tf.data.AUTOTUNE

train_dataset = train_dataset.cache().prefetch(buffer_size=AUTOTUNE)
validation_dataset = validation_dataset.cache().prefetch(buffer_size=AUTOTUNE)

# 将像素值从 [0, 255] 标准化到 [0, 1]
normalization_layer = layers.Rescaling(1./255)

train_dataset = train_dataset.map(lambda x, y: (normalization_layer(x), y))
validation_dataset = validation_dataset.map(lambda x, y: (normalization_layer(x), y))

数据加载过程会读取图像、调整大小、推断类别标签，并将数据分割成训练集和验证集。prefetch 和 cache 方法优化了数据管道。

# --- 步骤 4: 定义 CNN 模型 ---
model = keras.Sequential([
    layers.Input(shape=(IMG_HEIGHT, IMG_WIDTH, 3)), # 显式输入层
    layers.Conv2D(32, (3, 3), padding='same', activation='relu', kernel_constraint=keras.constraints.MaxNorm(3)),
    layers.Dropout(0.2),
    layers.Conv2D(32, (3, 3), activation='relu', padding='same', kernel_constraint=keras.constraints.MaxNorm(3)),
    layers.MaxPooling2D(pool_size=(2, 2)),
    layers.Flatten(),
    layers.Dense(512, activation='relu', kernel_constraint=keras.constraints.MaxNorm(3)),
    layers.Dropout(0.5),
    layers.Dense(num_classes, activation='softmax') # 类别数量从数据中推断得出
])

# --- 步骤 5: 编译模型 ---
# 使用原始示例中的 SGD 优化器，但 Adam 通常是一个不错的默认选择
epochs_train = 10 # 为加快演示而减少
lrate = 0.01
decay = lrate / epochs_train
sgd_optimizer = keras.optimizers.SGD(learning_rate=lrate, momentum=0.9, decay=decay, nesterov=False)

model.compile(loss='categorical_crossentropy',
              optimizer=sgd_optimizer,
              metrics=['accuracy'])

model.summary()

# --- 步骤 6: 训练模型 ---
print("\n正在训练模型...")
# 添加 TensorBoard 回调以进行监控 (可选)
tensorboard_callback = keras.callbacks.TensorBoard(log_dir="./logs_image_rec", histogram_freq=1)

history = model.fit(
    train_dataset,
    validation_data=validation_dataset,
    epochs=epochs_train,
    callbacks=[tensorboard_callback] # 在此处添加回调
)

# --- 步骤 7: 评估模型 ---
print("\n正在评估模型...")
loss, accuracy = model.evaluate(validation_dataset, verbose=0)
print(f"验证集损失: {loss:.4f}")
print(f"验证集准确率: {accuracy*100:.2f}%")

# --- 步骤 8: 保存模型 (可选) ---
# 保存整个模型的现代方式（架构、权重、优化器状态）
model.save("my_image_recognition_model.keras")
print("已将模型保存到磁盘文件 my_image_recognition_model.keras")

# 原始方式（将架构保存到 JSON，权重保存到 HDF5）
model_json = model.to_json()
with open("model_face_arch.json", "w") as json_file:
    json_file.write(model_json)
model.save_weights("model_face_weights.h5")
print("已将模型架构保存到 model_face_arch.json，权重保存到 model_face_weights.h5")

# 加载这种类型的已保存模型:
# from tensorflow.keras.models import model_from_json
# with open('model_face_arch.json', 'r') as json_file:
#     loaded_model_json = json_file.read()
# loaded_model = model_from_json(loaded_model_json)
# loaded_model.load_weights('model_face_weights.h5')
# print("已从磁盘加载模型")

输出将显示模型摘要、训练进度和最终评估指标。如果您使用虚拟数据集运行此代码，准确率可能接近随机（例如，对于 2 个类别，准确率约为 50%），因为图像是随机噪声。使用真实数据集，您会期望获得有意义的准确率。

此更新后的示例使用了 TensorFlow 2.x 的现代功能进行数据加载和模型训练，使过程更加流畅和高效。

要运行 TensorBoard 并可视化训练过程，请打开终端并运行：tensorboard --logdir ./logs_image_rec