自定义数据集图像分类实现

模型训练
要使用自己的图片分类数据集进行训练,这意味着数据集应该包含一个目录,其中每个子目录代表一个类别,子目录中包含该类别的所有图片。以下是一个使用Keras和TensorFlow加载自定义图片数据集进行分类训练的例子。
我们自己创建的数据集结构如下:

data/
    train/
        class1/
            img1.jpg
            img2.jpg
            ...
        class2/
            imga.jpg
            imgb.jpg
            ...
        ...
    validation/
        class1/
            img3.jpg
            img4.jpg
            ...
        class2/
            imgc.jpg
            imgd.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是训练数据所在的目录,这样类别顺序才是正确的。如果训练数据和测试数据不在同一个目录下,你需要确保测试数据的类别顺序与训练数据相同。

  • 9
    点赞
  • 1
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值