导入所需模块

import tensorflow as tf from tensorflow.keras.preprocessing.image import ImageDataGenerator import matplotlib.pyplot as plt import numpy as np import random import cv2

设置随机种子

random.seed(1) np.random.seed(1) tf.random.set_seed(1)

家具分类数据集信息可视化

家具类别标签

classes = ['chair', 'desk', 'lamp', 'sofa', 'wardrobe']

统计训练集和测试集中各个类别的样本数量

train_samples = [] test_samples = [] for c in classes: train_samples.append(len(os.listdir(os.path.join(train_dir, c)))) test_samples.append(len(os.listdir(os.path.join(test_dir, c))))

绘制条形图

plt.bar(classes, train_samples, label='train') plt.bar(classes, test_samples, label='test') plt.legend() plt.title('Furniture Classification Dataset') plt.xlabel('Classes') plt.ylabel('Number of samples') plt.show()

数据增强、归一化、one-hot编码、数据集划分

定义数据增强函数

train_datagen = ImageDataGenerator( rescale=1./255, rotation_range=30, # 随机旋转角度 width_shift_range=0.1, # 随机水平平移 height_shift_range=0.1, # 随机垂直平移 shear_range=0.1, # 随机剪切 zoom_range=0.1, # 随机放缩 horizontal_flip=True, # 随机水平翻转 fill_mode='nearest' # 填充方式 )

test_datagen = ImageDataGenerator(rescale=1./255)

定义训练集和测试集生成器

train_generator = train_datagen.flow_from_directory( train_dir, target_size=(224, 224), # 图像尺寸 batch_size=32, class_mode='categorical' # 多分类 )

test_generator = test_datagen.flow_from_directory( test_dir, target_size=(224, 224), batch_size=32, class_mode='categorical' )

定义one-hot编码函数

def to_one_hot(labels, classes): n = len(labels) one_hot = np.zeros((n, classes)) for i in range(n): one_hot[i, labels[i]] = 1 return one_hot

获取训练集和测试集的标签

train_labels = train_generator.classes test_labels = test_generator.classes

将标签进行one-hot编码

train_labels_one_hot = to_one_hot(train_labels, len(classes)) test_labels_one_hot = to_one_hot(test_labels, len(classes))

定义数据集划分函数

def split_dataset(data, labels, train_ratio=0.8): n = len(data) indices = np.random.permutation(n) train_size = int(n * train_ratio) train_indices = indices[:train_size] test_indices = indices[train_size:] x_train = data[train_indices] y_train = labels[train_indices] x_test = data[test_indices] y_test = labels[test_indices] return x_train, y_train, x_test, y_test

调用数据集划分函数获取训练集和验证集

x_train, y_train, x_val, y_val = split_dataset(train_generator.filenames, train_labels_one_hot)

模型搭建

model = tf.keras.models.Sequential([ tf.keras.layers.Conv2D(32, (3, 3), activation='relu', input_shape=(224, 224, 3)), tf.keras.layers.MaxPooling2D((2, 2)), tf.keras.layers.Conv2D(64, (3, 3), activation='relu'), tf.keras.layers.MaxPooling2D((2, 2)), tf.keras.layers.Conv2D(128, (3, 3), activation='relu'), tf.keras.layers.MaxPooling2D((2, 2)), tf.keras.layers.Flatten(), tf.keras.layers.Dense(256, activation='relu'), tf.keras.layers.Dropout(0.5), tf.keras.layers.Dense(len(classes), activation='softmax') ])

模型参数优化

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

模型编译、训练

history = model.fit( train_generator, steps_per_epoch=len(train_generator), epochs=20, validation_data=test_generator, validation_steps=len(test_generator) )

模型结构显示

model.summary()

准确率及Loss曲线图

plt.plot(history.history['accuracy'], label='train_accuracy') plt.plot(history.history['val_accuracy'], label='val_accuracy') plt.title('Accuracy') plt.xlabel('Epochs') plt.ylabel('Accuracy') plt.legend() plt.show()

plt.plot(history.history['loss'], label='train_loss') plt.plot(history.history['val_loss'], label='val_loss') plt.title('Loss') plt.xlabel('Epochs') plt.ylabel('Loss') plt.legend() plt.show()

图像识别展示

从测试集中随机选取一些样本进行预测

n = 5 indices = np.random.choice(len(test_generator), n) x_test = [] y_true = [] y_pred = [] for i in indices: x, y = test_generator[i] y_true.append(y) x_test.append(x) y_pred.append(model.predict(x))

显示预测结果

fig, axes = plt.subplots(n, 2, figsize=(8, 16)) for i in range(n): # 显示原始图像和标签 axes[i, 0].imshow(x_test[i][0]) axes[i, 0].axis('off') axes[i, 0].set_title('True Label: {}'.format(classes[np.argmax(y_true[i][0])])) # 显示预测结果和置信度 pred = np.argmax(y_pred[i][0]) confidence = np.max(y_pred[i][0]) axes[i, 1].imshow(x_test[i][0]) axes[i, 1].axis('off') axes[i, 1].set_title('Predicted Label: {}\nConfidence: {:.2f}'.format(classes[pred], confidence)) plt.show(