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

Python实例题：基于 TensorFlow 的图像识别与分类系统

news 2025/8/11 16:14:22

Python实例题

题目

基于 TensorFlow 的图像识别与分类系统

问题描述

开发一个基于 TensorFlow 的图像识别与分类系统，包含以下功能：

图像分类模型：基于预训练模型的图像分类器
数据处理与增强：图像预处理和数据增强
模型训练与评估：自定义数据集上的模型训练
API 服务：提供图像识别的 RESTful API
前端界面：用户上传图像并获取分类结果

解题思路

使用 TensorFlow 和 Keras 构建深度学习模型
基于预训练模型（如 ResNet、VGG、EfficientNet）进行迁移学习
设计数据处理和增强管道
使用 Flask 或 FastAPI 构建 API 服务
开发前端界面实现图像上传和结果展示

关键代码框架

import tensorflow as tf
from tensorflow.keras.applications import ResNet50, EfficientNetB0
from tensorflow.keras.layers import Dense, GlobalAveragePooling2D
from tensorflow.keras.models import Model
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint
import numpy as np
import os
from flask import Flask, request, jsonify, render_template
from PIL import Image
import io
import base64# 配置参数
IMAGE_SIZE = (224, 224)
BATCH_SIZE = 32
NUM_CLASSES = 10  # 根据实际数据集调整
EPOCHS = 50
BASE_MODEL = 'efficientnet'  # 可选 'resnet' 或 'efficientnet'# 创建数据增强和预处理
def create_data_generators(train_dir, val_dir):# 训练数据生成器（包含数据增强）train_datagen = ImageDataGenerator(rescale=1./255,rotation_range=20,width_shift_range=0.2,height_shift_range=0.2,shear_range=0.2,zoom_range=0.2,horizontal_flip=True,fill_mode='nearest')# 验证数据生成器（仅缩放）val_datagen = ImageDataGenerator(rescale=1./255)# 生成训练数据train_generator = train_datagen.flow_from_directory(train_dir,target_size=IMAGE_SIZE,batch_size=BATCH_SIZE,class_mode='categorical')# 生成验证数据val_generator = val_datagen.flow_from_directory(val_dir,target_size=IMAGE_SIZE,batch_size=BATCH_SIZE,class_mode='categorical')return train_generator, val_generator# 构建模型
def build_model(input_shape, num_classes, base_model_type='efficientnet'):# 选择基础模型if base_model_type == 'resnet':base_model = ResNet50(weights='imagenet',include_top=False,input_shape=input_shape)else:  # efficientnetbase_model = EfficientNetB0(weights='imagenet',include_top=False,input_shape=input_shape)# 冻结基础模型的所有层for layer in base_model.layers:layer.trainable = False# 添加自定义层x = base_model.outputx = GlobalAveragePooling2D()(x)x = Dense(1024, activation='relu')(x)predictions = Dense(num_classes, activation='softmax')(x)# 构建最终模型model = Model(inputs=base_model.input, outputs=predictions)# 编译模型model.compile(optimizer='adam',loss='categorical_crossentropy',metrics=['accuracy'])return model# 模型微调（解冻部分层）
def fine_tune_model(model, num_layers_to_unfreeze=20):# 解冻最后几层for layer in model.layers[-num_layers_to_unfreeze:]:layer.trainable = True# 重新编译模型，使用较低的学习率model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=1e-5),loss='categorical_crossentropy',metrics=['accuracy'])return model# 训练模型
def train_model(model, train_generator, val_generator, epochs=EPOCHS, model_path='model.h5'):# 设置回调函数callbacks = [EarlyStopping(monitor='val_loss', patience=5, restore_best_weights=True),ModelCheckpoint(model_path, monitor='val_accuracy', save_best_only=True)]# 训练模型history = model.fit(train_generator,steps_per_epoch=train_generator.samples // BATCH_SIZE,validation_data=val_generator,validation_steps=val_generator.samples // BATCH_SIZE,epochs=epochs,callbacks=callbacks)return history, model# 预测函数
def predict_image(model, image_path=None, image_bytes=None, class_names=None):# 从文件路径或字节数据加载图像if image_path:img = Image.open(image_path).convert('RGB')elif image_bytes:img = Image.open(io.BytesIO(image_bytes)).convert('RGB')else:raise ValueError("必须提供图像路径或图像字节数据")# 调整图像大小img = img.resize(IMAGE_SIZE)# 转换为numpy数组并归一化img_array = np.array(img) / 255.0img_array = np.expand_dims(img_array, axis=0)# 预测predictions = model.predict(img_array)predicted_class = np.argmax(predictions[0])confidence = np.max(predictions[0])# 获取类别名称if class_names and predicted_class < len(class_names):class_name = class_names[predicted_class]else:class_name = f"Class {predicted_class}"return {"class": class_name,"confidence": float(confidence),"all_predictions": predictions.tolist()[0]}# 创建Flask应用
app = Flask(__name__)# 加载模型和类别名称
model = None
class_names = None@app.before_first_request
def load_model_and_classes():global model, class_names# 加载训练好的模型model = tf.keras.models.load_model('model.h5')# 加载类别名称（从训练数据生成或手动定义）if os.path.exists('class_names.txt'):with open('class_names.txt', 'r') as f:class_names = [line.strip() for line in f.readlines()]@app.route('/')
def index():return render_template('index.html')@app.route('/predict', methods=['POST'])
def predict():try:# 获取上传的图像file = request.files['image']if not file:return jsonify({"error": "未提供图像文件"}), 400# 读取图像数据image_bytes = file.read()# 进行预测result = predict_image(model, image_bytes=image_bytes, class_names=class_names)return jsonify(result)except Exception as e:return jsonify({"error": str(e)}), 500@app.route('/train', methods=['POST'])
def train():try:# 获取训练配置data = request.jsontrain_dir = data.get('train_dir', 'data/train')val_dir = data.get('val_dir', 'data/val')epochs = data.get('epochs', EPOCHS)base_model = data.get('base_model', BASE_MODEL)# 创建数据生成器train_generator, val_generator = create_data_generators(train_dir, val_dir)# 构建模型model = build_model((*IMAGE_SIZE, 3), train_generator.num_classes, base_model)# 训练模型history, model = train_model(model, train_generator, val_generator, epochs)# 保存类别名称class_names = list(train_generator.class_indices.keys())with open('class_names.txt', 'w') as f:f.write('\n'.join(class_names))return jsonify({"message": "模型训练完成","classes": class_names})except Exception as e:return jsonify({"error": str(e)}), 500# 前端模板 (index.html)
<!DOCTYPE html>
<html>
<head><title>图像分类系统</title><style>body {font-family: Arial, sans-serif;max-width: 800px;margin: 0 auto;padding: 20px;text-align: center;}.container {background-color: #f5f5f5;padding: 20px;border-radius: 10px;box-shadow: 0 0 10px rgba(0,0,0,0.1);}h1 {color: #333;}.upload-area {margin: 20px 0;}.upload-btn {background-color: #4CAF50;color: white;padding: 10px 20px;border: none;border-radius: 5px;cursor: pointer;}.upload-btn:hover {background-color: #45a049;}.result-area {margin-top: 20px;padding: 15px;background-color: #fff;border-radius: 5px;min-height: 100px;}.image-preview {max-width: 100%;height: auto;margin: 20px 0;border-radius: 5px;}</style>
</head>
<body><div class="container"><h1>图像分类系统</h1><div class="upload-area"><input type="file" id="imageUpload" accept="image/*" style="display: none;"><button class="upload-btn" onclick="document.getElementById('imageUpload').click()">选择图像</button><button class="upload-btn" id="predictBtn" onclick="predictImage()" disabled>预测</button></div><div><img id="imagePreview" class="image-preview" src="" alt="图像预览"></div><div class="result-area" id="resultArea"><p>请上传一张图像进行分类</p></div></div><script>let selectedImage = null;document.getElementById('imageUpload').addEventListener('change', function(e) {if (this.files && this.files[0]) {const reader = new FileReader();reader.onload = function(e) {document.getElementById('imagePreview').src = e.target.result;document.getElementById('resultArea').innerHTML = '<p>图像已加载，请点击预测按钮</p>';document.getElementById('predictBtn').disabled = false;selectedImage = this.files[0];}.bind(this);reader.readAsDataURL(this.files[0]);}});function predictImage() {if (!selectedImage) {alert('请先选择一张图像');return;}const formData = new FormData();formData.append('image', selectedImage);document.getElementById('resultArea').innerHTML = '<p>正在预测，请稍候...</p>';fetch('/predict', {method: 'POST',body: formData}).then(response => response.json()).then(data => {if (data.error) {document.getElementById('resultArea').innerHTML = `<p>错误: ${data.error}</p>`;} else {let resultHTML = `<h3>预测结果</h3>`;resultHTML += `<p>类别: ${data.class}</p>`;resultHTML += `<p>置信度: ${(data.confidence * 100).toFixed(2)}%</p>`;document.getElementById('resultArea').innerHTML = resultHTML;}}).catch(error => {document.getElementById('resultArea').innerHTML = `<p>错误: ${error.message}</p>`;});}</script>
</body>
</html># 训练脚本示例
if __name__ == "__main__":# 创建数据生成器train_generator, val_generator = create_data_generators('data/train', 'data/val')# 构建模型model = build_model((*IMAGE_SIZE, 3), train_generator.num_classes, BASE_MODEL)# 训练模型print("开始训练基础模型...")history, model = train_model(model, train_generator, val_generator, EPOCHS, 'base_model.h5')# 模型微调print("开始微调模型...")model = fine_tune_model(model)history, model = train_model(model, train_generator, val_generator, EPOCHS//2, 'fine_tuned_model.h5')# 保存类别名称class_names = list(train_generator.class_indices.keys())with open('class_names.txt', 'w') as f:f.write('\n'.join(class_names))print("模型训练完成！")