ARTICLE DETAIL

资讯详情

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

36类果蔬图像分类数据集详解:从PyTorch加载到模型训练实践

36类果蔬图像分类数据集详解:从PyTorch加载到模型训练实践 简介这是一份面向图像分类任务的水果蔬菜数据集涵盖香蕉、苹果、梨、葡萄、橙子、猕猴桃等36个常见类别总计约3400张已标注图像。数据已完成预处理可直接作为分类网络输入方便学习者跳过采集与标注环节集中验证模型结构或训练技巧。包体共2000个文件以JPG格式图像为主1998张另附1个Python可视化脚本和1个JSON标签文件整体大小约94.47MB。数据已划分训练集与验证集各同类图像集中存放JSON中提供完整类别清单运行附带脚本即可直观查看数据集样例便于检查类别平衡与图像质量。目前已有213人学习浏览比较适合计算机视觉初学者、课程设计以及果蔬识别相关算法对比实验的快速起步。1. 36 类果蔬图像分类数据集一包能直接进 DataLoader 的已标注数据做图像分类最难拿的往往不是模型而是一份干净、已标注、能直接丢进 DataLoader 的数据集。这份 36 类果蔬图像分类数据集一共有约 3400 张图片类别覆盖香蕉、苹果、梨、葡萄、橙子、猕猴桃、西瓜、石榴、菠萝、芒果等 36 种常见果蔬带 json 标注已经划分好训练集和验证集内置 show 脚本可以预览。它的价值在于省掉了标注和清洗这两个最耗时的环节不用自己爬图、不用写标注工具拿到手就能当分类网络输入。如果你正在做图像分类课程设计或者拿公开模型跑果蔬识别项目这份数据适合用来做 baseline 和算法对比如果你想搞明白一套标注数据集是怎么组织的它的 json 结构和目录划分同样值得拆开看一遍。2. 数据集结构拆解json 类别映射、目录划分与 show 脚本用法这个包的核心不是图片本身而是它替你完成了“标注 划分 可预览”这三步。先看标签文件再看目录结构最后跑一遍可视化脚本确认数据没毛病再谈训练。很多人拿到压缩包就急着写模型结果类别编号对不上、目录缺文件夹返工两次才老实——数据检查这一步省不掉。2.1 从 json 读类别清单36 类的命名与编号规则解压后第一步不是急着数图片而是先打开 json 文件。里面记录着 36 个类别和对应的类别索引这是整个数据集的唯一标签依据。文件结构大概是类别名作为 key类别索引作为 value按 0 到 35 递增。类别里既有香蕉、苹果、梨、葡萄、橙子、猕猴桃、西瓜、石榴、菠萝、芒果这类水果也有黄瓜、胡萝卜、辣椒、洋葱、马铃薯、番茄、萝卜、甜菜根、卷心菜、生菜、菠菜、大豆、花椰菜、甜椒、玉米、甘薯、姜、大蒜、豌豆、茄子这类蔬菜还有玉米和甜玉米、辣椒和甜椒、萝卜和甜菜根这种外观接近的“难分对”后续训练时混淆矩阵里要重点盯这几组。import json with open(labels.json, r, encodingutf-8) as f: label_map json.load(f) print(类别数量:, len(label_map)) # 36 for name, idx in label_map.items(): print(idx, name)这段代码把 json 读成一个 dictkey 是类别名value 是整数索引。输出类别数量和每一类的编号。这里要记住一点json 里的 key 顺序和数字索引不一定跟目录排序一致后面做 Dataset 时要以这个 dict 为准别自己重排。有些类名本身带中文歧义比如辣椒和辣椒粉是两种状态甜玉米和玉米是两种作物json 里都会区分成独立类别。如果发现训练时某两类反复互相误判大概率就是它们。建议把 json 打印出来的结果直接保存一份后续训练、评估、推理都用同一份映射。2.2 目录结构train/val 划分与同类图片存放数据集已经把图片划分成训练集和验证集同类图片放在同一个子目录下整体结构是这样dataset/ ├── labels.json ├── show.py ├── train/ │ ├── banana/ │ │ ├── Image_1.jpg │ │ ├── Image_7.jpg │ │ └── ... │ ├── apple/ │ └── ... └── val/ ├── banana/ └── ...train 和 val 目录下各有 36 个类别子目录子目录名就是类别名。这个设计有两个好处一是直接用 torchvision 的 ImageFolder 就能读二是可以随时用系统命令统计各类样本数不用解析标注文件。约 3400 张图按 36 类平均下来每类不到 100 张属于典型的小样本分类任务后面配置网络时不要一开始就上太重的主干模型。import os from collections import Counter def count_samples(root): counts Counter() for cls in os.listdir(root): cls_dir os.path.join(root, cls) if os.path.isdir(cls_dir): counts[cls] len(os.listdir(cls_dir)) return counts train_counts count_samples(train) val_counts count_samples(val) print(训练集类别数:, len(train_counts)) print(验证集类别数:, len(val_counts)) print(单类最多:, max(train_counts.values())) print(单类最少:, min(train_counts.values()))这段脚本遍历 train 目录下的所有类文件夹用 Counter 统计每类图片数量。打印训练集和验证集的类别数以及最多最少类的样本量用来判断是否存在明显的类别不均匀。如果某类只有二三十张而另一类有一百多张训练时就要考虑采样策略。提示train 和 val 的类别子目录名必须完全一致少一个类或者名字拼错训练时会直接报 KeyError。2.3 可视化脚本 show先看数据再谈训练资源里带了 show 脚本直接运行就能预览数据集作用是快速检查图片是否损坏、类别标签和图片内容是否对得上。我一般会自己写一段更短的预览脚本每类取一张图拼成 6×6 网格扫一眼就能发现问题。import os import matplotlib.pyplot as plt from PIL import Image root train classes sorted([d for d in os.listdir(root) if os.path.isdir(os.path.join(root, d))]) fig, axes plt.subplots(6, 6, figsize(16, 16)) for i, cls in enumerate(classes[:36]): cls_dir os.path.join(root, cls) img_name sorted(os.listdir(cls_dir))[0] img Image.open(os.path.join(cls_dir, img_name)) ax axes[i // 6][i % 6] ax.imshow(img) ax.set_title(cls, fontsize10) ax.axis(off) plt.tight_layout() plt.savefig(preview.png, dpi150) plt.show()plt.subplots(6, 6) 创建 36 个子图正好对应 36 个类别i // 6 计算行号i % 6 计算列号把每一类第一张图填进对应位置。标题直接写类别名一眼能看出标签是否错位、分类是否合理。跑完这步基本可以确认数据可用再进训练环节。顺便建议做一次图片完整性校验避免训练时读到损坏文件导致 epoch 中断。这类问题在数据集里不常出现但一旦出现排查成本比提前跑一遍校验高得多。from PIL import Image bad [] for cls in classes: for f in os.listdir(os.path.join(root, cls)): p os.path.join(root, cls, f) try: Image.open(p).verify() except Exception: bad.append(p) print(损坏图片数量:, len(bad)) if bad: print(bad[:5])Image.open().verify() 只检查文件头不会真的解码整张图速度很快。返回异常的文件路径会收集到 bad 列表。如果发现损坏文件优先从数据集中剔除而不是等训练中途崩了再回头查。这一步在 Windows 和 Linux 上行为一致放心跑。3. 接进 PyTorch 训练流程自定义 Dataset、预处理参数与 DataLoader 配置数据看完了接下来把它接进训练代码。这里不走捷径直接用 ImageFolder而是自己写一个 Dataset。原因不复杂ImageFolder 虽然省事但它会按目录名的字典序自动生成类别索引这个顺序和 json 里的编号不一定一致。如果训练用 ImageFolder 的编号验证时又用 json 的编号混淆矩阵会乱到没法看。3.1 为什么自定义 Dataset 而不是直接用 ImageFolder自定义 Dataset 的核心思想很简单把 json 里的 label_map 作为标签唯一来源遍历每个类别的目录把图片路径和对应索引存成列表。这样训练和推理用同一份映射不会出现标签错位。import os import json from PIL import Image from torch.utils.data import Dataset class FruitVegDataset(Dataset): def __init__(self, root, label_map, transformNone): self.samples [] self.transform transform for cls, idx in label_map.items(): cls_dir os.path.join(root, cls) if not os.path.isdir(cls_dir): # 该类别在 train/val 中不存在就跳过避免路径错误 continue for fname in sorted(os.listdir(cls_dir)): self.samples.append((os.path.join(cls_dir, fname), idx)) def __len__(self): return len(self.samples) def __getitem__(self, index): path, label self.samples[index] img Image.open(path).convert(RGB) if self.transform: img self.transform(img) return img, label这份 Dataset 在初始化时一次性把图片路径和标签读进内存列表getitem按索引返回图像和标签。有几个设计点一是 label_map 直接来自 json标签编号完全可控二是类别目录不存在的容错train 和 val 子目录不完全一致时不会直接崩溃三是在初始化时就把文件列表准备好训练时不用反复遍历目录。初始化之后Dataset 里 sample 的量就是实际参与训练的图片数。这里有个容易被忽略的细节label_map 中的 key 顺序和 os.listdir 返回的顺序无关因为代码是显式遍历 label_map.items()所以 json 里第 5 个类别是西瓜训练样本里的标签就是 4不会有歧义。3.2 预处理参数与 DataLoader 配置预处理要遵循 ImageNet 预训练权重的惯例尺寸、归一化参数都必须跟预训练时保持一致否则加载预训练权重后特征分布对不上。from torch.utils.data import DataLoader from torchvision import transforms train_tf transforms.Compose([ transforms.Resize((224, 224)), transforms.RandomHorizontalFlip(p0.5), 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_tf transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]), ]) with open(labels.json, r, encodingutf-8) as f: label_map json.load(f) train_ds FruitVegDataset(train, label_map, train_tf) val_ds FruitVegDataset(val, label_map, val_tf) train_loader DataLoader(train_ds, batch_size32, shuffleTrue, num_workers2) val_loader DataLoader(val_ds, batch_size32, shuffleFalse, num_workers2)Resize((224, 224)) 直接拉伸会改变长宽比但 224 是 ImageNet 预训练模型最常见的输入尺寸。如果果蔬图像本身不是正方形可以改成 Resize((256, 256)) 再配 CenterCrop((224, 224))减少形变对识别的影响。Normalize 里的 0.485/0.456/0.406 是 ImageNet 统计值使用预训练权重时不要改。验证集的变换里没有随机增强只有确定性操作这样评估结果稳定可复现。3.3 参数调优的几个细节batch_size 在 36 类、每类平均不到 100 张的情况下32 是比较稳妥的选择64 也可以梯度更稳但每轮更新次数变少。num_workers 在 Linux 下可以设到 4Windows 下建议 0 或 1否则容易被多进程数据加载的问题卡住。使用 GPU 训练时把 pin_memory 打开能省去一部分 CPU 到 GPU 的数据拷贝时间。train_loader DataLoader( train_ds, batch_size32, shuffleTrue, num_workers2, pin_memoryTrue, drop_lastTrue )drop_lastTrue 在训练集上可以开它把最后不足一个 batch 的样本丢弃对 loss 曲线平滑有一点帮助但验证集不要开否则会漏掉尾部样本导致评估不完整。还有一个很多人忽略的点验证精度要算整个 val 目录的准确率而不是拿最后一个 batch 的准确率糊弄过去。下面这个评估函数是后续所有实验的基础建议直接沿用保证不同模型之间的对比口径一致。torch.no_grad() def evaluate(model, loader, device): model.eval() correct, total 0, 0 for imgs, labels in loader: imgs, labels imgs.to(device), labels.to(device) logits model(imgs) pred logits.argmax(dim1) correct pred.eq(labels).sum().item() total labels.size(0) return correct / totaltorch.no_grad() 关闭梯度计算推理时省显存也提速度。argmax(dim1) 取每个样本概率最大的类别索引pred.eq(labels) 返回布尔张量sum().item() 统计正确数。函数返回的是整体准确率而不是分批平均——小数据集上整体准确率才是最可靠的评估口径。4. 启动训练模型选型、训练脚本与混淆矩阵误差分析数据流程通了接下来是训练本身。模型选型上这份数据集的规模摆在那里3400 张图分 36 类模型太大容易过拟合模型太小精度上不去。合理的做法是拿一个轻量网络先跑通 baseline再根据混淆矩阵去改进。4.1 模型选型ResNet18 作为起点ResNet18 参数量约 11M在这种规模的数据集上训练稳定收敛速度也快非常适合做第一版 baseline。MobileNetV2 更轻适合后续做端侧部署验证但训练收敛会慢一些。ResNet50 这类大网络需要配合更强的正则和数据增强才能发挥优势不建议一上来就尝试。如果将来想验证 transformer 结构在这份数据上的表现可以再对比 Swin Transformer 的小版本但前提是先把 ResNet18 的 baseline 跑扎实。模型参数量建议用途说明ResNet18约 11.2M第一版 baseline训练快泛化稳定MobileNetV2约 3.4M端侧部署验证轻量精度略低ResNet50约 25.6M增强后尝试容易过拟合需更多技巧使用 ImageNet 预训练权重做迁移学习是这种小数据集最省事的选择。初始化模型后把最后一层全连接替换成输出 36 类的新层即可。4.2 训练脚本与关键超参数训练脚本我习惯封装成 20 轮左右跑一轮就能看出收敛趋势。关键超参数用下面这组AdamW 学习率 1e-4weight_decay 1e-4CosineAnnealingLR 在 20 轮内把学习率从初始值余弦降到接近 0。这组参数在大多数迁移学习任务上表现稳定不需要做太多额外调整。import torch import torch.nn as nn import torch.optim as optim from torchvision import models device torch.device(cuda if torch.cuda.is_available() else cpu) model models.resnet18(weightsmodels.ResNet18_Weights.IMAGENET1K_V1) model.fc nn.Linear(model.fc.in_features, 36) model model.to(device) criterion nn.CrossEntropyLoss() optimizer optim.AdamW(model.parameters(), lr1e-4, weight_decay1e-4) scheduler optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max20) for epoch in range(1, 21): model.train() total_loss, correct, total 0.0, 0, 0 for imgs, labels in train_loader: imgs, labels imgs.to(device), labels.to(device) optimizer.zero_grad() logits model(imgs) loss criterion(logits, labels) loss.backward() optimizer.step() total_loss loss.item() * imgs.size(0) correct logits.argmax(1).eq(labels).sum().item() total imgs.size(0) train_acc correct / total val_acc evaluate(model, val_loader, device) current_lr optimizer.param_groups[0][lr] print(fepoch {epoch:02d} | loss {total_loss / total:.4f} | train_acc {train_acc:.4f} | val_acc {val_acc:.4f} | lr {current_lr:.2e}) scheduler.step()这段代码有几点值得展开。model.fc 替换的是 ResNet18 最后的全连接层ImageNet 原本是 1000 类输出必须替换成 36 才能适配数据集。CrossEntropyLoss 内部自带 softmax网络最后一层不需要额外接激活函数。AdamW 的 weight_decay 是权重衰减用来缓解过拟合配合 CosineAnnealingLR 能在训练后期把学习率降下来让损失在收敛点附近稳定。训练几轮后你会发现 train_acc 比 val_acc 高不少这是正常的过拟合信号。此时不要急着换更大的模型优先增强数据变换把 RandomHorizontalFlip 保留增加 RandomCrop 或 RandomRotation或者把 weight_decay 提到 5e-4。数据增强是小数据集的后悔药多试组合比换模型更划算。4.3 混淆矩阵与难分对分析训练结束后val_acc 只能告诉你整体水平不能告诉你哪些类别容易混。果蔬分类里玉米和甜玉米、辣椒和甜椒、萝卜和甜菜根外观差异很小需要混淆矩阵定位具体错误。import numpy as np from sklearn.metrics import confusion_matrix all_pred, all_label [], [] model.eval() for imgs, labels in val_loader: imgs, labels imgs.to(device), labels.to(device) logits model(imgs) all_pred.extend(logits.argmax(1).cpu().numpy()) all_label.extend(labels.cpu().numpy()) cm confusion_matrix(all_label, all_pred) cm_norm cm / cm.sum(axis1, keepdimsTrue) print(cm_norm.shape) # (36, 36)confusion_matrix 返回一个 36×36 的矩阵第 i 行代表实际类别 i第 j 列代表预测为类别 j。cm_norm 把每一行归一化成 0 到 1对角线越高代表该类准确率越高非对角线数值大的格子就是重点排查的错分对。如果确认某两类高频混淆常见做法是检查图像里是否存在背景干扰或者单独给这两类补充样本。数据集本身是固定的但你可以通过混淆矩阵决定是否要用加权损失函数。提示混淆矩阵的类别索引顺序要和 label_map 一致否则画出来的矩阵行列标题全是错的。5. 避坑小数据集上的五个翻车点与排查思路数据量小、类别数多翻车点往往不在模型而在数据流程本身。下面五条是这类小样本分类数据集上比较常见的问题按现象、原因、解决的顺序写可以对照自查。5.1 现象训练集精度高、验证集精度低差距拉到 20 个百分点原因3400 张图片按 36 类平均下来每类不到 100 张模型容量稍高就很容易把训练集背下来。ResNet50 甚至更大模型在这种规模下过拟合速度非常快两三个 epoch 后 train_acc 就会接近 1而 val_acc 卡在某个位置不动。解决换小网络优先用 ResNet18 或 MobileNetV2训练轮数控制在 20 到 30 轮把 weight_decay 从 1e-4 提到 5e-4增强里加 RandomHorizontalFlip 和 ColorJitter。方向是降低模型容量、强化正则和增强而不是加数据。5.2 现象json 里明明 36 类训练完画混淆矩阵却少了两种原因label_map 用 json 生成目录却不是按 json 顺序排的。如果代码里用 os.listdir 遍历目录生成类别列表顺序就会和 json 不一致导致标签和类别名错位。只看预测索引而不对照 label_map 转回类别名就会觉得类别“变少”或“对不上”。解决全程以 json 的 label_map 为唯一标签来源。自定义 Dataset 时用 label_map.items() 生成样本列表推理时用 dict 反查类别名而不是重新排序目录。核心原则是一份映射哪里都用它。5.3 现象运行 show 脚本报 IndexError 或图片区域空白原因脚本假设每个类子目录里至少有一张可读图片但某类文件夹里可能混进了 .txt 文件或缩略图缓存也可能图片本身损坏无法解码。另一种情况是某类名称出现在 train 目录而 val 目录缺少同名文件夹。解决先跑第 2 章的 verify 脚本清理损坏图片再统计 train 和 val 的类别集合是否完全一致。不一致时把缺失类别的目录补齐或者从 label_map 中剔除该类别二选一不要两边都留着。5.4 现象loss 在下降但 val_acc 一直抖动不升像随机乱猜原因验证集预处理里混入了随机增强比如把 RandomHorizontalFlip 用在验证集上导致同一个模型对同一张图每次预测结果都不一样。另一个更隐蔽的原因是验证集和训练集有重叠同一张图被同时分进两边评估结果虚高且波动大。解决验证集只保留 Resize ToTensor Normalize 三类确定性变换。检查 train 和 val 的文件名集合确保没有交集。这两条排查完val_acc 的抖动通常会明显改善。5.5 现象把训练集文件夹丢给 YOLO 训练提示找不到 labels 目录原因这是分类数据集json 里记录的是类别标签不是目标检测的目标框坐标。YOLO 训练需要每张图对应一个 txt 标注文件格式是 class x_center y_center width height。直接套用检测流程必然失败。解决这份数据定位是图像分类用在分类网络里最顺手。如果你确实需要检测框就必须先对图片做目标框标注再用脚本转换成 YOLO 格式拿分类标签硬套检测训练行不通。图像分类和检测是两个任务评估口径也不同分类看 top-1 精度检测看 mAP别混在一起。6. 把实验固化下来随机种子、权重包与单图推理的一个习惯数据集只有 3400 张训练结果天然存在随机性。同一个脚本跑两遍val_acc 可能有 1 到 2 个点的波动。要让结论可复现我一般在训练脚本入口固定随机种子。这样换机器、换人跑至少能复现出相近的结果。import random import numpy as np import torch def set_seed(seed42): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed)set_seed 放在 import 之后、创建 DataLoader 之前能稳定绝大多数随机源。CUDA 侧的卷积实现仍可能引入少量不确定性但对 3400 张图的数据集来说影响不大。训练结束后别只存 state_dict。要把 label_map、验证精度和模型结构信息一起存进同一个权重文件这样推理时不需要额外代码去恢复类别顺序。torch.save({ model_state: model.state_dict(), label_map: label_map, val_acc: val_acc, }, fruit_veg_36_resnet18.pth)推理单张图时加载权重后用同一套 val_tf 做预处理再反查类别名。反查逻辑很简单把 label_map 的 key 和 value 互换就能得到索引到类别名的映射。def predict_one(model, img_path, label_map, device): from PIL import Image img Image.open(img_path).convert(RGB) img val_tf(img).unsqueeze(0).to(device) logits model(img) pred_idx logits.argmax(1).item() return {v: k for k, v in label_map.items()}[pred_idx]从那以后我每拿到一套新数据都会先跑一遍样本统计和图片校验再进训练流程每次实验结束把权重、类名映射和验证精度打成一个包和训练日志放同一目录。这个习惯在项目交付或复现旧结果时能省下大量时间毕竟模型可以重训但想知道当初是用哪份映射训出来的往往比重训还麻烦。希望帮到你。本文还有配套的精品资源点击获取
返回列表