ARTICLE DETAIL

资讯详情

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

基于深度学习的垃圾分类实战:从数据集构建到模型部署的完整指南

基于深度学习的垃圾分类实战:从数据集构建到模型部署的完整指南 简介这份资源是面向深度学习初学者、课程期末大作业与毕业设计需求者准备的垃圾分类实战项目包围绕图像识别与自动分类场景帮助读者理解从数据预处理、模型定义到应用部署的完整链路。压缩包共43个文件约55KB以17个Python脚本为核心辅以15个pyc编译文件、5个xml配置、3个ipynb实验笔记及说明文档涵盖模型构建、数据转换、工具函数与推理入口等模块目录结构清晰便于按功能检索与二次修改。目前已有45人学习下载。读者可从中获得一套可运行的垃圾分类方案骨架包括神经网络结构定义、数据集处理逻辑、图像预处理流程以及应用端交互程序同时可参考多份实验笔记与工具脚本快速复现训练与推理过程并在此基础上替换数据集或调整网络结构完成自己的课程设计或研究验证。1. 垃圾分类项目为什么总在“其他垃圾”上翻车“基于深度学习的垃圾分类.zip”这个标题很多人第一反应是不就是个图像分类吗拿 ResNet 跑一遍完事。但真做过就知道垃圾分类的难点从来不在模型结构而在数据本身——同一张外卖盒洗干净是“可回收物”沾着油渍就是“其他垃圾”一杯没喝完的奶茶杯身、盖子、液体分属三个类别。这类边界样本会让任何在干净数据集上刷到 99% 的模型一上真实场景就集体翻车。这个方向适合两类人一是想拿垃圾分类做深度学习入门实战的因为它数据获取门槛低、类别定义清晰、可视化直观二是要做毕设或课程项目的垃圾分类是个能讲清楚“数据预处理—模型选型—训练调参—部署推理”完整链路的载体。但前提是你得先把数据集的坑摸清楚而不是急着 import torch。2. 数据集构建从爬虫到标注的完整链路2.1 垃圾分类的数据从哪来公开数据集里TrashNet 是最常被引用的一个约 2500 张图分玻璃、纸、 cardboard、塑料、金属、trash 六类。但它的问题很明显背景干净、光照均匀、每类样本量不均衡纸类 594 张而 trash 只有 137 张。直接拿它训练模型学到的更多是“背景颜色”而不是“物体材质”。我的做法是公开数据集打底再用自己拍的和爬的补充。爬虫部分不展开敏感内容核心思路是用关键词在图片搜索引擎里批量下载然后人工过一遍。这里给一个用 Python 做图片去重和尺寸过滤的脚本这是数据清洗的第一步import os import hashlib from PIL import Image from tqdm import tqdm def file_hash(filepath): 计算文件的 MD5 值用于去重 with open(filepath, rb) as f: return hashlib.md5(f.read()).hexdigest() def clean_dataset(raw_dir, min_size224): 清洗数据集去重 过滤过小图片 raw_dir: 原始图片目录 min_size: 最小边长小于此值的图片会被删除 seen_hashes {} removed_dup 0 removed_small 0 for root, _, files in os.walk(raw_dir): for fname in tqdm(files): if not fname.lower().endswith((.jpg, .jpeg, .png)): continue fpath os.path.join(root, fname) # 去重 h file_hash(fpath) if h in seen_hashes: os.remove(fpath) removed_dup 1 continue seen_hashes[h] fpath # 过滤过小图片 try: with Image.open(fpath) as img: w, h_img img.size if min(w, h_img) min_size: os.remove(fpath) removed_small 1 except Exception: os.remove(fpath) print(f去重删除: {removed_dup}, 过小删除: {removed_small}) clean_dataset(./raw_images, min_size224)这段代码的逻辑很直接MD5 去重解决的是爬虫重复下载的问题尺寸过滤解决的是缩略图混入的问题。min_size224是因为主流 CNN 输入是 224×224小于这个尺寸的图片放大后会模糊反而引入噪声。实际跑下来爬来的 5000 张图通常能删掉 15%20%。2.2 标注格式的选择与转换垃圾分类的标注比目标检测简单因为一张图通常只有一个主体。但如果你要做“一张图里多个垃圾”的场景就需要目标检测格式。这里分两种情况说。单标签分类用文件夹结构就够了dataset/train/可回收物/xxx.jpg。但如果你用 LabelImg 或 Labelme 标了 XML/JSON就需要转成分类文件夹或 YOLO 格式。下面是把 Labelme 的 JSON 转成分类数据集的脚本import json import os import shutil def labelme_to_cls(json_dir, out_dir): 将 Labelme 标注的 JSON 转为分类数据集 每个 JSON 对应一张图取第一个 label 作为类别 os.makedirs(out_dir, exist_okTrue) for fname in os.listdir(json_dir): if not fname.endswith(.json): continue jpath os.path.join(json_dir, fname) with open(jpath, r, encodingutf-8) as f: data json.load(f) # 取第一个标注的 label 作为类别 if not data.get(shapes): continue label data[shapes][0][label] # 找到对应图片 img_name data[imagePath] img_src os.path.join(json_dir, img_name) if not os.path.exists(img_src): continue # 复制到类别文件夹 cls_dir os.path.join(out_dir, label) os.makedirs(cls_dir, exist_okTrue) shutil.copy(img_src, os.path.join(cls_dir, img_name)) print(转换完成) labelme_to_cls(./annotations, ./cls_dataset)参数说明json_dir是 Labelme 标注文件所在目录out_dir是输出的分类数据集根目录。这段代码只取第一个 label适合单主体场景。如果一张图有多个垃圾应该走目标检测路线用 YOLO 格式而不是分类格式。注意垃圾分类的类别定义各地标准不同上海分干垃圾/湿垃圾/可回收物/有害垃圾北京分厨余/可回收/有害/其他。做项目前先确认你的类别体系不然后面全要返工。3. 模型选型从 ResNet 到轻量级网络的取舍3.1 为什么我不推荐一上来就用 ViT2024 年了Transformer 在 CV 领域确实火但垃圾分类这个任务ViT 的性价比不高。原因有三第一ViT 需要大量数据预训练你的垃圾分类数据集通常只有几千到几万张从头训 ViT 基本训不动第二垃圾分类的判别特征偏局部材质纹理、颜色CNN 的归纳偏置反而更匹配第三部署时 ViT 的推理延迟和显存占用都比同精度的 CNN 高。我的建议是ResNet50 或 EfficientNet-B0 打底如果要做移动端部署换 MobileNetV3 或 ShuffleNetV2。下面是一个用 PyTorch 构建 ResNet50 分类器的完整代码import torch import torch.nn as nn from torchvision import models, transforms from torch.utils.data import DataLoader from torchvision.datasets import ImageFolder # 数据增强训练集用强增强验证集只做 resize 和归一化 train_transform transforms.Compose([ transforms.RandomResizedCrop(224, scale(0.7, 1.0)), transforms.RandomHorizontalFlip(), transforms.RandomRotation(15), transforms.ColorJitter(brightness0.2, contrast0.2, saturation0.2), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) val_transform transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) def build_model(num_classes4, pretrainedTrue): 构建 ResNet50 分类模型 num_classes: 垃圾分类类别数 pretrained: 是否加载 ImageNet 预训练权重 model models.resnet50(pretrainedpretrained) # 替换最后的全连接层 in_features model.fc.in_features model.fc nn.Sequential( nn.Dropout(0.3), nn.Linear(in_features, num_classes) ) return model # 数据集加载 train_dataset ImageFolder(./dataset/train, transformtrain_transform) val_dataset ImageFolder(./dataset/val, transformval_transform) train_loader DataLoader(train_dataset, batch_size32, shuffleTrue, num_workers4) val_loader DataLoader(val_dataset, batch_size32, shuffleFalse, num_workers4) model build_model(num_classeslen(train_dataset.classes)) print(f类别: {train_dataset.classes})关键参数说明RandomResizedCrop(224, scale(0.7, 1.0))的 scale 下限设 0.7 而不是 0.08是因为垃圾分类的物体通常占据画面主体裁得太狠会把物体裁没。ColorJitter的幅度控制在 0.2太大反而会让模型把颜色变化当成类别差异。Dropout(0.3)放在全连接前是为了缓解小数据集上的过拟合。3.2 训练策略与学习率调度垃圾分类数据集通常不大训练策略比模型结构更重要。我一般用 AdamW CosineAnnealing配合 warmup。下面是一个训练循环的核心代码import torch.optim as optim from torch.optim.lr_scheduler import CosineAnnealingLR def train(model, train_loader, val_loader, epochs50, lr1e-3, devicecuda): model model.to(device) criterion nn.CrossEntropyLoss(label_smoothing0.1) optimizer optim.AdamW(model.parameters(), lrlr, weight_decay1e-4) scheduler CosineAnnealingLR(optimizer, T_maxepochs, eta_min1e-6) best_acc 0.0 for epoch in range(epochs): model.train() running_loss 0.0 for imgs, labels in train_loader: imgs, labels imgs.to(device), labels.to(device) optimizer.zero_grad() outputs model(imgs) loss criterion(outputs, labels) loss.backward() optimizer.step() running_loss loss.item() scheduler.step() # 验证 model.eval() correct, total 0, 0 with torch.no_grad(): for imgs, labels in val_loader: imgs, labels imgs.to(device), labels.to(device) outputs model(imgs) _, pred outputs.max(1) correct (pred labels).sum().item() total labels.size(0) acc correct / total print(fEpoch {epoch1}/{epochs}, Loss: {running_loss/len(train_loader):.4f}, Val Acc: {acc:.4f}) if acc best_acc: best_acc acc torch.save(model.state_dict(), best_model.pth) return best_acclabel_smoothing0.1是垃圾分类任务里我必加的参数。因为边界样本的标签本身就有歧义硬标签会让模型过度自信。weight_decay1e-4配合 AdamW 是小数据集的常规操作。CosineAnnealingLR的eta_min1e-6保证学习率不会降到零最后几个 epoch 还能微调。提示如果你的类别不均衡比如“有害垃圾”样本很少把CrossEntropyLoss换成带权重的版本权重设为类别频率的倒数。4. 避坑与排查垃圾分类项目里最容易翻车的五件事4.1 验证集准确率很高一上真实场景就废现象验证集 98%拿手机拍几张垃圾去测错得离谱。原因验证集和训练集来自同一批数据背景、光照、拍摄角度都相似。模型学到的是数据集特有的偏差不是垃圾本身的特征。解决划验证集时按“拍摄批次”或“来源”划分而不是随机划分。比如公开数据集全部做训练自己拍的做验证。另外在验证集里故意加入一些“脏数据”——模糊的、暗光的、背景杂乱的看模型真实表现。4.2 模型把“背景”当成了分类依据现象把垃圾桶放在白色桌面上拍模型全预测成“可回收物”放在木桌上全预测成“其他垃圾”。原因训练数据里某个类别的背景高度一致模型走了捷径。解决数据增强里加RandomResizedCrop和ColorJitter是一方面更有效的是做背景替换——用分割模型把物体抠出来贴到随机背景上。如果没精力做分割至少保证每个类别的背景多样化。4.3 训练 loss 不下降或者下降后突然爆炸现象前几个 epoch loss 正常下降第 10 个 epoch 突然变成 NaN。原因学习率太大或者数据里有损坏图片导致梯度异常。解决先检查数据用前面的清洗脚本过一遍。然后加梯度裁剪torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm5.0)这行加在loss.backward()之后、optimizer.step()之前。max_norm5.0是经验值垃圾分类任务里通常够用。4.4 类别不均衡导致模型只预测多数类现象训练集里“其他垃圾”占 60%模型把所有输入都预测成“其他垃圾”准确率还有 60%。原因交叉熵损失被多数类主导。解决用加权交叉熵或 Focal Loss。加权交叉熵的权重设为1 / 类别频率归一化后传给CrossEntropyLoss(weightweights)。Focal Loss 更适合极端不均衡但需要调gamma参数一般从 2.0 开始试。4.5 部署时预处理和训练时不一致现象训练时验证集 95%部署到服务端或移动端后准确率掉到 70%。原因训练时用了RandomResizedCrop等增强部署时只做了 resize或者归一化的 mean/std 写错了。解决把验证集的 transform 单独抽成一个函数部署时严格复用。归一化参数用 ImageNet 的mean[0.485, 0.456, 0.406]std[0.229, 0.224, 0.225]不要自己改。如果用了自定义归一化训练和推理必须完全一致。5. 把模型塞进手机ONNX 导出与推理加速的一个关键技巧训练完的 PyTorch 模型要落地绕不开格式转换。我一般走 ONNX 路线因为兼容性好从服务端到移动端都能用。但导出时有个坑PyTorch 的model.eval()和torch.no_grad()只影响 dropout 和 batchnorm不影响导出计算图。真正影响 ONNX 推理结果的是动态轴设置和算子版本。下面是一个完整的导出脚本附带验证导出前后输出一致性的代码import torch import onnx import onnxruntime as ort import numpy as np def export_onnx(model, save_pathgarbage_cls.onnx, input_size(1, 3, 224, 224)): 导出 ONNX 模型并验证一致性 model.eval() dummy_input torch.randn(*input_size) torch.onnx.export( model, dummy_input, save_path, export_paramsTrue, opset_version11, do_constant_foldingTrue, input_names[input], output_names[output], dynamic_axes{input: {0: batch_size}, output: {0: batch_size}} ) # 验证 ONNX 模型 onnx_model onnx.load(save_path) onnx.checker.check_model(onnx_model) # 用 onnxruntime 推理和 PyTorch 输出对比 ort_session ort.InferenceSession(save_path) with torch.no_grad(): torch_out model(dummy_input).numpy() ort_inputs {ort_session.get_inputs()[0].name: dummy_input.numpy()} ort_out ort_session.run(None, ort_inputs)[0] np.testing.assert_allclose(torch_out, ort_out, rtol1e-3, atol1e-5) print(ONNX 导出验证通过) # 假设 model 已经加载了训练好的权重 model.load_state_dict(torch.load(best_model.pth, map_locationcpu)) export_onnx(model)opset_version11是我在移动端部署时最常用的版本兼容性好支持动态 batch。dynamic_axes把 batch 维度设为动态这样服务端可以一次推理多张图。do_constant_foldingTrue会把能合并的常量算子提前算掉减小模型体积。导出后移动端用 ONNX Runtime Mobile 或 NCNN 加载。如果追求极致速度可以把 ONNX 再转成 TensorRT服务端或 TFLite安卓端。但转之前一定先跑通 ONNX 推理确认输出和 PyTorch 一致否则后面出了问题你都不知道是哪一步引入的。一个我踩过的坑PyTorch 的RandomHorizontalFlip在训练时用了但垃圾分类里有些类别翻转后语义会变——比如“有害垃圾”标志上的箭头方向。如果你的类别里有方向敏感的物体把翻转增强关掉或者只对方向不敏感的类别开。最后说个习惯我每次训完模型都会拿 20 张自己手机拍的、没进过数据集的图跑一遍人工看错在哪。这个“土办法”比任何指标都管用因为真实场景的分布只有你的手机知道。希望帮到你。本文还有配套的精品资源点击获取
返回列表