本文主要是介绍自定义数据集图像分类实现，希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！
模型训练
要使用自己的图片分类数据集进行训练，这意味着数据集应该包含一个目录，其中每个子目录代表一个类别，子目录中包含该类别的所有图片。以下是一个使用Keras和TensorFlow加载自定义图片数据集进行分类训练的例子。
我们自己创建的数据集结构如下：
data/train/class1/img1.jpgimg2.jpg...class2/imga.jpgimgb.jpg......validation/class1/img3.jpgimg4.jpg...class2/imgc.jpgimgd.jpg......
以下是训练模型的代码：
import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense, Dropout
from tensorflow.keras.preprocessing.image import ImageDataGenerator
# 定义数据集的路径
train_data_path = 'data/train'
validation_data_path = 'data/validation'
# 设置图片大小
img_height = 150
img_width = 150
# 创建ImageDataGenerator实例
train_datagen = ImageDataGenerator(rescale=1./255,rotation_range=40,width_shift_range=0.2,height_shift_range=0.2,shear_range=0.2,zoom_range=0.2,horizontal_flip=True,fill_mode='nearest')
validation_datagen = ImageDataGenerator(rescale=1./255)
# 使用ImageDataGenerator读取数据
train_generator = train_datagen.flow_from_directory(train_data_path,target_size=(img_height, img_width),batch_size=32,class_mode='categorical'
)
validation_generator = validation_datagen.flow_from_directory(validation_data_path,target_size=(img_height, img_width),batch_size=32,class_mode='categorical'
)
# 创建CNN模型
model = Sequential([Conv2D(32, (3, 3), activation='relu', input_shape=(img_height, img_width, 3)),MaxPooling2D(2, 2),Conv2D(64, (3, 3), activation='relu'),MaxPooling2D(2, 2),Conv2D(128, (3, 3), activation='relu'),MaxPooling2D(2, 2),Flatten(),Dense(512, activation='relu'),Dropout(0.5),Dense(train_generator.num_classes, activation='softmax')
])
# 编译模型
model.compile(optimizer='adam',loss='categorical_crossentropy',metrics=['accuracy'])
# 训练模型
history = model.fit(train_generator,steps_per_epoch=train_generator.samples // train_generator.batch_size,epochs=50,validation_data=validation_generator,validation_steps=validation_generator.samples // validation_generator.batch_size
)
# 保存模型
model.save('my_model.h5')
在这个例子中，我们首先定义了训练和验证数据的路径，然后创建了ImageDataGenerator实例来处理图片增强和归一化。接着，我们使用flow_from_directory方法从目录中读取图片数据，并生成批量数据供模型训练使用。然后，我们创建了一个简单的CNN模型，并使用训练和验证数据进行了训练。最后，我们将训练好的模型保存为my_model.h5文件。
模型使用
import tensorflow as tf
from tensorflow.keras.preprocessing import image
import numpy as np
import os
# 加载训练好的模型
model = tf.keras.models.load_model('my_model.h5')
# 定义包含图片的目录
img_directory = 'path_to_directory_with_images'
# 创建ImageDataGenerator实例来获取类别的顺序
datagen = ImageDataGenerator()
generator = datagen.flow_from_directory('path_to_training_data_directory', # 训练数据所在的目录target_size=(150, 150),batch_size=32,class_mode='categorical'
)
# 获取类别到索引的映射
class_indices = generator.class_indices
# 反转映射，从索引到类别名称
inverse_map = {v: k for k, v in class_indices.items()}
# 加载所有图片并进行预处理
img_paths = [os.path.join(img_directory, f) for f in os.listdir(img_directory) if f.endswith(('.png', '.jpg', '.jpeg'))]
img_arrays = [image.load_img(img_path, target_size=(150, 150)) for img_path in img_paths] # 确保这个尺寸与训练时使用的尺寸一致
img_arrays = [image.img_to_array(img) for img in img_arrays]
img_arrays = [img / 255.0 for img in img_arrays] # 归一化
img_array = np.array(img_arrays)
# 使用模型进行预测
predictions = model.predict(img_array)
# 将预测结果转换为类别索引
predicted_classes = np.argmax(predictions, axis=-1)
# 打印预测结果
for i, img_path in enumerate(img_paths):class_index = predicted_classes[i]class_name = inverse_map[class_index]print(f"Image: {img_path}, Predicted class: {class_name}")
我们首先创建了ImageDataGenerator实例并使用了flow_from_directory方法来获取类别到索引的映射。然后，我们将这个映射反转，以便可以从索引得到类别名称。在预测完成后，我们使用predicted_classes中的索引通过inverse_map来获取每个图片对应的类别名称，并打印出来。
请确保path_to_training_data_directory是训练数据所在的目录，这样类别顺序才是正确的。如果训练数据和测试数据不在同一个目录下，你需要确保测试数据的类别顺序与训练数据相同。
这篇关于自定义数据集图像分类实现的文章就介绍到这儿，希望我们推荐的文章对编程师们有所帮助！
