PyTorch垃圾图片分类数据集：ImageFolder加载与实战解析

简介面向深度学习图像分类任务这份资源提供了已做好划分的4类别垃圾图片数据集覆盖厨余垃圾、有害垃圾等常见大类别可直接用ImageFolder加载训练省去数据清洗与划分环节。包内共2000个文件主体为1998张jpg图像另含一个可视化Python脚本及配套json配置压缩包整体约591MB适合图像分类入门与实践项目。数据规模方面训练集39116张、测试集9777张总计近5万张图像分辨率在400至1000像素之间类别分布与尺寸差异能较好模拟真实场景。数据已按类别文件夹整理分别存放于train与test目录脚本无需修改即可运行随机传入一张图片即可绘制边界框并保存结果便于快速预览与验证。已有465人学习使用对于需要标准垃圾分类数据集或快速搭建分类流程的开发者具有直接参考价值。1. 一份到手就能跑的垃圾图片分类数据集省掉最麻烦的数据环节做图像分类时真正卡住进度的往往不是网络结构选型而是数据本身。这份垃圾图片分类数据集一共包含 4 个大类别、接近 4.9 万张图片数据总量 653 MB分辨率集中在 400 到 1000 之间全部是 RGB 彩色图。最难能可贵的是它已经按 train/test 完成了数据集划分train 有 39,116 张test 有 9,777 张比例大约是 4:1拿过来直接用 torchvision 的ImageFolder就能加载连random_split都省了。对于做深度学习入门、垃圾分类识别 demo或者只想在图片数据集上快速验证一个分类模型的从业者来说这套数据的价值在于把数据清洗和划分这两个最琐碎的环节直接抹掉了。下面从目录结构、类别分布、可视化脚本到训练链路逐一拆开看。2. 数据集目录结构与 ImageFolder 加载原理2.1 目录组织方式为什么按文件夹归档是分类数据的默认形态先看数据集解压之后的目录骨架readme 和可视化 py 脚本外层放数据全部集中在data目录下内部严格区分训练集与测试集。这种组织方式不是随手排的而是分类任务里最常规的归档惯例类别名作为目录名某一个类别的所有图片放进同名目录。这样一来有没有标签、标签是什么、标签对应的图片在哪全部通过目录路径就能回答不需要额外维护一份 CSV 映射表。data/ ├── train/ │ ├── kitchen_waste/ # 厨余垃圾 │ ├── harmful_waste/ # 有害垃圾 │ ├── recyclable/ # 可回收物 │ └── other_waste/ # 其他垃圾 └── test/ ├── kitchen_waste/ ├── harmful_waste/ ├── recyclable/ └── other_waste/这里类别名的具体拼写以压缩包内实际名字为准原理是固定的。torchvision.datasets.ImageFolder的加载逻辑就是扫描指定 root 下的一级子目录把子目录名解析为类别标签子目录下的每一张图片自动归类到对应标签。使用时要特别注意根路径不能选错ImageFolder(rootdata/train)与ImageFolder(rootdata/test)是两个独立的数据集对象二者共享相同的类别目录结构但数据互不重叠。2.2 ImageFolder 加载三行代码实例化训练集与测试集下面这段代码就是这份数据集的标准打开方式。训练集和测试集分别做一次实例化后续直接交给 DataLoader。import torch from torchvision import datasets, transforms # 预处理流水线统一尺寸 - 转张量 - 按 ImageNet 均值方差归一化 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]) ]) # 直接用 ImageFolder 加载不需要手动解析文件名或构造标签向量 train_ds datasets.ImageFolder(rootdata/train, transformtransform) test_ds datasets.ImageFolder(rootdata/test, transformtransform) print(f训练集类别数: {len(train_ds.classes)}) print(f训练集样本数: {len(train_ds)}) print(f测试集样本数: {len(test_ds)})逻辑说明ImageFolder先扫描子目录生成类别列表再把图片路径与类别索引一一配对。transform在每次取样本时同步执行Resize((224, 224))将原始分辨率 400 到 1000 的图片统一压缩为 224 见方ToTensor会把 HWC 的 0-255 像素值转为 CHW 的 0.0-1.0 浮点张量Normalize则按通道做标准化让输入分布落在激活函数敏感区间。这两步顺序不能颠倒先转张量再归一化是固定约定。参数说明中值得关注的是Resize((224, 224))虽然会改变原始宽高比但原始图片分辨率普遍在 400 以上压缩后目标对象依然占足够像素不会像小图拉伸那样产生严重畸变。若你的实验对比例敏感可以改用transforms.Resize(256)配合CenterCrop(224)不过这份数据集本身的图像内容相对居中直接 Resize 的精度损失可接受。2.3 类别映射机制class_to_idx 与样本配对关系ImageFolder实例化后会自动生成类别索引字典和样本列表这是模型训练时标签编码的依据所在。print(类别 - 索引映射:, train_ds.class_to_idx) # 输出示例: {kitchen_waste: 0, harmful_waste: 1, recyclable: 2, other_waste: 3} print(单个样本结构:, train_ds.samples[0]) # 输出示例: (data/train/kitchen_waste/22723.jpg, 0) print(样本总数:, len(train_ds.samples))逻辑说明train_ds.samples是一个列表每个元素为(图片绝对路径, 类别索引)的元组。ImageFolder内部依靠__getitem__按下标读取图片并应用 transform因此无需手写一个Dataset子类。class_to_idx的字典序取决于os.listdir的返回顺序不同操作系统下可能不同但这不影响训练因为索引与类别名的对应关系已经固定下来。提示如果解压后在 Windows 下运行发现样本数少了几个优先检查每个类别目录里是否有Thumbs.db一类的隐藏系统文件以 jpg 为过滤条件是最稳妥的做法。3. 数据集划分比例分析与类别均衡性验证3.1 4:1 划分意味着什么训练集 39,116 张、测试集 9,777 张两者之比约 4:1。这个比例在分类任务中属于相当标准的划分策略训练数据占 80%测试数据留出 20% 做最终评估。相比常见的 7:3 或 8:1:1这个划分在总量上保证了每个类别在训练阶段都有充足的样本同时留出足够的测试空间来评估泛化能力。数据集图片数量占比作用train39,11680.0%模型参数学习test9,77720.0%最终泛化评估合计48,893100%全量数据对垃圾图像分类这类类间相似度偏高的任务4:1 的划分在工程上是合理的。四个大类之间厨余、有害、可回收、其他存在部分外观交叉比如塑料瓶既是可回收物也可能被误分类为其他垃圾训练数据多一点能缓解边界样本带来的混淆。真正需要警觉的反而是类别间的均衡性如果某个类别的图片数量明显偏低模型会偏向样本量大的类。3.2 用脚本验证每个类别的图片数量拿到数据集后先别急着训练第一步要确认四个类别的图片数量是否均衡。下面这段脚本用标准库遍历目录统计 train 和 test 下每个类别的图片文件数不需要安装额外依赖。import os from collections import Counter def count_images_per_class(root_dir): 遍历一级子目录统计每个目录下的图片数量 counter Counter() for class_name in os.listdir(root_dir): class_path os.path.join(root_dir, class_name) if os.path.isdir(class_path): # 仅统计 jpg/jpeg 文件防止混入系统隐藏文件 num len([f for f in os.listdir(class_path) if f.lower().endswith((.jpg, .jpeg))]) counter[class_name] num return counter train_counts count_images_per_class(data/train) test_counts count_images_per_class(data/test) print(训练集各类别数量:, dict(train_counts)) print(测试集各类别数量:, dict(test_counts)) print(训练集合计:, sum(train_counts.values())) print(测试集合计:, sum(test_counts.values()))逻辑说明os.listdir只做目录内容遍历没法判断路径类型所以先用os.path.isdir过滤目录再用后缀名过滤图片。这段代码同时可以验证摘要中的总样本数是否吻合。Counter在此处的场景非常合适设定 key 为类别名value 直接累计图片数量。如果统计结果出现某个类别占比超过 40% 而其他类别不足 15%就要考虑类别均衡策略。常见做法是WeightedRandomSampler先根据样本数计算权重让少量类别获得更高的采样概率。但就这份数据集而言四个大类的分布相对平均直接在 DataLoader 中开启 shuffle 即可没有额外加权的必要。3.3 训练集内部还要不要再切一份验证集data/test只能用于最终评估不能作为训练过程中的调参依据。反复使用 test 集判断模型好坏会造成信息泄漏最终精度会失真。所以在data/train内部再划分出一部分作为验证集是常规操作。from torch.utils.data import random_split train_len int(len(train_ds) * 0.9) val_len len(train_ds) - train_len train_sub, val_sub random_split(train_ds, [train_len, val_len]) print(f子训练集: {len(train_sub)}, 验证集: {len(val_sub)})逻辑说明random_split按照给定长度比例切分数据集不需要动磁盘上的文件。子训练集用于梯度回传验证集用于每轮结束后的指标监测比如判断是否过拟合或决定何时降低学习率。9:1 的内部分割保留下来后最终测试仍然使用test_ds这样整个数据使用链路就是原始数据 - 90% 训练 / 10% 验证 - 独立测试集。提示random_split默认使用全局随机种子如果希望每次运行结果可复现在主程序开头加上torch.manual_seed(0)。4. 可视化脚本细节边界框绘制与数据抽查4.1 为什么分类数据集也要画边界框分类任务本身不需要边界框标注做监督信号但可视化脚本在这里解决的是「人的确认成本」。计算机把目录名当作标签人却需要眼睛核对图片内容是否真的属于那个类别。当图片分辨率为 400 到 1000 且目标物体占画面比例不一致时直接肉眼看原图很容易漏掉角落里的异常样本如果在图上用一个矩形框把目标区域圈出来抽查效率会高很多。该脚本的定位就是随机抽图、绘制边界框、保存到当前目录不依赖 GPU也不需要读取任何标注文件。4.2 随机抽图与矩形框绘制实现脚本的核心逻辑分三步遍历得到全部图片路径、随机取一张、用 PIL 绘制矩形框并保存。边界框的坐标在这里是人为给定的示意区域因为数据集本质是分类任务没有标签文件可读取。import glob import random from PIL import Image, ImageDraw # 递归收集 train 目录下所有 jpg 图片 image_paths glob.glob(data/train/*/*.jpg, recursiveTrue) if len(image_paths) 0: raise FileNotFoundError(没有找到任何图片请检查 data/train 路径) # 随机选一张模拟传入一张图片的入口 img_path random.choice(image_paths) img Image.open(img_path).convert(RGB) draw ImageDraw.Draw(img) w, h img.size # 以图片中心为锚点绘制占画面 80% 区域的矩形框 box (int(w * 0.1), int(h * 0.1), int(w * 0.9), int(h * 0.9)) draw.rectangle(box, outlinered, width4) # 保存到当前目录不覆盖原图 output_name visual_check.jpg img.save(output_name) print(f已保存: {output_name}, 原图来自: {img_path})逻辑说明glob.glob(data/train/*/*.jpg, recursiveTrue)匹配一级子目录下的全部 jpg这里递归标志实际不改变匹配行为因为通配符已经指定了两层目录。random.choice负责随机性每次运行都会挑出不同样本。ImageDraw.Draw是基于 PIL 的绘图句柄rectangle接收左上角和右下角坐标构成的对角线元组outline指边框颜色width控制线宽。边界框的参数调整思路如果图片中物体经常偏左或偏右可以改为水平居中、垂直拉伸的框比如(int(w*0.2), int(h*0.05), int(w*0.8), int(h*0.95))。不过该资源附带脚本的定位只是通用可视化工具不需要精确匹配目标位置。如果你的任务是识别目标类别后再框选那就不是这套脚本的工作范围了。4.3 脚本运行与输出效果验证命令行直接执行python visualize.py脚本会随机抽取一张训练图片在当前目录下生成visual_check.jpg同时控制台打印原始图片路径。打开输出图片确认三件事图片内容与目录名是否一致、边界框是否完整落在画面内、颜色是否正常。若图片偏暗或偏色确认原始采集设备与格式后再决定是否加入颜色校正但这不影响分类训练因为归一化阶段会削弱光照影响。如果脚本运行时报ModuleNotFoundError: No module named PIL执行pip install pillow即可。该脚本只依赖 Pillow 和标准库不需要安装 torch所以在没有 CUDA 的机器上也能跑。5. 预划分数据直接接入 PyTorch 训练链路5.1 DataLoader 参数配置与样本批处理ImageFolder实例化完成的是单样本读取实际训练还需要DataLoader将样本组装成 batch并在训练时做打乱、并行读取等调度。from torch.utils.data import DataLoader train_loader DataLoader(train_sub, batch_size32, shuffleTrue, num_workers4) val_loader DataLoader(val_sub, batch_size32, shuffleFalse, num_workers4) test_loader DataLoader(test_ds, batch_size32, shuffleFalse, num_workers4)逻辑说明train_sub来自前面random_split得到的子训练集shuffleTrue让每个 epoch 的梯度更新顺序不同避免模型记住样本排列验证集和测试集不需要打乱只做前向计算。num_workers控制子进程数量Windows 下建议设为 0 或 2因为多进程数据加载在 Windows 的 spawn 模式下会频繁拉起进程反而更慢Linux 下设为 CPU 核心数的一半即可。5.2 基于 ResNet18 的快速 baseline拿到预划分数据后快速搭建一个分类 baseline最直接的选择是 ImageNet 预训练的 ResNet18它能在参数量和精度之间取得平衡适合 4 分类任务在较小计算资源下跑通。import torch.nn as nn from torchvision import models model models.resnet18(weightsmodels.ResNet18_Weights.IMAGENET1K_V1) in_features model.fc.in_features num_classes len(train_ds.classes) model.fc nn.Linear(in_features, num_classes) print(f特征维度: {in_features}, 分类数: {num_classes})逻辑说明model.fc原本输出 1000 类在迁移学习场景下替换为 4 输出的全连接层前面的卷积特征提取层全部保留预训练权重。in_features从原 fc 层中取出保证替换后维度匹配。冻结骨干网络的requires_gradFalse策略只在样本量极小的情况下使用本数据集接近 4 万张训练图更适合全量微调。5.3 用独立测试集完成最终验证训练完成后在test_loader上做一次完整前向推理打印准确率这是判断整份数据可用性的直接方式。model.eval() correct, total 0, 0 with torch.no_grad(): for images, labels in test_loader: outputs model(images) _, predicted torch.max(outputs, dim1) total labels.size(0) correct (predicted labels).sum().item() print(f测试集准确率: {correct / total:.4f})逻辑说明model.eval()关闭 Dropout 和 BatchNorm 的统计更新确保推理结果稳定torch.no_grad()禁止梯度跟踪减少显存占用和计算开销。torch.max(outputs, dim1)返回每个样本在 4 个类别上的最大得分索引与真实标签做逐元素比较后累加计数最终准确率反映这份预划分数据在跨类泛化上的真实水平。本文还有配套的精品资源点击获取