简介面向需要对农产品外观质量做图像识别初探的开发者这份基于PyTorch的苹果腐烂识别工程以Python为主要语言从数据集整理、模型训练到界面演示形成一条完整可运行的流程。资源共601个文件压缩包64.41MB其中531张jpg、63张png和1张jpeg构成新鲜苹果与腐烂苹果的分类图像数据集3个py脚本分别对应数据集文本生成、深度学习模型训练、PyQt界面启动3个txt文件记录环境依赖与标签信息。预处理环节会在较短边补灰边使图片变为正方形并加入角度旋转扩增样本可减少背景与尺度差异对识别的影响训练完成后模型保存在本地方便继续调参或迁移到其他农产品外观分类场景。已有130人学习适合刚配置好PyTorch环境、想获得一份带图片数据的完整分类项目作为入门参照的读者。1. 这套苹果腐烂识别资源包解决的不只是分类问题苹果腐烂识别很多人以为难点在模型实际上拿到这套资源后你会发现数据集准备和推理链路才是真正占用时间的地方。这套基于 Python PyTorch 的深度学习二分类代码包把整个流程拆成了三个可独立运行的脚本01 脚本扫描数据集生成 txt 标签清单02 脚本读取清单完成模型训练并保存权重03 脚本用 PyQt5 做一个拖图片即可识别的桌面界面。适合三类人第一次跑图像分类完整流程的初学者想快速搭建腐烂检测 Demo 的从业者以及想复用它来做其他果蔬外观检测的开发者。目录和接口都算清楚整体改造工作量不大。2. 数据集预处理与增强灰边填充和翻转旋转的三个关键细节拿到资源第一件事我建议先打开数据集文件夹把文件列表翻一遍。你会看到类似freshApple (184_flip.jpg、rottenApple (463_rotated45.jpg这种命名规律非常直白原始图片名 下划线 增强方式。这说明该数据集的预处理是离线做的先把每张原始图生成了水平翻转和旋转 45 度两个副本再连同原图一起放进训练样本池。相比每轮随机增强的在线扩增这种离线扩增优点在于样本稳定、实验结果可复现缺点是占用磁盘空间而且如果翻转图被同时分到训练集和验证集验证集指标会虚高一点后面避坑章节会再提到。2.1 为什么必须先把苹果图片变成正方形ResNet18 这类带全连接层的分类网络输入尺寸必须固定。如果直接把非正方形图片 resize 成 224×224椭圆形的苹果会被拉宽或压扁腐烂区域的纹理比例也跟着失真。补灰边的做法是把短边所在的画布加宽到与长边一致苹果的物理比例完全不变多余部分用灰色填充。这样做既保住了形态信息又把背景干扰控制在边缘区域思路和 YOLO 系列一直在用的 letterbox 一致图像分类预处理里非常稳妥。注意补灰边不是把整张图硬塞进固定画布而是先算出原图的长边以长边作为正方形边长把原图居中贴上去。如果原图本身就是正方形这步直接跳过不加边框。灰边颜色也有讲究推荐用 RGB(114, 114, 114) 而不是黑色或白色。纯黑会在边缘形成强烈对比模型容易学到边缘有黑块就是腐烂这种伪特征纯白在浅色背景的苹果图片上同理。灰色介于两者之间对后续归一化也更友好。2.2 灰边填充代码最稳妥的正方形化写法虽然资源包里已经生成了预处理后的图片但你要换自己的数据集时这步必须自己实现。我一般用 OpenCV 加 NumPy 这样写import cv2 import numpy as np def pad_to_square(img, fill_value114): h, w img.shape[:2] if h w: return img side max(h, w) canvas np.full((side, side, 3), fill_value, dtypenp.uint8) y_offset (side - h) // 2 x_offset (side - w) // 2 canvas[y_offset:y_offset h, x_offset:x_offset w] img return canvas逻辑说明先判断图片是否已经是正方形是就直接返回避免多此一举加边框。然后取长边作为画布尺寸创建一个三通道的灰色画布。偏移量用整除计算目的是让原图居中当差值不是偶数时会差一个像素对识别结果没有实际影响。最后把原图像素拷贝到画布对应区域。参数说明fill_value是灰边颜色值三通道都是 114这就是 YOLO 系列默认的填充灰值适合绝大多数自然图像。如果你做的是医学影像或红外图可以把灰边改成该场景的背景均值自己统计一下背景像素平均值填入即可。dtypenp.uint8不能省否则画布是 float64 类型保存图片时要么转回 uint8 要么直接报错。如果你更习惯用 OpenCV 自带的函数也可以用cv2.copyMakeBorder但要注意borderType必须设成BORDER_CONSTANTvalue(114, 114, 114)。两种写法效果一样函数式写法更容易控制偏移量建议直接复刻上面这段。2.3 翻转和旋转 45 度离线增强的正确打开方式从文件命名看这套数据集对每张原始图做了两个增强水平翻转和旋转 45 度。水平翻转对苹果这类没有左右语义差别的目标非常合适比旋转更安全不会把腐烂斑块扭曲得脱离真实形态。用 PIL 做这两个增强最小实现是这样from PIL import Image img Image.open(rottenApple (463).JPG) img.transpose(Image.FLIP_LEFT_RIGHT).save(rottenApple (463_flip.jpg) img.rotate(45, expandTrue, fillcolor(114, 114, 114)).save(rottenApple (463_rotated45.jpg)逻辑说明FLIP_LEFT_RIGHT是水平翻转FLIP_TOP_BOTTOM是垂直翻转对水果外观检测来说水平翻转就够用。旋转 45 度时expandTrue会扩大画布保证苹果的四个角不会被裁掉但扩展出的角落是空白的必须用fillcolor填上灰色否则默认是黑色模型会学到有黑角就是坏苹果的错误关联。参数说明旋转 45 度是一个相对保守的角度不会造成严重畸变。如果你想做更激进的数据增强常见做法是随机旋转正负 45 度但固定 45 度对复现实验更友好。还有一种常见误用是只旋转不扩画布导致苹果边缘被切掉标签内容变了模型被迫从残缺形状里学特征。实操时建议旋转前后都做一次尺寸检查确认苹果主体完整。3. PyTorch 训练脚本txt 标签体系、ResNet18 与模型保存01 脚本扫描数据集生成 txt 文件02 脚本只认这个 txt这种以文本为中间产物的设计有一个明显好处换数据集时不需要动训练逻辑只要改 01 脚本里的类别名和路径就行。我改过类似项目把二分类扩成五分类时只动了 01 脚本和最后界面的标签映射02 脚本一行没改。对这种小规模视觉项目先用 txt 把路径和标签固定下来比直接用 ImageFolder 文件夹结构更直观训练脚本里的依赖也更少。3.1 从 txt 文本到自定义 Dataset 类01 脚本的核心任务是遍历每个类别文件夹把图片路径和对应标签写到 txt 里每行一条记录。注意数据集里同时存在.JPG和.jpeg后缀判断图片时要做大小写归一化import os def generate_txt(root, out_path): with open(out_path, w) as f: for label_name in [freshApple, rottenApple]: folder os.path.join(root, label_name) label_id 0 if label_name freshApple else 1 for filename in os.listdir(folder): if filename.lower().endswith((.jpg, .jpeg, .png)): img_path os.path.join(folder, filename).replace(\\, /) f.write(f{img_path} {label_id}\n)逻辑说明遍历两个类别文件夹把类别字符串映射成数字标签freshApple对应 0rottenApple对应 1。注意这里label_id的赋值顺序要和后期界面显示保持一致否则会出现预测结果和实际显示名称相反的诡异情况。参数说明filename.lower()是这一步最容易漏的细节否则.JPG大写在 Linux 上会漏扫。.replace(\\, /)是为了保证路径分隔符统一避免反斜杠在 txt 解析时被当成转义符。txt 里每行格式是路径 空格 标签后面 Dataset 解析时也要按这个格式拆分。对应的 Dataset 类写起来很标准from torch.utils.data import Dataset from PIL import Image class AppleDataset(Dataset): def __init__(self, txt_path, transformNone): self.pairs [] with open(txt_path, r) as f: for line in f: img_path, label line.strip().rsplit( , 1) self.pairs.append((img_path, int(label))) self.transform transform def __len__(self): return len(self.pairs) def __getitem__(self, idx): img_path, label self.pairs[idx] img Image.open(img_path).convert(RGB) if self.transform: img self.transform(img) return img, label逻辑说明__getitem__是 PyTorch Dataset 的核心接口DataLoader 每次批量取数据时都会按索引调它。这里先打开图片并强制转成 RGB避免灰度图和 RGBA 图混进来导致通道数不一致。标签从 txt 读出来是字符串int(label)转成数字后再返回如果不转CrossEntropyLoss 会在计算时直接报类型错误。参数说明rsplit( , 1)比split( )更保险。因为图片路径里可能有括号和空格用split( )会把freshApple (184_flip.jpg这种路径拆成两段路径就废了。rsplit从右侧拆分只拆最后一位保证路径完整。这一点放在避坑章节还会再强调。3.2 ResNet18 迁移学习与训练参数选择二分类任务不需要上很大的网络ResNet18 在这个资源规模下比 ResNet50 更好训收敛快且不容易过拟合。VGG16 也可以但参数太多CPU 训练会很慢。常见做法是直接用 PyTorch 自带的 torchvision 预训练权重做迁移学习import torch.nn as nn from torchvision import models model models.resnet18(weightsmodels.ResNet18_Weights.DEFAULT) model.fc nn.Linear(model.fc.in_features, 2)逻辑说明把最后一层全连接替换成输出 2 类的新全连接层前面的卷积特征提取部分完全沿用 ImageNet 预训练权重。这样模型一开始就已经具备识别纹理和边缘的能力只需要微调对苹果腐烂特化的部分几百张小样本也能在十几分钟内收敛到可用水平。参数说明weightsmodels.ResNet18_Weights.DEFAULT是新版 PyTorch 推荐写法比pretrainedTrue更明确旧版写法在 PyTorch 2.x 里还会弹出弃用警告。如果训练机器不能联网下载预训练权重可以把这行改成weightsNone但收敛速度和最终准确率都会差不少。然后冻结骨干网络只训练新加的 fc 层是迁移学习里先冷后热的标准打法for param in model.parameters(): param.requires_grad False for param in model.fc.parameters(): param.requires_grad True optimizer torch.optim.Adam(model.fc.parameters(), lr1e-3) criterion nn.CrossEntropyLoss()逻辑说明冻结所有卷积层权重后梯度只流向最后的全连接层训练速度快、显存占用低。跑通基线后如果想继续提升再解冻最后几个残差块做一轮低学习率微调一般学习率降为 1e-4。参数说明optimizer 只传入model.fc.parameters()这一步很关键。如果直接传model.parameters()预训练权重一开始就会被大规模更新几轮后特征提取部分就被破坏小数据集上很容易过拟合。3.3 模型保存格式state_dict 与标签映射文件训练过程中要保存验证集准确率最高的模型而不是最后一个 epoch 的模型。我一般按验证准确率做判断模型更好就覆盖保存一次import json best_acc 0 if val_acc best_acc: best_acc val_acc torch.save(model.state_dict(), best_model.pth) with open(class_names.json, w) as f: json.dump({0: freshApple, 1: rottenApple}, f)逻辑说明保存state_dict而不是整个 model 对象文件体积更小跨环境迁移不会带着 Python 版本和类定义的依赖。加载时要先用代码里同一个网络结构实例化模型再调用load_state_dict。参数说明class_names.json必须和模型配套保存。界面脚本加载模型后需要根据数字索引反查类别名称如果丢了映射文件模型能跑但显示的结果永远是 0 或 1用户根本看不懂。后期改分类类别数时这个 json 也要同步改。4. PyQt5 界面加载模型推理链路与两类预处理对齐03 脚本用 PyQt5 把训练好的模型封装成桌面界面选一张苹果图片点击识别返回新鲜或腐烂的类别和置信度。这章核心不是 PyQt5 控件怎么写而是推理预处理必须和训练时完全对齐。很多人在这一步翻车训练准确率很高界面识别却一塌糊涂。4.1 单张图片推理的完整流程推理链路本质上是把一张图片从原始状态逐步转换成模型输入先做灰边正方形化再 resize 到 224×224然后 ToTensor 和 Normalize最后加一个 batch 维度送入模型。完整代码大致这样def predict(image_path, model, transform, class_names): img Image.open(image_path).convert(RGB) img pad_to_square(np.array(img)) img Image.fromarray(img) img transform(img) img img.unsqueeze(0) with torch.no_grad(): outputs model(img) prob torch.softmax(outputs, dim1) idx prob.argmax(1).item() conf prob[0][idx].item() return class_names[str(idx)], conf逻辑说明pad_to_square必须和训练脚本用的同一个函数灰边颜色、居中方式都不能改。transform里面包括 Resize(224, 224)、ToTensor、Normalizemean 和 std 必须复刻训练时的值。unsqueeze(0)是把三维的 C×H×W 变成四维的 N×C×H×W因为模型一次处理一个 batch。新手报错Expected 4-dimensional input for batchnormalization八成就是漏了这一步。参数说明推理时一定要包住torch.no_grad()否则 PyTorch 会为每个中间张量保存梯度信息显存很快被占满。softmax把输出归一化到 0 到 1argmax(1)取置信度最高类别的索引conf就是那个类别的置信度。界面显示时可以直接把这个 conf 转成百分数。4.2 推理 transform 和训练 transform 的对齐检查训练脚本中如果有随机翻转、随机旋转这类在线增强推理时不能带进去否则每次识别结果可能不一样。我通常把训练和验证共用一套 transform 定义在同一个 dict 里推理时单独写一个不含随机操作的版本infer_transform transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ])逻辑说明这套 mean/std 是 ResNet 系列常用的 ImageNet 归一化参数。如果你的训练脚本自定义了 mean/std推理脚本一定要同步替换否则模型看到的输入分布完全不一样。参数说明Resize((224, 224))是在灰边填充之后再做的先把图片变成正方形再 resize这样苹果不会被拉伸。顺带提醒如果训练脚本里开了RandomHorizontalFlip(p0.5)推理时不能加因为推理结果应该只取决于原图本身。5. 常见问题与避坑环境配置、显存和标签错位的实战记录这套代码我拿到的第一反应是先把 requirement.txt 里列出的包跟版本逐条核了一遍。训练脚本坑不算多但环境、显存、标签映射、文件大小写这四个位置最容易出问题。每一条我都实际踩过或者帮别人排查过按现象、原因、解决三步记录如下。5.1 requirement.txt 逐行安装后 import torchvision 报错现象按照 requirement.txt 里列的包名逐个 pip install装完执行import torchvision直接报找不到某个 C 扩展或者 PyQt5 打开界面时弹出 qt.qpa.plugin 找不到窗口插件。原因requirement.txt 如果不锁定版本pip 会默认安装当前最新版本而 PyTorch 和 torchvision 必须严格配套两者 C 扩展的 ABI 是绑定的。另外 PyQt5 和某些版本的 OpenCV 会共享 Qt 动态库装混了容易出现插件冲突。解决装 PyTorch 系列不要用普通 pip直接用官方源安装配套版本比如pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118。装好后先做一次自检python -c import torch, torchvision; print(torch.__version__, torchvision.__version__)能正常输出版本号说明环境基本通了。如果只是 CPU 环境就选 cu 前缀改为 cpu 的版本训练速度慢一些但结果可用。PyQt5 和 OpenCV 如果冲突优先统一通过 requirements 安装别混着 conda 和 pip 一起装。5.2 训练刚跑几轮就报 CUDA out of memory现象batch_size 设成 32训练没到十个 epoch 就报显存不足把 batch_size 降到 8 甚至 4 还是偶尔会炸。原因OOM 不一定只是 batch_size 的锅。如果训练脚本在验证阶段忘了包torch.no_grad()验证时也会为中间激活张量保留梯度缓存显存占用直接翻倍。另外每个 epoch 结束后频繁地torch.save也会触发显存碎片累积。解决先把验证循环用with torch.no_grad():包住这是最容易被忽略的一行。训练循环里适当位置加torch.cuda.empty_cache()释放缓存。然后把num_workers设为 0 排查 DataLoader 多线程带来的显存压力。如果仍然报错做一个最小化实验把输入图改成 128×128 跑几十个 batch能正常跑就说明瓶颈在激活值占用继续降 batch_size 或输入尺寸128 也跑不了则要认真检查代码里是否有张量泄漏。5.3 训练集准确率高界面识别结果却和真实标签相反现象模型在训练集上 acc 接近百分之百打开界面识别一张明显新鲜的红苹果返回结果却是 rottenApple置信度还有九十多完全反了。原因txt 生成脚本里 label 顺序和界面脚本里 class_names 顺序没有对齐。01 脚本里 freshApple 映射成 002 脚本训练时也读的是 0但 03 脚本的 class_names 字典写成了{1: freshApple, 0: rottenApple}模型输出的 0 被界面解释成了腐烂。解决txt 里不要只写数字标签直接写成类别名Dataset 里再用字典映射这样中间任何一环出问题字符标签都不会漂移label_map {freshApple: 0, rottenApple: 1} label label_map[label_str]对应class_names也用同一个字典反查生成。如果你习惯在 txt 里写数字那就把 01、02、03 三个脚本的标签映射固定写在同一处不要每个脚本各自定义一遍。5.4 数据集图片上千张生成的 txt 却少了一半现象01 脚本跑完txt 里只有几百行数据集里明明有一千多张图。进一步排查发现部分图片文件明明存在训练时却报 FileNotFoundError。原因文件后缀大小写不统一endswith(.jpg)无法命中.JPG和.jpeg另外部分文件名里有括号和空格如果 txt 解析时用split( )或split(\t)路径在空格处被拆断后半截被当成标签读入图片 URL 无效。解决判断后缀时统一做大小写归一化也就是filename.lower().endswith((.jpg, .jpeg))。解析 txt 时用rsplit( , 1)代替split( )只从右侧切一次保证路径完整。生成路径时把反斜杠统一替换成/避免在 Linux 容器里出现转义问题。6. 模型验证的上分技巧先用混淆矩阵和难例回放确认错在哪训练完只看准确率是不够的二分类准确率 95% 和 99% 在生产环境里是两种完全不同的体验。我习惯在验证阶段把每个 batch 的预测结果和真实标签都收集起来最后用混淆矩阵和分类报告定位模型到底在错哪类。6.1 混淆矩阵与分类报告在验证循环里把预测索引和标签分别追加到两个列表最后用 sklearn 输出报告from sklearn.metrics import confusion_matrix, classification_report cm confusion_matrix(all_labels, all_preds) print(classification_report( all_labels, all_preds, target_names[freshApple, rottenApple] ))看混淆矩阵时重点看两个位置新鲜苹果被误判成腐烂的数量以及腐烂苹果被漏检的数量。如果 freshApple 的 precision 高但 recall 低说明模型把一部分新鲜苹果误杀成了腐烂这对质检场景影响大因为用户拿着好苹果去识别却被说成坏的。6.2 难例回放与阈值调整把预测错的图片单独存到一个文件夹重新打开看一遍。我遇到过两类典型难例一张图里有多个苹果其中只有一个局部腐烂模型按整图分类就容易犹豫还有背景里的棕色纸箱或木纹被模型学成了腐烂特征。难点往往不在苹果本身而在背景干扰。二分类的默认决策阈值是 0.5如果你做的是宁可多报腐烂也不能漏检的质检逻辑把腐烂类别的置信度阈值降到 0.4 甚至 0.3再用新的阈值跑一遍验证集通常能提高召回率。代价是误检率会上升界面里把两个方向的指标都打印出来再权衡。从那以后我每次训完这类检测模型都会强制走一遍混淆矩阵加难例文件夹确认错的到底是苹果特征还是背景干扰再决定要不要调阈值、加样本或者换网络。希望帮到你。本文还有配套的精品资源点击获取