当前位置：首页 > article >正文

基于Python实现的完整解决方案，用于对包含四个类别的1500张图像数据集进行分割、训练模型，并提供简易前端和可视化结果

article 2025/2/10 11:37:59

以下是一个基于Python实现的完整解决方案，用于对包含四个类别的1500张图像数据集进行分割、训练模型，并提供简易前端和可视化结果。我们将使用Keras构建一个简单的卷积神经网络（CNN）模型，使用Streamlit创建简易前端，使用Scikit-learn进行模型评估和可视化。

步骤概述

数据准备：加载和预处理数据集，将其划分为训练集、验证集和测试集。
模型构建：构建一个简单的CNN模型。
模型训练：使用训练集对模型进行训练。
模型评估：使用测试集评估模型，计算F1分数、准确率、混淆矩阵和PR曲线。
前端开发：使用Streamlit创建简易前端，展示训练结果和可视化图表。

代码实现

import os
import numpy as np
import tensorflow as tf
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense
from sklearn.metrics import f1_score, confusion_matrix, precision_recall_curve
import matplotlib.pyplot as plt
import seaborn as sns
import streamlit as st

# 数据准备
data_dir = 'your_data_directory'  # 替换为你的数据集目录
img_height = 150
img_width = 150
batch_size = 32

# 数据增强和归一化
train_datagen = ImageDataGenerator(
    rescale=1./255,
    shear_range=0.2,
    zoom_range=0.2,
    horizontal_flip=True,
    validation_split=0.2
)

train_generator = train_datagen.flow_from_directory(
    data_dir,
    target_size=(img_height, img_width),
    batch_size=batch_size,
    class_mode='categorical',
    subset='training'
)

validation_generator = train_datagen.flow_from_directory(
    data_dir,
    target_size=(img_height, img_width),
    batch_size=batch_size,
    class_mode='categorical',
    subset='validation'
)

# 模型构建
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(128, activation='relu'),
    Dense(4, activation='softmax')
])

# 模型编译
model.compile(optimizer='adam',
              loss='categorical_crossentropy',
              metrics=['accuracy'])

# 模型训练
epochs = 10
history = model.fit(
    train_generator,
    steps_per_epoch=train_generator.samples // batch_size,
    validation_data=validation_generator,
    validation_steps=validation_generator.samples // batch_size,
    epochs=epochs
)

# 模型评估
test_generator = train_datagen.flow_from_directory(
    data_dir,
    target_size=(img_height, img_width),
    batch_size=batch_size,
    class_mode='categorical',
    shuffle=False
)

y_pred = model.predict(test_generator)
y_pred_classes = np.argmax(y_pred, axis=1)
y_true = test_generator.classes

# 计算F1分数和准确率
f1 = f1_score(y_true, y_pred_classes, average='weighted')
accuracy = np.mean(y_true == y_pred_classes)

# 混淆矩阵
cm = confusion_matrix(y_true, y_pred_classes)

# PR曲线
precision = dict()
recall = dict()
for i in range(4):
    precision[i], recall[i], _ = precision_recall_curve(y_true == i, y_pred[:, i])

# 前端开发
st.title('图像分类模型训练结果')

# 显示F1分数和准确率
st.write(f'F1分数: {f1:.2f}')
st.write(f'准确率: {accuracy:.2f}')

# 显示混淆矩阵
fig_cm, ax_cm = plt.subplots()
sns.heatmap(cm, annot=True, fmt='d', ax=ax_cm)
ax_cm.set_xlabel('预测标签')
ax_cm.set_ylabel('真实标签')
st.pyplot(fig_cm)

# 显示PR曲线
fig_pr, ax_pr = plt.subplots()
for i in range(4):
    ax_pr.plot(recall[i], precision[i], lw=2, label=f'类别 {i}')
ax_pr.set_xlabel('召回率')
ax_pr.set_ylabel('精确率')
ax_pr.legend()
st.pyplot(fig_pr)