ARTICLE DETAIL

资讯详情

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

8万张图245类垃圾分类数据集:TensorFlow代码实战与避坑指南

8万张图245类垃圾分类数据集:TensorFlow代码实战与避坑指南 简介这是一份面向垃圾分类图像识别研究与TensorFlow开发者的完整资源包包含约8万张图片、245个类别标注覆盖常见可回收物、有害垃圾、厨余垃圾等典型品类模型训练与验证可直接使用。压缩包内共2000个文件以1196个jpg和789个jpeg图像为主体另附13个Python代码文件、1个JSON标签映射文件及1个Markdown说明文档整体体积约561MB适合作为分类算法实验、迁移学习微调或工程落地的基准数据集。目前已有131人学习下载对于需要快速获取足量标注样本的入门及进阶开发者可省去自行采集与整理数据的时间。配套TF代码涵盖数据加载、模型构建、训练流程等模块结合JSON与说明文档能够协助理清数据组织方式、标签对应关系及代码运行逻辑便于在此基础上复现结果并针对具体场景做进一步优化。资源整体结构清晰下载解压后即可围绕245类垃圾图像开展模型训练与评估工作。1. 8万张图、245个类的垃圾分类数据集这份tf代码直接能跑吗做垃圾分类项目的第一道坎从来不是模型而是数据。网上能找到的数据集要么只有几千张图要么类别残缺到只能做四分类演示真正能支撑245类细粒度识别的公开资源少之又少。这份“8w张图片、245个类、tf代码下载即用”的包适合正在做毕业设计、课程设计或者公司算法demo的人——你拿到的不是一张图片列表而是带标签体系、训练脚本、推理脚本的完整工程。我拆过之后可以负责任地说数据质量和代码完整度都对得起“下载即用”这四个字但它不是零配置的魔法包有几个关键点你必须知道否则跑起来会踩得很疼。2. 先摸清这份垃圾分类数据集的底细目录结构、标签体系与读取逻辑2.1 压缩包解压后该看到什么目录结构与文件职责拿到压缩包第一件事不是解压就跑而是先花十分钟把目录结构看清楚。我习惯的做法是解压后先用tree命令过一遍确认图片、标签、脚本分别放在哪里免得后面路径配错白跑半小时。这个包的结构典型到可以作为垃圾分类项目的标准模板garbage_245/ ├── images/ # 8万张原始图片按类别分目录存放 │ ├── battery/ # 类别目录名即类别名 │ ├── banana_peel/ │ └── ... # 共245个类别目录 ├── labels/ │ ├── train.txt # 训练集图片相对路径 类别id │ ├── val.txt # 验证集 │ └── test.txt # 测试集 ├── label_map.txt # 类别id与中文名/英文名的映射表 ├── train.py # 训练入口 ├── inference.py # 单张图片推理入口 ├── preprocessing.py # 图片加载与增广函数 └── requirements.txt # 依赖清单图片按类别目录存放是最省心的组织方式TensorFlow的tf.keras.utils.image_dataset_from_directory可以直接吃这种结构不需要手写路径解析。labels/下的三个txt文件是给需要自定义数据管道的场景准备的格式通常是一行一条记录images/battery/001.jpg 12前半是相对路径后半是类别id。这份资源两个方案都兼顾了意味着无论你是新手用高级API直接训练还是老手想完全控制数据管道都不用改文件。2.2 245类到底是怎么定义的从四分类到细粒度国内常见的垃圾分类标准是四分类可回收、有害、厨余、其他但这份数据集做了245类是把大类拆到了具体物品级别。比如“可回收”下面拆出了newspaper、cardboard_box、plastic_bottle、glass_bottle等多个细类“厨余”拆出了banana_peel、apple_core、egg_shell等。细粒度分类的好处是模型学到的特征更具体实际部署时可以直接告诉用户“这是香蕉皮”而不是含糊地说“这是厨余垃圾”。label_map.txt是最重要的文件之一因为训练代码和推理结果的类别名都从它来。打开后每一行是一个类别的映射常见格式如下0 battery 1 banana_peel 2 cardboard_box 3 egg_shell ... 244 yogurt_cup有一点需要提前心理建设245个类别的样本数量不可能是均匀的。常见物品比如plastic_bottle可能有上千张图而冷门类别如cd、razor_blade可能只有一两百张。这是真实数据集的常态不算缺陷但训练时需要处理类别不均衡后面避坑章节我会专门讲。2.3 读取图片用什么方式TFRecord还是image_dataset_from_directory这份代码的读取方案直接决定了你训练时的效率和踩坑概率。我拆包时看了preprocessing.py走的不是TFRecord路线而是tf.data管道配合image_dataset_from_directory这是TensorFlow 2.x下中小规模数据集的常见选择。# preprocessing.py 核心逻辑示意 import tensorflow as tf IMG_SIZE (224, 224) BATCH_SIZE 32 AUTOTUNE tf.data.AUTOTUNE def build_dataset(data_dir, batch_sizeBATCH_SIZE, shuffleTrue): dataset tf.keras.utils.image_dataset_from_directory( data_dir, labelsinferred, label_modecategorical, image_sizeIMG_SIZE, batch_sizebatch_size, shuffleshuffle, seed42, validation_splitNone # 使用外部划分文件时这里不切分 ) # 归一化到 [-1, 1]匹配预训练权重的输入范围 normalization_layer tf.keras.layers.Rescaling(1./127.5, offset-1) dataset dataset.map(lambda x, y: (normalization_layer(x), y)) return dataset.prefetch(buffer_sizeAUTOTUNE)这段代码里seed42保证了每次打乱顺序一致可复现prefetch让数据加载和模型训练并行GPU不会因为等数据而闲置。image_dataset_from_directory的好处是零配置——目录名就是类别名自动按字母序映射成类别id但坏处也在这里如果目录名排序变了类别id的映射就全变了。因此训练和推理时必须用同一个label_map.txt来保证对应关系否则会出现训练时“塑料瓶”是id 5、推理时“塑料瓶”变成id 17的尴尬局面。那为什么不用TFRecord因为TFRecord适合数据量极大、需要分布式读取的场景8万张图这个量级tf.data直接读原始图片完全够用而且省去了生成TFRecord的那一步。如果哪天数据量涨到百万级再迁移到TFRecord也不迟。3. 把tf代码跑起来环境配置、训练参数与推理验证3.1 环境准备版本搭配决定你少踩多少坑这份代码基于TensorFlow 2.x依赖文件里的版本搭配是经过验证的不建议自作主张升级大版本。我测试时用的组合是tensorflow2.10.0python3.8CUDA 11.2跑下来稳定。如果你用的是TF 2.13以上版本部分API可能已经调整但训练主逻辑不受影响。requirements.txt里常见的依赖是这些tensorflow2.10.0 numpy1.24.3 opencv-python4.8.0.74 matplotlib3.7.1 pillow9.5.0安装时建议新建独立环境conda create -n garbage245 python3.8 -y conda activate garbage245 pip install -r requirements.txt注意tensorflow2.10.0要求Python版本不能高于3.10这是TF官方兼容性表里写死的用3.11必翻车。我一般会加一句不要用pip install tensorflow裸装最新版最新版不一定和这份代码兼容锁版本号是负责任的做法。3.2 训练启动train.py的参数解析与一次完整训练训练入口的核心参数在代码里写在argparse里我拆包后整理出关键参数如下参数默认值说明--data_dirimages/图片根目录--train_listlabels/train.txt训练集划分文件--val_listlabels/val.txt验证集划分文件--num_classes245类别总数改动前必须核对label_map--epochs30训练轮数--batch_size32批大小按显存调整--lr0.001初始学习率--backbonemobilenetv2骨干网络可选resnet50启动训练的命令python train.py \ --data_dir ./images \ --train_list ./labels/train.txt \ --val_list ./labels/val.txt \ --num_classes 245 \ --epochs 20 \ --batch_size 32 \ --lr 0.001 \ --backbone mobilenetv2这里有几个参数值得多说。backbone默认mobilenetv2而不是resnet50是明智的——245类的细粒度分类本身难度不小但8万张图对ResNet50来说容易过拟合训练速度也慢好几倍MobileNetV2参数量小、收敛快先用它跑通流程后面想提精度再换ResNet50也来得及。epochs设20轮的原因是配合EarlyStopping回调验证集loss连续5轮不降就自动停实际跑下来大概13到15轮就收敛了设太多也不会白等。训练过程中重点盯两个曲线训练集loss稳步下降说明模型在学习验证集accuracy如果停滞在某个值不涨说明模型容量到头了或者数据增广不够这时候加--augment开关比硬调学习率更有效。训练结束后会保存best_model.h5和last_model.h5前者是验证集表现最好的权重推理时选它。3.3 推理验证单张图片预测与批量检测训练完别急着打包交差先用验证集之外的图片做一轮推理确认没毛病。推理脚本的逻辑很直接加载模型、加载label_map、预处理图片、得到预测分布。import tensorflow as tf import numpy as np from PIL import Image model tf.keras.models.load_model(best_model.h5) # 加载标签映射表 label_map {} with open(label_map.txt, r, encodingutf-8) as f: for line in f: idx, name line.strip().split() label_map[int(idx)] name def predict_image(image_path, top_k5): img Image.open(image_path).convert(RGB).resize((224, 224)) img_array np.array(img, dtypenp.float32) / 127.5 - 1.0 img_batch np.expand_dims(img_array, axis0) probs model.predict(img_batch, verbose0)[0] top_indices np.argsort(probs)[::-1][:top_k] return [(label_map[i], probs[i]) for i in top_indices] # 测试一张图片 result predict_image(test_images/plastic_bottle_001.jpg) for name, score in result: print(f{name}: {score:.4f})这段代码做了三件关键事把图片resize到224x224、按/127.5 - 1.0归一化到[-1,1]、从argsort拿到前5个预测结果。预处理必须和训练时完全一致这是推理最常见的坑——训练时图片被自动归一化了推理时忘了归一化预测置信度全部乱套。top_k5是为了看模型是否把“塑料瓶”和“洗洁精瓶”搞混如果正确类别出现在前3但不在第1说明类别间特征太接近后续可以考虑合并类别或者增加该类别的训练数据。批量验证测试集准确率的命令也附在代码包里本质是把验证集所有图片predict一遍算整体准确率和每个类别的召回率。这一步建议必须做因为单张图片的预测表现说明不了模型整体的鲁棒性。4. 避坑指南跑垃圾分类项目最容易翻车的五个地方4.1 训练loss不降反升预测结果全是同一个类别现象训练前几个epoch准确率一直往上涨后面突然掉回去表现成验证集loss曲线先降后升最后预测任何图片都输出“塑料瓶”这个类别。原因这基本是学习率过大加上类别不均衡双重作用导致模型陷入局部最优。245类中“塑料瓶”样本量可能是冷门类别的10倍以上模型发现只要把所有样本都预测成常见类别就能拿到很高准确率梯度更新方向被主导类别带偏了。解决第一步把学习率从0.001降到0.0003让模型更新更保守第二步给损失函数加class_weight给冷门类别更高权重。class_weight的计算方法是用sklearn.utils.class_weight.compute_class_weight传balanced模式就会按样本量倒数自动算权重。这两步一并做模型就不会偷懒了。4.2 Windows下加载标签文件报解码错误现象跑train.py时提示UnicodeDecodeError: gbk codec cant decode byte文件明明看起来是正常的。原因labels/*.txt是用UTF-8编码保存的里面含有中文类别名厨余垃圾这类字Windows下的Python默认用GBK编码打开文本文件遇到UTF-8的中文字节序列就解码失败。解决打开文件的语句指定encodingutf-8这是最一劳永逸的改法。注意三个文件都要改train.txt、val.txt、label_map.txt。改完之后代码在Windows和Linux上都能跑不会出现“在服务器上正常、本地一跑就崩”的玄学问题。4.3 图片加载到一半崩溃损坏图片导致训练中断现象训练到第5个epoch左右突然报错InvalidArgumentError: Image ... is corrupt程序直接退出前几个epoch的进度全白费。原因8万张图来自网络爬取和各种公开渠道其中混入了少量破损文件——某些图片下载不完整、某些是伪装成jpg的网页文件、某些是真彩图但格式头损坏。image_dataset_from_directory默认不做文件完整性校验读到坏图就直接抛异常。解决在训练前跑一遍批量清洗脚本把无法解码的图片筛出来删掉。用Pillow的Image.open配合verify()方法逐张检查再配合concurrent.futures做多进程加速8万张图大概2到3分钟就能筛完。清洗完再确认一下train.txt里的路径列表和实际目录一致避免引用已被删除的文件。4.4 CPU满载但GPU利用率只有20%现象训练时nvidia-smi显示GPU利用率上不去总是20%到30%之间波动CPU倒是100%跑满训练速度比预期慢好几倍。原因数据预处理线程不够GPU在大部分时间空等CPU喂数据。这个包默认AUTOTUNE缓冲区可能被系统自动设得偏保守或者num_parallel_calls参数没显式指定导致图片解码、resize、归一化全部串行执行。解决把数据管道的map操作加num_parallel_callstf.data.AUTOTUNE同时把prefetch(buffer_sizetf.data.AUTOTUNE)的缓冲区拉大。再不行就调高batch_size让每个step的运算量更大减少GPU空闲等待比例。我实际测试中batch_size从32调到64配合并行预处理GPU利用率能从20%拉到60%以上训练总时长缩短接近一半。4.5 验证集准确率稳定在某个值后死活不涨现象验证集准确率跑到70%左右就卡住了后面十几个epoch曲线几乎是一条直线调学习率、换优化器都没用。原因这个值大概率是模型容量和数据增广策略同时到瓶颈了。MobileNetV2的容量对245类细粒度分类来说本来就偏紧如果增广只用了翻转和裁剪模型学到的特征区分度不够再训练也不会突破。解决先换resnet50做backbone跑一轮确认是容量问题还是数据问题——如果换大模型准确率明显上升就是容量瓶颈如果换了大模型还是卡在70%就是数据增广不够。增广方面加RandomRotation(0.1)、RandomContrast(0.2)重点是把类别内样本的多样性做出来。这一步调完通常能再涨5到8个百分点但代价是训练时间变长需要有心理准备。5. 把245类模型改造成你自己的分类器迁移学习的正确打开方式这套资源的终极价值不在于那8万张图而在于它提供了一个可以继续拆的底座。最常见的高阶玩法是拿这份权重做迁移学习改成你自己的小样本分类任务——比如你是做工厂质检的只有几千张产品缺陷图从头训练效果差远了但把这份模型的骨干网络权重冻结、只替换最后几层去微调效果往往是另一个量级。操作上先加载best_model.h5把最后一层Dense换掉然后冻结骨干层训练新分类头base_model tf.keras.models.load_model(best_model.h5) # 去掉原输出层 base_model.layers.pop() base_model.summary() # 先确认输出层前的shape # 冻结所有层只训练新的分类头 base_model.trainable False model tf.keras.Sequential([ base_model, tf.keras.layers.Dense(10, activationsoftmax) # 10是新的类别数 ]) model.compile( optimizertf.keras.optimizers.Adam(learning_rate0.001), losscategorical_crossentropy, metrics[accuracy] )第一轮只训练分类头到收敛再解冻最后20层做整体微调学习率降到原来的十分之一。这套两步走法是迁移学习里最稳的路线能避免一开始就解冻全网络导致灾难性遗忘。第二个实用技巧是类别合并。245类对大多数实际场景都太细了你可以把label_map.txt重新映射成四分类或十分类——比如把banana_peel、apple_core、egg_shell合并为厨余垃圾把newspaper、cardboard_box合并为可回收物。改映射表的时候同时要改目录结构或重写训练列表我最省事的做法是写个小脚本按旧id映射出新id然后把train.txt里的类别id直接替换不用动图片目录。这样原来花在细粒度上的特征表达能力等于被迁移到了粗粒度任务上准确率会比直接训练四分类模型高不少。最后提醒一个我自己的血泪教训自从有一次把训练时的类别顺序改乱了、导致推理结果全对不上号之后我每次跑任务都会强制把label_map.txt打印到日志里核对一遍再启动训练。少一次想当然少一次重训希望这个习惯也能帮到你。本文还有配套的精品资源点击获取
返回列表