ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

基于PyTorch的垃圾分类图像识别系统:从CNN训练到Flask部署

基于PyTorch的垃圾分类图像识别系统:从CNN训练到Flask部署 简介基于Python与深度学习的垃圾分类系统完整源码已获老师指导并通过评审适合作为毕业设计、期末大作业或课程设计的参考项目。系统主要面向需要快速搭建AI项目完整流程的计算机相关专业学生尤其适合希望以较低成本完成高质量项目的初学者。压缩包内共包含6319个文件约18.98MB其中Python源码文件py/pyc达2900余个另有约2400个mo/po翻译文件以及html、js、css等前端资源与png、svg图片覆盖系统后端逻辑、多语言国际化和可视化界面。已有338人学习下载项目实战性强。源码涵盖数据准备、模型构建、训练评估与分类预测等关键环节目录结构规整代码注释清晰便于逐模块理解并进行二次开发可帮助读者快速搭建可演示的垃圾分类应用为课程任务和毕业答辩提供有力支撑。整体完成度较高适合直接运行验证效果也可结合自身课题思路进行功能扩展。1. 这个“高分项目”到底在做什么一个图像分类系统的一次完整落地一份名为“基于Python与深度学习的垃圾分类系统设计与实现源码”的压缩包解压后通常是一堆 .py 文件加一个图片文件夹而不是能双击运行的安装包。它的本质很朴素用深度学习里的 CNN 图像分类模型把一张垃圾照片映射到“可回收、厨余、有害、其他”这类标签上再套一层 Web 页面完成上传与结果展示。这类系统有真实落地场景小区智能分类桶、校园环保活动、工厂废料分拣演示后端跑的核心就是这套图像分类链路。它解决的不是算法创新问题而是“数据—训练—部署”这条线怎么在有限资源下串起来。适合谁有 Python 基础、想把深度学习从跑通 MNIST 推到交付一个可演示系统的人。新手照着顺序能把训练到部署完整跑一遍熟手可以直接拿推理封装和部署部分的参数边界去改。2. 从压缩包到系统架构先想清楚再动手别急着解压跑训练拿到这种高分项目源码第一反应通常是解压后直接python train.py。我建议先花半小时把系统拆成两条链路再动手不然调参时很容易困在“改了什么都没反应”的境地。所谓两条链路就是训练端和推理端它们对资源、速度和代码组织的要求完全不同。2.1 系统核心链路训练端与推理端是两套不同的代码逻辑垃圾分类系统从功能上分成两段。训练端的输入是海量图片输出是一个权重文件跑在开发机上一次训练几十分钟很正常推理端的输入是一张单图输出是类别和置信度要求在演示环境里几百毫秒出结果。一个项目如果在课设答辩时卡住多半是推理端没有做好独立封装而把训练逻辑和预测逻辑混在一个脚本里。常见的高分源码布局是train.py、predict.py、app.py三个文件各司其职。train.py管数据加载、模型定义、训练循环和保存权重predict.py管加载权重并对单张图片输出预测app.py把 predict 封装成 Web 接口。这个布局本身就是最好的学习材料先读predict.py理解它期望的输入尺寸和预处理方式再回来看train.py确认训练输出是否匹配。这里有一个高频坑训练时输入是 224×224推理时却忘了缩放直接拿手机原图喂给模型结果类别完全乱猜。所以拿到源码第一个动作永远是核对两端预处理是否一致包括 resize 尺寸、归一化的 mean/std 值、通道顺序。这些参数差一个精度就会雪崩到接近随机猜测。设备选择也要在系统设计时定好。训练端可以依赖 GPU推理端可能只有 CPU。推理代码里应该这样处理device torch.device(cuda if torch.cuda.is_available() else cpu)。这样在没 GPU 的机器上也能演示只是慢一点。很多源码会写这段不要觉得它多余它是保证项目在不同环境都能跑的关键。2.2 模型选型为什么 ResNet 和 MobileNet 是源码里的常客垃圾分类在算法层面是图像分类不是目标检测。这类项目的数据集通常一张图只包含一个主体比如一个塑料瓶、一张废纸不需要检测模型去画边界框。只有当需求变成“识别桌面上的塑料瓶和纸团各在什么位置”时才需要升级到 YOLO 系列。想清楚这一点选型就不会被带偏。图像分类里最常出现的是 ResNet 和 MobileNet。ResNet 用残差连接缓解了深层网络的梯度消失问题在 ImageNet 上表现可靠是课设、毕设里的稳妥牌。MobileNet 用深度可分离卷积大幅降低计算量模型小、速度快适合部署在老旧笔记本或树莓派上。如果演示电脑配置一般选 MobileNet 会比 ResNet 流畅很多。预训练权重是另一个决定性因素。ImageNet 上训练好的模型已经学会识别边缘、纹理、物体局部结构这些通用特征对垃圾图片同样有效。垃圾分类数据集通常只有几千张从零训练 CNN 很难收敛而加载预训练权重后只微调分类层往往几十个 epoch 就能达到可用效果。所以源码里出现resnet18(pretrainedTrue)或新版weights参数是正常操作不要改成随机初始化。输入尺寸上224×224 是预训练模型的标准输入。改小能提速但丢精度改大不一定涨点还费显存。我的习惯是先用 224×224 跑通再靠数据增强模拟尺度变化而不是一上来就上高分辨率。2.3 数据集准备类别不平衡比模型结构更容易决定成败垃圾分类数据集常见两种形态。一种是四分类可回收、厨余、有害、其他每类几百到几千张适合快速演示。另一种是细分类比如塑料瓶、玻璃、金属罐、果皮、电池等几十类贴近真实场景但对数据量和标注质量的要求高得多。拿到源码先看类别数再和你被要求的任务对比不一致就得重新准备数据。数据划分是第一道工程决策。我一般按 8:1:1 分成训练、验证、测试划分前先做全量随机打散。有人图省事直接按文件夹切结果某一类的图片全部来自同一个拍摄批次的同背景验证集准确率虚高拿到现场拍的真实照片就翻车。用random.shuffle或sklearn.train_test_split先打散一次成本很低收益很大。类别不平衡是垃圾分类里最常见的翻车原因。比如“其他”3000 张、“有害”200 张模型会直接忽略少数类。两个缓解手段加权采样让 DataLoader 给少数类更高的采样概率数据增强对少数类做翻转、旋转、色彩抖动变相扩充样本量。但数据增强不是越强越好。垃圾图片和 ImageNet 的自然图差别很大真实场景里背景复杂可能出现在桌面、草地、桶边。过强的随机遮挡会把瓶子主体特征盖掉模型学不到东西。我平时只保留随机水平翻转、小角度旋转、轻微色度调整以 loss 能平稳下降为准而不是把增强堆到极致。2.4 技术栈与版本先统一环境再谈跑通技术栈基本由标题锁定Python 负责脚本逻辑深度学习框架负责网络训练再加 Web 框架做界面展示。常见的组合是 PyTorch torchvision Flask。PyTorch 对中小型项目最友好torchvision 自带预训练模型Flask 轻量够用且容易改。版本是新手最容易踩的隐形坑。PyTorch 1.x 和 2.x 在 API 上有差异torchvision 的预训练权重加载方式也变过。拿到源码别急着pip install torch先看requirements.txt或代码里的 import 写法按源码锁定的版本装。如果源码没给版本我的做法是装当前稳定版遇到报错逐个查而不是来回升降级。环境建议用 Anaconda 建独立环境Python 选 3.8 或 3.9。深度学习框架依赖多、包体大混用系统 Python 容易搞坏别的项目。conda create -n garbage python3.9之后一路pip install就行。这一步看似浪费时间但能省掉后续大量的环境翻车。3. 手写训练循环从图片目录到模型文件的三个关键步骤训练部分真正难的往往不是网络结构而是输入管道。数据加载和预处理只要有一点和推理端不一致模型效果就大打折扣。这个章节按“数据加载 → 迁移学习 → 评估保存”的顺序把训练循环的几个关键步骤拆开讲。3.1 数据加载与预处理用 ImageFolder 和 DataLoader 搭好输入管道PyTorch 的torchvision.datasets.ImageFolder能直接把目录结构变成数据集前提是目录长成data/train/类别名/图片.jpg的样子。它会把每个子文件夹名称自动映射成类别索引比如0对应可回收、1对应厨余。# data_loader.py import torch from torch.utils.data import DataLoader from torchvision import datasets, transforms # 训练集的预处理缩放 轻量增强 归一化 train_transforms transforms.Compose([ transforms.Resize((224, 224)), # 与预训练模型输入保持一致 transforms.RandomHorizontalFlip(), # 随机翻转增加泛化能力 transforms.ToTensor(), # 像素值从 0~255 转到 0~1 transforms.Normalize( mean[0.485, 0.456, 0.406], # ImageNet 的 mean std[0.229, 0.224, 0.225] # ImageNet 的 std ) ]) # 验证集不做随机增强只做缩放和归一化 val_transforms transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) train_dataset datasets.ImageFolder(data/train, transformtrain_transforms) val_dataset datasets.ImageFolder(data/val, transformval_transforms) train_loader DataLoader(train_dataset, batch_size32, shuffleTrue, num_workers2, pin_memoryTrue) val_loader DataLoader(val_dataset, batch_size32, shuffleFalse)这段代码的逻辑ImageFolder会扫描data/train下的每个子目录并建立类别映射训练时shuffleTrue打乱样本顺序验证时shuffleFalse保证评估结果可复现。参数说明里最值得调的是batch_size和num_workers。显存不够时把batch_size从 32 降到 16 或 8num_workers在 Windows 上不宜超过 2否则容易报 IO 错误。Normalize的 mean/std 不要乱改除非你重新在垃圾分类数据上统计过。Resize((224, 224))会直接拉伸图片如果担心变形可以改成Resize(256)CenterCrop(224)但那样预处理和推理端也要同步改。3.2 迁移学习加载预训练权重替换分类头训练循环本身不复杂真正决定效果的是“从什么起点开始学”。迁移学习的做法是加载在 ImageNet 上训练好的 ResNet18把最后一层全连接替换成自己的分类头然后只用较小的学习率微调整个网络。# train.py import torch import torch.nn as nn import torch.optim as optim from torchvision import models num_classes 4 # 可回收 / 厨余 / 有害 / 其他 # 加载预训练权重 model models.resnet18(weightsmodels.ResNet18_Weights.IMAGENET1K_V1) # 替换最后一层全连接输出类别数改为自己的 num_classes model.fc nn.Linear(model.fc.in_features, num_classes) device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device) criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lr1e-4) for epoch in range(30): model.train() running_loss 0.0 for images, labels in train_loader: images, labels images.to(device), labels.to(device) optimizer.zero_grad() # 清空梯度 outputs model(images) # 前向传播 loss criterion(outputs, labels) loss.backward() # 反向传播 optimizer.step() # 更新权重 running_loss loss.item() print(fepoch {epoch1}, loss: {running_loss/len(train_loader):.4f})这段代码的逻辑model.fc.in_features读取原全连接层的输入维度替换后的输出维度等于类别数。CrossEntropyLoss内部已经包含 softmax所以前向输出不需要手动再做 softmax。参数上lr1e-4是迁移学习的常用起点比从头训练小一个数量级避免破坏预训练特征。30个 epoch 对几千张图的数据集通常够用判断标准不是 epoch 固定值而是验证集准确率是否连续 5 轮不再上升。如果训练 loss 一直不降优先调学习率先试 1e-3 和 1e-5再考虑换优化器。3.3 验证与保存准确率该看哪一项模型该存哪种格式训练时打印的 loss 只反映训练集的拟合程度真正衡量模型好坏要看验证集准确率。每个 epoch 结束后跑一遍验证同时把模型权重保存下来。# 每个 epoch 结束后执行 model.eval() correct 0 total 0 with torch.no_grad(): for images, labels in val_loader: images, labels images.to(device), labels.to(device) outputs model(images) _, predicted torch.max(outputs, 1) # 取概率最大的类别 total labels.size(0) correct (predicted labels).sum().item() print(fval accuracy: {correct / total:.4f}) # 保存成 state_dict 格式文件小加载灵活 torch.save(model.state_dict(), garbage_model.pth)要注意这里没有把整个模型对象保存只保存了权重字典state_dict。加载时先创建相同结构的模型再load_state_dict否则会报结构不匹配。model.eval()必须写它会关闭 dropout 和 batch norm 的统计更新否则验证结果不稳定。只盯着总准确率也不够。四分类里如果“其他”类样本占一半模型全猜“其他”也能有 50% 准确率。建议额外打印每一类的准确率或者保存一份混淆矩阵这样才能看到“有害”类是不是被模型忽略了。4. 把模型包成可演示的垃圾分类系统推理封装与 Web 界面训练完拿到garbage_model.pth下一步是让不懂 Python 的人也能用。这段路的重点是推理封装和 HTTP 接口做得好不好直接决定答辩演示顺不顺利。4.1 推理脚本先让单张图片在命令行跑通不要一上来就写 Web 界面先在命令行把单张图片的预测跑通。推理脚本必须包含和训练端完全一致的预处理这一步能提前暴露绝大多数“模型没用”的误会。# predict.py import torch from PIL import Image from torchvision import transforms, models import torch.nn as nn # 类别顺序要和训练时 ImageFolder 的目录顺序一致 class_names [可回收, 厨余, 有害, 其他] def load_model(num_classes4, weight_pathgarbage_model.pth): model models.resnet18(weightsNone) model.fc nn.Linear(model.fc.in_features, num_classes) model.load_state_dict(torch.load(weight_path, map_locationcpu)) model.eval() return model def predict(model, image_path): transform transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) image Image.open(image_path).convert(RGB) tensor transform(image).unsqueeze(0) # 增加 batch 维度 with torch.no_grad(): outputs model(tensor) probs torch.softmax(outputs, dim1)[0] top_idx torch.argmax(probs).item() return class_names[top_idx], probs[top_idx].item() if __name__ __main__: model load_model() label, prob predict(model, test_photo.jpg) print(f预测类别: {label}, 置信度: {prob:.2f})逻辑说明map_locationcpu让权重在无 GPU 机器上也能加载convert(RGB)处理灰度图或带透明通道的 PNG避免通道数不一致unsqueeze(0)把 3 维张量变成 4 维因为模型期望的是(batch, channel, height, width)。这里有一个值得注意的参数细节class_names的顺序必须与训练时ImageFolder自动生成的类别索引完全一致。验证方法是打印训练集里的train_dataset.class_to_idx然后按这个字典顺序写class_names。顺序错了模型会把“可回收”显示成“厨余”。4.2 用 Flask 包一个 HTTP 接口上传图片返回 JSON命令行跑通后用 Flask 起一个服务。前端传图片上来后端调用predict函数返回 JSON 给页面展示。# app.py from flask import Flask, request, jsonify, render_template import os from predict import load_model, predict app Flask(__name__) app.config[MAX_CONTENT_LENGTH] 8 * 1024 * 1024 # 限制上传最大 8MB # 全局只加载一次模型不要在每次请求里重复加载 model load_model(weight_pathgarbage_model.pth) app.route(/) def index(): return render_template(index.html) app.route(/predict, methods[POST]) def api_predict(): if file not in request.files: return jsonify({error: 未找到上传文件}), 400 file request.files[file] if file.filename : return jsonify({error: 文件名为空}), 400 save_path uploads/tmp.jpg file.save(save_path) label, prob predict(model, save_path) return jsonify({label: label, confidence: prob}) if __name__ __main__: os.makedirs(uploads, exist_okTrue) app.run(host0.0.0.0, port5000, debugFalse)这段代码的关键在于模型只加载一次。如果写在api_predict里每个用户上传图片都会重新读一遍几百 MB 的权重并发一上来直接卡死。MAX_CONTENT_LENGTH设成 8MB可以防止有人传大图拖垮服务。上传图片后建议立即用predict里的Image.open做一次完整读取如果文件损坏predict会抛异常。想要更稳妥就在predict内部加try/except并返回错误信息。host0.0.0.0允许局域网内其他设备访问演示时手机和电脑连同一个 WiFi 就能测。4.3 前端页面与联调别把图片路径写死Flask 默认从templates目录找模板从static目录找静态资源。一个最简单的页面只需要一个表单、一个预览图和一个结果区域。!-- templates/index.html -- !DOCTYPE html html langzh head meta charsetUTF-8 title垃圾分类识别系统/title /head body h1上传一张垃圾照片/h1 form iduploadForm enctypemultipart/form-data input typefile namefile acceptimage/* required button typesubmit开始识别/button /form img idpreview width224 styledisplay:none; p idresult/p script const form document.getElementById(uploadForm); form.addEventListener(submit, async function (event) { event.preventDefault(); const fileInput form.querySelector(input[namefile]); const file fileInput.files[0]; if (!file) return; const preview document.getElementById(preview); preview.src URL.createObjectURL(file); preview.style.display block; const formData new FormData(); formData.append(file, file); const response await fetch(/predict, { method: POST, body: formData }); const data await response.json(); if (data.error) { document.getElementById(result).textContent 错误: data.error; } else { document.getElementById(result).textContent 类别: data.label , 置信度: (data.confidence * 100).toFixed(1) %; } }); /script /body /html前端逻辑不复杂但有个坑值得注意表单的enctypemultipart/form-data必须设置否则文件不会真正上传。页面里用fetch提交而不是传统表单跳转这样能在当前页面直接展示结果。联调时先打开http://127.0.0.1:5000用一张测试图片走完整流程。如果页面卡住不动先去 Flask 终端看有没有异常堆栈。浏览器控制台的 Network 面板也能看到是请求超时还是返回 500。提示debugTrue会开启自动重载但也会起一个额外进程如果项目里模型加载占内存较高建议生产演示时改成debugFalse。5. 复现这个项目时最容易踩的五个坑现象、原因、解决这一章是实操里最常见的五类翻车记录。每一条都按“现象 → 原因 → 解决”拆开对照着排查即可。5.1 图片读取失败FileNotFoundError 或 OpenCV 返回 None现象ImageFolder加载数据时报FileNotFoundError或者用 OpenCV 读图时cv2.imread返回None但文件明明存在。原因Windows 上路径分隔符写成了\导致转义错误目录或文件名包含中文、空格、括号还有一部分图片后缀是.jpg实际编码是 PNG 或 WebPOpenCV 读不出。解决统一用pathlib.Path处理路径不要手写字符串拼接。图片入库前做一次完整性检查用PIL.Image.open逐张打开并verify()损坏或格式不对的直接剔除。数据集路径保持纯英文目录结构固定为data/train/类别名/图片.jpg。5.2 训练 loss 不降或反复震荡现象训练了十几个 epochloss 始终在 2.0 附近不动或者上下剧烈波动验证集准确率也一直原地踏步。原因学习率不合适是第一位。迁移学习用太大学习率会破坏预训练权重从头训练用太小学习率又学不动。另外预处理里少了ToTensor()或Normalize会把 0~255 的像素值直接送进网络梯度计算异常。解决先用一张图过一遍正向和反向确认流程跑通。迁移学习设lr1e-4起步从头训练设1e-3起步。如果 loss 卡住不动先把batch_size调小再把归一化的 mean/std 核对一遍。数据增强过强也会导致 loss 不降可以把增强全部关掉试一次确定基准再逐步加回来。5.3 GPU 显存不足或内存溢出现象训练中途报CUDA out of memory或者 Web 演示时每上传一张图就卡几秒最后进程崩溃。原因batch_size太大输入图片分辨率太高num_workers开太多导致系统内存被占满Web 端每次请求都重新加载模型。解决训练时把batch_size减半直到能跑为止或者用torch.cuda.empty_cache()释放缓存。推理服务里把load_model放到全局作用域只加载一次。上传图片前先在前端压缩到 512px 以内再传能显著降低内存压力。5.4 加载权重时报 missing keys 或 unexpected keys现象load_state_dict报错提示缺了某些 key或者多了某些 key模型完全无法加载。原因最常见的是保存的是整个模型torch.save(model, ...)加载时却用load_state_dict或者训练时类别数是 4推理时类别数写成 10最后的全连接层维度不一致。解决统一约定保存model.state_dict()。加载前先打印state_dict的 keys确认结构一致。如果只改了分类头可以只加载前缀匹配的权重比如model.load_state_dict(torch.load(...), strictFalse)但这种做法要特别小心错误 key 会静默忽略最好还是保证结构完全一致。5.5 Flask 上传大图后一直转圈或返回 413现象点击识别后页面一直等待最后报413 Request Entity Too Large或者 Flask 终端显示一个超长请求处理时间。原因MAX_CONTENT_LENGTH没设置用户一次传了十几 MB 的高清照片Flask 和模型都要处理很久也可能是请求没有正确设置enctype后端收到的文件为空。解决设置app.config[MAX_CONTENT_LENGTH] 8 * 1024 * 1024前端在fetch前用 canvas 把图片压缩到最大 512px、质量 0.8再转成 blob 上传。另外predict函数的预处理里加Image.open(...).convert(RGB)避免格式问题。6. 把“能跑”变成“高分”调优顺序、验证与交付技巧项目跑通只是及格线“高分项目”和“能跑的 demo”之间的差距通常在三处数据质量、验证方式、交付文档。调优顺序我建议固定为先清洗数据再调训练策略最后才换网络结构。数据清洗的性价比最高比如删掉标签错误的图、把重复的相似图去重、补上少数类样本这比把 ResNet18 换成 ResNet50 提升更明显。第二个动作是给训练加学习率衰减用StepLR或CosineAnnealingLR让模型在后期稳步收敛。如果数据量连每类 500 张都不到换大模型只会更早过拟合不如保持小模型加增强。验证时不要只依赖测试集。测试集里的图片和训练集来自同一批拍摄环境模型很容易“背答案”。我的做法是准备 20 到 30 张在手机自然光下现拍的照片包含不同背景、不同角度让模型认一遍。这组照片不在任何训练流程里出现才是判断真实泛化能力的唯一标准。交付时把这几样东西补齐requirements.txt里写明 Python 和 PyTorch 的确切版本README 按“环境安装 → 数据准备 → 训练 → 推理 → 启动 Web”的顺序写每条命令都带上模型文件单独放models/目录不要和临时上传图片混在一起。代码里每次print的信息也要有含义比如打印当前 epoch、验证准确率、模型保存路径评审老师打开终端就能看到完整过程。我自己做这类项目的习惯是先花十分钟读完全部代码再动手把所有print和model.eval()位置标出来。这种习惯救过我很多次因为源码包里的代码未必都按顺序执行有的脚本是给 Colab 写的有的变量名和注释不一致照着盲跑最容易在环境阶段卡死。希望帮到你让这份高分项目真正变成你自己能讲清楚的东西。本文还有配套的精品资源点击获取
返回列表