import os import tensorflow as tf import numpy as np import cv2 from sklearn.model_selection import train_test_split

设置参数

IMG_SIZE = 70 NUM_CLASSES = 5 EPOCHS = 20 BATCH_SIZE = 64 LEARNING_RATE = 0.001

读取图像路径并生成每张图像的路径和标签

def get_dataset(img_folder): images = [] labels = [] for label, pokemon in enumerate(os.listdir(img_folder)): for image_filename in os.listdir(os.path.join(img_folder, pokemon)): img = cv2.imread(os.path.join(img_folder, pokemon, image_filename)) if img is not None: img = cv2.resize(img, (IMG_SIZE, IMG_SIZE)) images.append(img) labels.append(label) return np.array(images), np.array(labels)

加载数据集

X, y = get_dataset('pokemon_dataset')

打乱数据集

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, stratify=y) X_train, X_val, y_train, y_val = train_test_split(X_train, y_train, test_size=0.2, stratify=y_train)

构建模型

model = tf.keras.Sequential([ tf.keras.layers.Conv2D(32, (3, 3), activation='relu', input_shape=(IMG_SIZE, IMG_SIZE, 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(512, activation='relu'), tf.keras.layers.Dropout(0.5), tf.keras.layers.Dense(NUM_CLASSES, activation='softmax') ])

编译模型

model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=LEARNING_RATE), loss='sparse_categorical_crossentropy', metrics=['accuracy'])

使用TensorBoard对模型进行可视化

tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir='./logs', histogram_freq=1)

训练模型

model.fit(X_train, y_train, epochs=EPOCHS, batch_size=BATCH_SIZE, validation_data=(X_val, y_val), callbacks=[tensorboard_callback])

评估模型

test_loss, test_acc = model.evaluate(X_test, y_test) print('Test Accuracy: {}'.format(test_acc))

保存模型

model.save('pokemon_model.h5')