免费获取学习方案
ARTICLE DETAIL

资讯详情

深耕编程基础知识与建站技术分享的一线实战洞察。

废物图像分类数据集与可视化脚本:从整理到YOLOv5训练实践

废物图像分类数据集与可视化脚本:从整理到YOLOv5训练实践 简介这套废物类别图像分类数据集专为垃圾分类识别与机器视觉入门设计已按训练集与验证集划分完毕可直接用于图像分类模型训练、评估也可作为yolov5分类任务的数据源。资源共2000个文件主体为1998张jpg图片另含1个类别字典json文件与1个Python可视化脚本压缩包约121.47MB。训练集含4200张图片、验证集含1800张图片覆盖Cardboard、Food_Waste、Glass、Metal、Other、Paper、Plastic共7个常见废物类别json字典可直接加载映射免去手工标注的麻烦。同时提供可视化脚本随机传入4张图片即可自动展示并保存在当前目录方便快速核验数据质量与类别分布。目前已有54人浏览学习适合开展垃圾分类、目标检测等视觉项目也可用于相关课程设计或毕业设计。1. 废物图像分类数据集与可视化脚本从数据整理到模型训练的一站式资源做机器视觉分类项目的人都有体会真正拖慢进度的不是网络结构调参而是数据准备、样本检查和验证集划分。这套废物类别图像分类数据集把 4200 张训练图和 1800 张验证图按 7 类整理好目录结构直接兼容 ImageFolder 和 YOLO 分类模式类别字典用 JSON 给出还带了一个随机抽取 4 张图片的可视化脚本。对需要快速验证图像分类算法、做垃圾分类识别预研或者想跑通 YOLOv5/YOLOv8 分类训练流程的工程师来说可以直接跳过最繁重的文件整理环节把精力放在模型和调参上。我拆过几套类似数据集坦白讲能同时提供标准目录结构和可视化工具的资源并不算多。2. 数据集目录结构与类别字典先看懂组织方式再动手训练2.1 train/val 目录与类别子目录的组织解压后的 data 目录下只保留 train 和 val 两个子目录分别存放训练集和验证集。每个子目录内部又按照类别名称建立二级目录同一类图片放在同一个文件夹中。这种「主目录/训练验证划分/类别子目录」的结构是 PyTorch 的 ImageFolder、Torchvision 的datasets.ImageFolder以及 YOLOv5 分类模块默认支持的目录格式不用额外写复杂的数据加载器。以这里的情况为例train 目录下会有 Cardboard、Food_Waste、Glass、Metal、Other、Paper、Plastic 这 7 个类别文件夹val 目录下同样维持这 7 个类别文件夹但图片数量不同。训练集总量 4200 张验证集 1800 张比例接近 7:3。对于普通的中小型机器视觉分类项目这个比例足够支撑模型训练不需要再额外划分测试集。下面是一个典型的目录树示意data/ ├── train/ │ ├── Cardboard/ │ ├── Food_Waste/ │ ├── Glass/ │ ├── Metal/ │ ├── Other/ │ ├── Paper/ │ └── Plastic/ └── val/ ├── Cardboard/ ├── Food_Waste/ ├── Glass/ ├── Metal/ ├── Other/ ├── Paper/ └── Plastic/需要特别注意的是类别目录名与 JSON 字典中的类别名必须完全一致大小写也要一致。YOLO 系列在读取分类目录时会直接拿目录名作为标签名一旦出现 Cardboard 和 cardboard 混用训练出来的模型类别顺序就会错位。我在处理类似数据集时第一件事就是检查两个目录下每个类别的图片数量确认有没有空目录。索引类别名称语义0Cardboard纸板1Food_Waste厨余垃圾2Glass玻璃3Metal金属4Other其他垃圾5Paper纸张6Plastic塑料实际使用时这份索引对应关系就是训练的监督信号。如果换用自定义网络softmax 输出维度设置为 7最后通过索引反向映射到类别名就可以输出完整的预测结果。2.2 类别字典 JSON 加载与映射逻辑数据集中附带一个类别字典文件内容是索引到类别名的映射。这种格式本质上是一条轻量的标签文件方便在训练脚本、评估脚本和可视化脚本之间共享。文件内容类似{ 0: Cardboard, 1: Food_Waste, 2: Glass, 3: Metal, 4: Other, 5: Paper, 6: Plastic }读取这个文件并在训练前打印出来是避免标签错乱的最直接方法。我一般用 Python 标准库json加载并逆转为类别名到索引的字典import json # 读取类别字典文件 with open(class_dict.json, r, encodingutf-8) as f: class_dict json.load(f) # 将字符串索引转为整数 class_to_idx {v: int(k) for k, v in class_dict.items()} idx_to_class {int(k): v for k, v in class_dict.items()} print(类别到索引的映射) for name, idx in class_to_idx.items(): print(f{name}: {idx})这段代码先把 JSON 中 key 为字符串的索引转换成整数再建立正反两个映射。正向映射用于在训练集加载时把文件夹名替换成标签数字反向映射则用在推理阶段将模型输出的数字索引还原为可读的类别名。这里要注意 JSON 默认读出的是字符串 key如果直接拿去和模型输出的 int 比较会得到错误结果所以先转换是关键一步。2.3 统计各类别样本量验证数据均衡性拿到数据集后下一步是统计每个类别的图片数量。这一步能暴露出类别不均衡、文件缺失、目录放错等问题。常见做法是直接用 Python 遍历目录并计数find data/train -type f -name *.jpg | cut -d/ -f3 | sort | uniq -c这个命令在 Linux 环境下可以快速统计data/train下每个子目录中的 jpg 文件数量。cut -d/ -f3提取的是三级路径中的类别名uniq -c做计数。如果某个类别数量明显偏少后续训练时需要考虑类别权重或者数据增强。对这套 7 分类数据集我建议至少确认每个类别的训练数量都不低于 400 张否则少数类很容易被多数类压制。3. 可视化脚本的实现逻辑随机抽 4 张图检查样本质量3.1 可视化脚本要解决什么问题标注好的图片并不是拿来就能直接用还得看照片里有没有拍错、模糊、被水印遮挡、光照异常或者与标签不符的情况。直接遍历几千张图片效率太低因此常见做法是写一个随机抽样脚本每次跑一遍都能看到不同的样本。这份资源附带的可视化脚本正是做这件事从不同类别中随机挑选图片拼成一张图展示并保存在当前目录下方便复盘。核心设计思路是随机性优先每次执行结果不同输出为图片文件方便在服务端或没有显示器的环境里查看代码短小依赖尽量少。通常只需要PIL和matplotlib两个库如果安装了opencv-python也可以用 cv2 实现。3.2 随机抽 4 张图片的 Python 实现假设脚本输入是一组图片路径或者一个数据集根目录需要从中随机选择 4 张并拼接显示。一个比较稳妥的实现方式如下import os import random import matplotlib.pyplot as plt from PIL import Image def visualize_random_images(data_dir, num_samples4, save_namevisualize.png): # 收集所有图片路径 all_images [] for root, dirs, files in os.walk(data_dir): for file in files: if file.lower().endswith((.jpg, .jpeg, .png)): all_images.append(os.path.join(root, file)) # 随机抽取 num_samples 张 if len(all_images) num_samples: raise ValueError(f图片总数不足 {num_samples} 张) sampled random.sample(all_images, num_samples) # 绘制子图 fig, axes plt.subplots(1, num_samples, figsize(15, 5)) for ax, img_path in zip(axes, sampled): try: img Image.open(img_path).convert(RGB) except Exception as e: print(f读取失败: {img_path}, 错误: {e}) continue ax.imshow(img) ax.axis(off) ax.set_title(os.path.basename(os.path.dirname(img_path))) plt.tight_layout() plt.savefig(save_name, dpi150) print(f可视化结果已保存到: {save_name}) if __name__ __main__: visualize_random_images(data/train)这段代码用os.walk递归遍历目录收集所有图片路径然后使用random.sample随机抽取 4 张。num_samples参数控制抽取数量save_name控制输出文件名。子图标题显示的是图片所在父目录名也就是类别名这样能快速确认该样本属于哪一类。convert(RGB)是为了避免灰度图或带透明通道的图片在显示时出现问题。参数说明data_dir需要传入指向数据集根目录的路径如果实际是 train 或 val 目录路径就对应传入data/trainnum_samples默认 4可以改为 9 或 16 来查看更多样本save_name默认保存为visualize.png如果不希望覆盖之前的输出可以在调用时加上时间戳或运行批次编号。3.3 参数调整与运行时机参数默认值建议调整场景num_samples4想一次看更多样本时设为 9 或 16save_namevisualize.png需要留存多轮结果时按时间命名data_dirdata/train检查验证集时改为 data/val运行脚本的时机也很重要。我通常在开始训练前跑一次用于确认数据格式训练结束后再跑一次结合模型输出做对照。如果发现某类图片明显模糊或方向错误就要考虑在预处理中增加筛选逻辑而不是盲目增强。4. 用 YOLOv5 分类模式直接训练命令、参数和遇到的坑4.1 将分类数据集转成 YOLOv5 能用的 YAMLYOLOv5 的分类训练脚本classify/train.py接受的数据集格式与 ImageFolder 一致根目录下包含 train 和 val 两个子目录每个子目录里是各个类别的图片文件夹。这一步已经满足不需要额外转换。但为了指定类别列表和路径需要写一个 YAML 配置文件内容如下# dataset.yaml path: ./data train: train val: val names: 0: Cardboard 1: Food_Waste 2: Glass 3: Metal 4: Other 5: Paper 6: Plasticpath是数据集根目录的相对路径或绝对路径train和val是相对于根目录的子目录名。names下的索引必须与类别字典保持一致YOLOv5 训练时会根据这个映射生成训练标签。这里容易踩坑的是names的 key 类型YAML 中写数字 0 和字符串 0 都合法但最好统一用数字避免与后续代码中的类型判断冲突。4.2 训练启动命令与关键参数说明进入 YOLOv5 代码目录激活对应虚拟环境后可以使用以下命令启动分类训练python classify/train.py \ --model yolov5s-cls.pt \ --data dataset.yaml \ --epochs 50 \ --batch-size 64 \ --img 224 \ --device 0这条命令会从官方权重yolov5s-cls.pt开始微调训练 50 个 epoch。--data指向刚才写好的 YAML 文件--batch-size根据 GPU 显存调整--img设置为 224 是分类任务中常用的输入尺寸。--device 0指定使用第一张 GPU如果使用 CPU 训练则改为--device cpu但速度会慢很多。参数表补充参数作用推荐值--model预训练权重或模型结构yolov5s-cls.pt / yolov5m-cls.pt--data数据集 YAML 文件路径dataset.yaml--epochs训练轮数50-100--batch-size批量大小32/64/128 视显存而定--img输入分辨率224 / 256--optimizer优化器Adam / SGD--lr学习率0.001 或默认训练结束后结果保存在runs/train-cls/exp*目录下包含训练过程中的损失曲线、准确率曲线、混淆矩阵以及类别预测样例。YOLOv5 分类模式没有单独的测试集参数验证集结果直接由--val控制如果没有单独指定默认使用 YAML 中val对应的目录。4.3 训练过程中常见的坑第一目录名与 YAML 的names不一致。比如数据集里文件夹叫 Food_Waste而 YAML 里写成了 Foodwaste训练时会报无法找到类别目录。检查方式是启动训练后看控制台输出的类别列表如果发现缺少某个类别立即终止训练。第二图片格式混用。YOLOv5 分类训练默认支持 jpg、png、bmp但混用时有可能出现读取失败。建议在用可视化脚本检查时留意图片后缀必要时统一转换为 jpg。第三类别不平衡。7 个类中如果 Cardboard 或 Other 数量明显偏多模型会偏向多数类。建议先看 2.3 节的统计结果如果发现不均衡可以在classify/train.py中启用标签平滑或加大--weight参数。第四验证集和训练集中出现同一张图。这种泄漏会让模型准确率虚高部署后表现严重下降。检查方法在下一章展开。5. 更进一步用可视化脚本做验证集泄漏检查与阈值调试5.1 泄漏检查思路卷积神经网络对重复图片非常敏感。如果训练集和验证集中出现了同一张图片模型会把这张图的特征强行记住最终验证准确率可能高达 99%但遇到真实数据时准确率骤降。对于这份废物图像数据集训练集和验证集在目录上已经分开但文件名是序列号无法直接判断是否有重叠。可视化脚本在这里可以当作人眼抽查工具交替传入 train 和 val 目录把两张可视化结果图拼在一起人工比对。更可靠的办法是对所有图片计算感知哈希或文件级哈希这里推荐用 Python 快速实现import os import hashlib from collections import defaultdict def compute_file_hash(path): hasher hashlib.md5() with open(path, rb) as f: for chunk in iter(lambda: f.read(4096), b): hasher.update(chunk) return hasher.hexdigest() train_dir data/train val_dir data/val hash_map defaultdict(list) for base_dir in [train_dir, val_dir]: for root, _, files in os.walk(base_dir): for file in files: if file.lower().endswith((.jpg, .jpeg, .png)): file_path os.path.join(root, file) h compute_file_hash(file_path) hash_map[h].append(file_path) duplicates {h: paths for h, paths in hash_map.items() if len(paths) 1 and any(p.startswith(train_dir) for p in paths) and any(p.startswith(val_dir) for p in paths)} print(f疑似跨集重复图片组数: {len(duplicates)})这段代码先遍历 train 和 val 目录计算每个文件的 MD5 值。如果同一个哈希值同时出现在 train 和 val 中就说明存在跨集重复。MD5 对图片内容完全相同的情况是有效的即使文件被重命名也能识别。如果发现重复优先将对应的验证集图片移出或直接从数据集中剔除。5.2 一个可复用的小技巧把随机抽样改成按类别抽样可视化脚本的随机抽样有一个盲区它不保证每次都能覆盖全部 7 个类别。做模型调试时我更倾向于改成按类别分别抽样每类取 1 张最终拼成 1 行 7 列的图。这样一眼就能看出每个类别的图片质量也能观察类别内的差异。修改方式非常简单只需要把visualize_random_images函数改为接收类别列表并在每个类别目录里单独随机挑选def visualize_by_class(data_dir, class_names, save_namevisualize_by_class.png): fig, axes plt.subplots(1, len(class_names), figsize(20, 5)) for ax, cls in zip(axes, class_names): cls_dir os.path.join(data_dir, cls) images [os.path.join(cls_dir, f) for f in os.listdir(cls_dir) if f.lower().endswith((.jpg, .jpeg, .png))] sample_path random.choice(images) img Image.open(sample_path).convert(RGB) ax.imshow(img) ax.axis(off) ax.set_title(cls) plt.savefig(save_name, dpi150) print(f按类别可视化结果已保存: {save_name})这种方式比纯随机抽样更适合训练前检查。调用时传入class_names list(class_to_idx.keys())即可自动覆盖所有 7 个类别。最后一个调试技巧是使用可视化脚本输出验证模型的分类置信度。推理阶段将脚本稍作修改把ax.set_title(cls)改为ax.set_title(f{cls}: {confidence:.2f})就能直观看到模型对每个类别样本的置信度分布辅助确定合理的分类阈值。这种方法在机器视觉工程落地时非常实用尤其当某些类别之间视觉特征接近时阈值微调能显著降低误检率。本文还有配套的精品资源点击获取
返回列表