首页 / 资讯中心 / 文章详情

真实废弃物分类数据集实战:4,800张图从训练到部署

真实废弃物分类数据集实战:4,800张图从训练到部署 ★ FEATURED ARTICLE
简介本资源为面向计算机视觉初学者与图像分类实践者的真实废弃物图像分类数据集覆盖纸板、食品有机物、玻璃、金属、杂项垃圾、纸张、塑料、纺织品垃圾和植被共9个类别适合用于分类网络训练、迁移学习验证及垃圾分类相关课程设计。数据已完成预处理可直接作为分类网络输入并已划分训练集与测试集各类别图片分目录存放便于快速构建实验流程。压缩包共2000个文件以1998张jpg图像为主体另含1个json标注文件与1个Python可视化脚本整体约155.99MB运行show脚本即可直观查看样本分布与类别效果。目前已有65人学习下载读者可借此省去数据采集与清洗成本将精力集中于模型结构改进、参数调优与结果对比同时结合标注文件理解类别映射关系为后续分割或分类任务提供可靠的数据基础。1. 真实废弃物分类数据集4,800 张标注图能跑出什么名堂拿到「生活中真实废弃物图像分类数据集」这个标题多数人第一反应是去找下载链接但真正决定项目成败的是标注质量与类别体系能不能对上你要落地的场景。这个数据集约 4,800 张、已完成标注属于中小规模图像分类数据集适合做垃圾分类识别、智能回收箱、环卫巡检等方向的模型验证与原型开发。它解决的核心问题是让你跳过最耗时的采集与标注环节直接进入模型训练和效果调优。适合谁想快速验证图像分类算法的新手、需要 baseline 做对比的算法工程师、以及做课程设计或产品 demo 的开发者。但别指望它直接产出生产级模型——4,800 张的体量决定了它更适合跑通流程、验证思路而不是追求 SOTA 精度。2. 废弃物图像分类的数据集拆解与模型选型2.1 先搞清楚 4,800 张图里到底有什么在动手写任何训练代码之前必须先把数据集的结构摸清楚。真实废弃物图像分类数据集通常按类别分文件夹存放目录结构类似dataset/train/plastic/、dataset/train/paper/这种形式。你需要确认三件事类别数量、每类样本数、图像尺寸分布。import os from collections import Counter from PIL import Image data_dir dataset/train class_counts {} size_dist Counter() for cls in sorted(os.listdir(data_dir)): cls_path os.path.join(data_dir, cls) if not os.path.isdir(cls_path): continue imgs [f for f in os.listdir(cls_path) if f.lower().endswith((.jpg, .png, .jpeg))] class_counts[cls] len(imgs) for img_name in imgs[:20]: # 每类抽样20张看尺寸 with Image.open(os.path.join(cls_path, img_name)) as im: size_dist[im.size] 1 print(类别分布, class_counts) print(尺寸分布抽样, size_dist.most_common(5)) total sum(class_counts.values()) print(f总图片数{total}) for cls, cnt in class_counts.items(): print(f {cls}: {cnt} ({cnt/total*100:.1f}%))这段脚本做的是类别均衡性检查。参数说明data_dir指向训练集根目录脚本会自动遍历子文件夹。重点看输出里的百分比——如果某个类别占比超过 40% 或低于 5%就存在类别不平衡问题后续训练需要加WeightedRandomSampler或做数据增强补偿。尺寸分布则决定你统一 resize 到多少常见做法是 224×224 或 256×256。2.2 模型选型从 ResNet 到 Transformer 的取舍2024 年之后图像分类模型的选择面比几年前宽了很多。对于 4,800 张这个量级我的建议是分三档模型参数量适用场景预期准确率5折ResNet-1811M快速 baselineCPU 也能跑82%88%EfficientNet-B05.3M精度与速度平衡85%91%ViT-B/16预训练86M数据增强充分时88%93%Swin-Tiny28M想要 Transformer 但显存有限87%92%选型的核心逻辑不是「哪个最新」而是「你的数据量和算力撑得住哪个」。4,800 张图直接训 ViT 从零开始基本会过拟合但如果用 ImageNet 预训练权重做迁移学习ViT 系列反而可能比 CNN 高 23 个点。我一般会先用 ResNet-18 跑一个 baseline确认数据 pipeline 没问题再换 EfficientNet 或 Swin 做精度提升。import torch import torchvision.models as models import torch.nn as nn def build_model(num_classes, archresnet18, pretrainedTrue): if arch resnet18: model models.resnet18(weightsmodels.ResNet18_Weights.DEFAULT if pretrained else None) model.fc nn.Linear(model.fc.in_features, num_classes) elif arch efficientnet_b0: model models.efficientnet_b0(weightsmodels.EfficientNet_B0_Weights.DEFAULT if pretrained else None) model.classifier[1] nn.Linear(model.classifier[1].in_features, num_classes) elif arch swin_t: model models.swin_t(weightsmodels.Swin_T_Weights.DEFAULT if pretrained else None) model.head nn.Linear(model.head.in_features, num_classes) else: raise ValueError(f不支持的架构: {arch}) return model # 假设有6个废弃物类别 model build_model(num_classes6, archefficientnet_b0) print(f可训练参数{sum(p.numel() for p in model.parameters() if p.requires_grad):,})关键点在于替换分类头ResNet 换fcEfficientNet 换classifier[1]Swin 换head。pretrainedTrue时加载 ImageNet 权重这是小数据集能训出可用模型的前提。参数量打印出来是为了确认你改对了层——如果可训练参数接近全量参数说明预训练权重没加载成功。2.3 数据增强策略别让 4,800 张图白瞎了4,800 张图做分类数据增强不是可选项而是必选项。但废弃物图像有个特殊性翻转和旋转通常安全颜色抖动要谨慎——塑料瓶被调成奇怪色调可能影响模型对材质的判断。from torchvision import transforms train_transform transforms.Compose([ transforms.RandomResizedCrop(224, scale(0.7, 1.0)), transforms.RandomHorizontalFlip(p0.5), transforms.RandomVerticalFlip(p0.2), transforms.RandomRotation(15), transforms.ColorJitter(brightness0.2, contrast0.2, saturation0.1, hue0.05), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]), transforms.RandomErasing(p0.1, scale(0.02, 0.1)), ]) 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]), ])参数说明RandomResizedCrop的scale(0.7, 1.0)表示随机裁剪原图 70%100% 的区域再缩放模拟不同拍摄距离。ColorJitter的hue0.05控制得很小就是为了避免颜色失真。RandomErasing模拟遮挡场景对废弃物识别很有用——实际场景中垃圾经常被部分遮挡。验证集只用 resize centercrop保证评估一致性。3. 从零跑通训练命令行、参数与监控3.1 训练脚本的核心结构一个能复现的训练脚本需要包含数据加载、模型构建、损失函数、优化器、学习率调度、训练循环、验证循环、模型保存。下面是一个精简但完整的版本。import torch from torch.utils.data import DataLoader from torchvision.datasets import ImageFolder from torch.optim import AdamW from torch.optim.lr_scheduler import CosineAnnealingLR import time def train_one_epoch(model, loader, criterion, optimizer, device): model.train() running_loss, correct, total 0.0, 0, 0 for imgs, labels in 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() * imgs.size(0) correct (outputs.argmax(1) labels).sum().item() total imgs.size(0) return running_loss / total, correct / total torch.no_grad() def evaluate(model, loader, criterion, device): model.eval() running_loss, correct, total 0.0, 0, 0 for imgs, labels in loader: imgs, labels imgs.to(device), labels.to(device) outputs model(imgs) loss criterion(outputs, labels) running_loss loss.item() * imgs.size(0) correct (outputs.argmax(1) labels).sum().item() total imgs.size(0) return running_loss / total, correct / total # 主流程 device torch.device(cuda if torch.cuda.is_available() else cpu) train_ds ImageFolder(dataset/train, transformtrain_transform) val_ds ImageFolder(dataset/val, transformval_transform) train_loader DataLoader(train_ds, batch_size32, shuffleTrue, num_workers4, pin_memoryTrue) val_loader DataLoader(val_ds, batch_size32, shuffleFalse, num_workers4, pin_memoryTrue) num_classes len(train_ds.classes) model build_model(num_classes, archefficientnet_b0).to(device) criterion torch.nn.CrossEntropyLoss(label_smoothing0.1) optimizer AdamW(model.parameters(), lr3e-4, weight_decay1e-4) scheduler CosineAnnealingLR(optimizer, T_max30) best_acc 0.0 for epoch in range(30): t0 time.time() tr_loss, tr_acc train_one_epoch(model, train_loader, criterion, optimizer, device) val_loss, val_acc evaluate(model, val_loader, criterion, device) scheduler.step() if val_acc best_acc: best_acc val_acc torch.save(model.state_dict(), best_model.pth) print(fEpoch {epoch1:02d} | train_loss{tr_loss:.4f} acc{tr_acc:.3f} | fval_loss{val_loss:.4f} acc{val_acc:.3f} | {time.time()-t0:.1f}s) print(f最佳验证准确率{best_acc:.4f})逻辑说明label_smoothing0.1缓解过拟合对小数据集效果明显。AdamW的lr3e-4是迁移学习的常用起点weight_decay1e-4做正则化。CosineAnnealingLR让学习率按余弦曲线下降T_max30对应总 epoch 数。每个 epoch 结束后比较验证准确率只保存最好的模型——这是防止过拟合的后悔药。3.2 训练过程中的关键监控指标光看 loss 和 accuracy 不够你还需要关注训练/验证 loss 的差距判断过拟合、每类准确率判断类别不平衡影响、学习率变化确认调度器生效。建议加一个简单的混淆矩阵输出。from sklearn.metrics import classification_report, confusion_matrix torch.no_grad() def detailed_eval(model, loader, device, class_names): model.eval() all_preds, all_labels [], [] for imgs, labels in loader: imgs imgs.to(device) preds model(imgs).argmax(1).cpu().numpy() all_preds.extend(preds) all_labels.extend(labels.numpy()) print(classification_report(all_labels, all_preds, target_namesclass_names)) print(混淆矩阵) print(confusion_matrix(all_labels, all_preds)) detailed_eval(model, val_loader, device, train_ds.classes)classification_report会输出每类的 precision、recall、f1-score。如果某个类别 recall 特别低说明模型对这个类别识别能力差可能需要补充该类样本或调整采样权重。混淆矩阵则告诉你哪些类别容易被搞混——比如「纸类」和「纸板」如果混淆严重说明类别定义本身可能有问题。3.3 用 TensorBoard 或 wandb 记录实验命令行打印只能看当前状态做对比实验需要可视化工具。最轻量的方案是 TensorBoard。from torch.utils.tensorboard import SummaryWriter writer SummaryWriter(runs/waste_classify_exp1) for epoch in range(30): tr_loss, tr_acc train_one_epoch(model, train_loader, criterion, optimizer, device) val_loss, val_acc evaluate(model, val_loader, criterion, device) writer.add_scalars(Loss, {train: tr_loss, val: val_loss}, epoch) writer.add_scalars(Accuracy, {train: tr_acc, val: val_acc}, epoch) writer.add_scalar(LR, optimizer.param_groups[0][lr], epoch) scheduler.step() writer.close()启动命令tensorboard --logdirruns。参数说明add_scalars把训练和验证曲线画在同一张图上方便对比。add_scalar(LR, ...)确认学习率按预期下降。如果 val loss 在某个 epoch 后开始上升而 train loss 继续下降就是过拟合信号需要提前停止或加强正则化。4. 避坑与排查废弃物分类训练中最容易翻车的 5 个点4.1 类别文件夹命名混乱导致标签错位现象训练时准确率始终在随机水平附近比如 6 分类一直在 16% 左右loss 不下降。原因ImageFolder按文件夹名排序生成标签如果文件夹命名有中文、空格或大小写不一致可能导致标签映射混乱。更隐蔽的情况是某些文件夹里混入了其他类别的图片。解决训练前先打印train_ds.classes和train_ds.class_to_idx确认类别列表和映射关系符合预期。再用脚本检查每个文件夹内是否有异常图片比如尺寸为 0 或无法打开的文件。for cls in train_ds.classes: cls_path os.path.join(dataset/train, cls) for f in os.listdir(cls_path): fp os.path.join(cls_path, f) try: with Image.open(fp) as im: im.verify() except Exception as e: print(f损坏文件: {fp} - {e})4.2 验证集与训练集分布不一致现象验证准确率远低于训练准确率且差距持续扩大。原因划分验证集时没有做分层采样导致某些类别在验证集中占比过高或过低。或者验证集的图片来自不同拍摄条件比如训练集是白底图验证集是实景图。解决用sklearn.model_selection.train_test_split的stratify参数做分层划分保证每个类别在训练集和验证集中的比例一致。如果数据集本身已经分好 train/val先检查两边的类别分布是否接近。from sklearn.model_selection import train_test_split import numpy as np all_files, all_labels [], [] for idx, cls in enumerate(sorted(os.listdir(dataset/train))): cls_path os.path.join(dataset/train, cls) for f in os.listdir(cls_path): all_files.append(os.path.join(cls_path, f)) all_labels.append(idx) X_train, X_val, y_train, y_val train_test_split( all_files, all_labels, test_size0.2, stratifyall_labels, random_state42 ) print(训练集类别分布, np.bincount(y_train)) print(验证集类别分布, np.bincount(y_val))4.3 图像尺寸不统一导致 DataLoader 报错现象训练时报RuntimeError: stack expects each tensor to be equal size。原因数据集中存在尺寸差异极大的图片而 transform 中没有做统一 resize或者某些图片是灰度图/ RGBA 四通道图。解决在 transform 中强制 resize 到固定尺寸并在数据集类中加一个转换步骤把灰度图和 RGBA 图统一转成 RGB。from PIL import Image class WasteDataset(torch.utils.data.Dataset): def __init__(self, file_list, labels, transformNone): self.file_list file_list self.labels labels self.transform transform def __len__(self): return len(self.file_list) def __getitem__(self, idx): img Image.open(self.file_list[idx]).convert(RGB) # 关键统一转RGB label self.labels[idx] if self.transform: img self.transform(img) return img, label4.4 学习率设太大导致 loss 震荡不收敛现象训练初期 loss 剧烈震荡甚至出现 NaN。原因迁移学习时用了从零训练的学习率比如 0.1而预训练模型需要更小的学习率来微调。解决迁移学习场景下AdamW 的 lr 建议从 1e-4 到 3e-4 开始试。如果 loss 仍然震荡降到 1e-5。另外可以加梯度裁剪。torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0)这行代码放在loss.backward()之后、optimizer.step()之前防止梯度爆炸。4.5 显存不够导致训练中断现象CUDA out of memory。原因batch_size 太大或者模型参数量超出显存容量。解决优先降 batch_size从 32 降到 16 或 8其次考虑混合精度训练。from torch.cuda.amp import autocast, GradScaler scaler GradScaler() for imgs, labels in train_loader: imgs, labels imgs.to(device), labels.to(device) optimizer.zero_grad() with autocast(): outputs model(imgs) loss criterion(outputs, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()混合精度训练能省 30%50% 显存对 4,800 张图这个量级完全够用。注意autocast只包前向传播反向传播用scaler处理。5. 把 4,800 张图用到极致进阶技巧与验证方法5.1 用交叉验证榨干小数据集的每一张图4,800 张图做单次划分验证集可能只有 8001,000 张评估结果波动大。5 折交叉验证能让每张图都参与验证一次得到更稳定的性能估计。from sklearn.model_selection import StratifiedKFold skf StratifiedKFold(n_splits5, shuffleTrue, random_state42) fold_results [] for fold, (train_idx, val_idx) in enumerate(skf.split(all_files, all_labels)): train_files [all_files[i] for i in train_idx] val_files [all_files[i] for i in val_idx] train_labels [all_labels[i] for i in train_idx] val_labels [all_labels[i] for i in val_idx] train_ds WasteDataset(train_files, train_labels, transformtrain_transform) val_ds WasteDataset(val_files, val_labels, transformval_transform) train_loader DataLoader(train_ds, batch_size32, shuffleTrue, num_workers4) val_loader DataLoader(val_ds, batch_size32, shuffleFalse, num_workers4) model build_model(num_classes, archefficientnet_b0).to(device) optimizer AdamW(model.parameters(), lr3e-4, weight_decay1e-4) scheduler CosineAnnealingLR(optimizer, T_max20) best_acc 0.0 for epoch in range(20): train_one_epoch(model, train_loader, criterion, optimizer, device) _, val_acc evaluate(model, val_loader, criterion, device) scheduler.step() best_acc max(best_acc, val_acc) fold_results.append(best_acc) print(fFold {fold1}: best_acc{best_acc:.4f}) print(f5折平均准确率{np.mean(fold_results):.4f} ± {np.std(fold_results):.4f})这个脚本会跑 5 次完整训练每次用不同的验证集划分。最终报告的是平均准确率和标准差。标准差如果超过 3 个点说明模型对数据划分敏感需要检查数据分布或增加正则化。5.2 用混淆矩阵定位「重灾区」类别交叉验证跑完后挑一个 fold 的模型输出混淆矩阵重点看哪些类别被系统性搞混。import matplotlib.pyplot as plt import seaborn as sns from sklearn.metrics import confusion_matrix cm confusion_matrix(all_labels_fold, all_preds_fold) plt.figure(figsize(8, 6)) sns.heatmap(cm, annotTrue, fmtd, cmapBlues, xticklabelstrain_ds.classes, yticklabelstrain_ds.classes) plt.xlabel(预测) plt.ylabel(真实) plt.title(废弃物分类混淆矩阵) plt.tight_layout() plt.savefig(confusion_matrix.png, dpi150)如果发现「塑料瓶」和「玻璃瓶」混淆严重说明模型对材质纹理的区分能力不足。解决办法针对这两个类别补充更多样本或者在数据增强中加强纹理相关的变换比如RandomGrayscale让模型更关注形状而非颜色。5.3 导出 ONNX 做推理验证训练完的模型最终要部署导出 ONNX 是验证模型可用性的关键一步。import torch.onnx model.eval() dummy_input torch.randn(1, 3, 224, 224).to(device) torch.onnx.export( model, dummy_input, waste_classify.onnx, input_names[input], output_names[output], dynamic_axes{input: {0: batch_size}, output: {0: batch_size}}, opset_version13 ) print(ONNX 导出完成) # 验证 ONNX 推理结果与 PyTorch 一致 import onnxruntime as ort import numpy as np sess ort.InferenceSession(waste_classify.onnx) test_img torch.randn(1, 3, 224, 224).numpy() onnx_out sess.run(None, {input: test_img})[0] torch_out model(torch.from_numpy(test_img).to(device)).cpu().detach().numpy() print(f最大差异{np.abs(onnx_out - torch_out).max():.6f})dynamic_axes让导出的模型支持可变 batch size部署时更灵活。最后对比 ONNX 和 PyTorch 的输出差异正常应该在 1e-5 以内。如果差异过大检查opset_version是否兼容你的推理环境。5.4 一个我踩过的坑别在验证集上调超参数早期做这个数据集时我习惯在验证集上试不同的学习率和 batch_size选验证集准确率最高的那组。结果模型在测试集上表现远低于预期——因为验证集被「偷看」了太多次已经失去了评估的客观性。正确做法是从训练集里再切一小块做验证集用于调参真正的测试集只在最后评估一次。如果数据量实在不够至少保证调参时看的是交叉验证的平均结果而不是单次划分的验证集。希望帮到你。本文还有配套的精品资源点击获取
阅读完成 · 觉得有帮助?
咨询建站