import tensorflow as tf import numpy as np import os from sklearn.model_selection import train_test_split from tensorflow.keras.preprocessing.image import ImageDataGenerator from tensorflow.keras.applications.vgg16 import VGG16 from tensorflow.keras.layers import Dense, Flatten, Dropout from tensorflow.keras.models import Sequential, Model from tensorflow.keras.optimizers import Adam from tensorflow.keras.callbacks import TensorBoard, ModelCheckpoint

读取图片路径和标签

def load_data(): datapath = './pokemon' imgpaths = [] labels = [] for dirname in os.listdir(datapath): label = int(dirname.split('_')[0]) for filename in os.listdir(os.path.join(datapath, dirname)): imgpaths.append(os.path.join(datapath, dirname, filename)) labels.append(label) imgpaths = np.array(imgpaths) labels = np.array(labels) return imgpaths, labels

随机打乱数据

def shuffle_data(imgpaths, labels): indices = np.arange(len(imgpaths)) np.random.shuffle(indices) imgpaths = imgpaths[indices] labels = labels[indices] return imgpaths, labels

划分训练集、测试集和验证集

def split_data(imgpaths, labels): x_train, x_test, y_train, y_test = train_test_split(imgpaths, labels, test_size=0.2, stratify=labels, random_state=42) x_train, x_val, y_train, y_val = train_test_split(x_train, y_train, test_size=0.2, stratify=y_train, random_state=42) return x_train, y_train, x_val, y_val, x_test, y_test

数据增强

def data_augmentation(x_train): datagen_train = ImageDataGenerator(rotation_range=20, width_shift_range=0.2, height_shift_range=0.2, horizontal_flip=True, vertical_flip=False) datagen_train.fit(x_train) return datagen_train

加载预训练模型VGG16

def load_pretrained_model(): base_model = VGG16(include_top=False, weights='imagenet', input_shape=(224,224,3)) x = base_model.output x = Flatten()(x) x = Dense(1024, activation='relu')(x) x = Dropout(0.5)(x) predictions = Dense(5, activation='softmax')(x) model = Model(inputs=base_model.input, outputs=predictions) for layer in base_model.layers: layer.trainable = False return model

训练模型

def train_model(model, x_train, y_train, x_val, y_val): datagen_train = data_augmentation(x_train) optimizer = Adam(lr=0.001) model.compile(optimizer=optimizer, loss='sparse_categorical_crossentropy', metrics=['accuracy']) tensorboard = TensorBoard(log_dir='./logs', histogram_freq=0, write_graph=True, write_images=False) checkpoint = ModelCheckpoint(filepath='./models/model.h5', monitor='val_accuracy', save_best_only=True) history = model.fit(datagen_train.flow(x_train, y_train, batch_size=32), validation_data=(x_val, y_val), epochs=50, callbacks=[tensorboard, checkpoint]) return history

主函数

def main(): imgpaths, labels = load_data() imgpaths, labels = shuffle_data(imgpaths, labels) x_train, y_train, x_val, y_val, x_test, y_test = split_data(imgpaths, labels) model = load_pretrained_model() history = train_model(model, x_train, y_train, x_val, y_val)

if name == 'main': main()