PyTorch苹果腐烂识别：从数据集预处理到PyQt5部署全流程-平芜编程栈

简介：面向需要对农产品外观质量做图像识别初探的开发者，这份基于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_value=114): h, w = img.shape[:2] if h == w: return img side = max(h, w) canvas = np.full((side, side, 3), fill_value, dtype=np.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 系列默认的填充灰值，适合绝大多数自然图像。如果你做的是医学影像或红外图，可以把灰边改成该场景的背景均值，自己统计一下背景像素平均值填入即可。dtype=np.uint8不能省，否则画布是 float64 类型，保存图片时要么转回 uint8 要么直接报错。

如果你更习惯用 OpenCV 自带的函数，也可以用cv2.copyMakeBorder，但要注意borderType必须设成BORDER_CONSTANT，value=(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, expand=True, fillcolor=(114, 114, 114)).save('rottenApple (463_rotated45.jpg')

逻辑说明：FLIP_LEFT_RIGHT是水平翻转，FLIP_TOP_BOTTOM是垂直翻转，对水果外观检测来说水平翻转就够用。旋转 45 度时expand=True会扩大画布，保证苹果的四个角不会被裁掉，但扩展出的角落是空白的，必须用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对应 0，rottenApple对应 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, transform=None): 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(weights=models.ResNet18_Weights.DEFAULT) model.fc = nn.Linear(model.fc.in_features, 2)

逻辑说明：把最后一层全连接替换成输出 2 类的新全连接层，前面的卷积特征提取部分完全沿用 ImageNet 预训练权重。这样模型一开始就已经具备识别纹理和边缘的能力，只需要微调对苹果腐烂特化的部分，几百张小样本也能在十几分钟内收敛到可用水平。

参数说明：weights=models.ResNet18_Weights.DEFAULT是新版 PyTorch 推荐写法，比pretrained=True更明确，旧版写法在 PyTorch 2.x 里还会弹出弃用警告。如果训练机器不能联网下载预训练权重，可以把这行改成weights=None，但收敛速度和最终准确率都会差不少。然后冻结骨干网络，只训练新加的 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(), lr=1e-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, dim=1) idx = prob.argmax(1).item() conf = prob[0][idx].item() return class_names[str(idx)], conf

逻辑说明：pad_to_square必须和训练脚本用的同一个函数，灰边颜色、居中方式都不能改。transform里面包括 Resize(224, 224)、ToTensor、Normalize，mean 和 std 必须复刻训练时的值。unsqueeze(0)是把三维的 C×H×W 变成四维的 N×C×H×W，因为模型一次处理一个 batch。新手报错"Expected 4-dimensional input for batchnormalization"，八成就是漏了这一步。

参数说明：推理时一定要包住torch.no_grad()，否则 PyTorch 会为每个中间张量保存梯度信息，显存很快被占满。softmax把输出归一化到 0 到 1，argmax(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(p=0.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 映射成 0，02 脚本训练时也读的是 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 低，说明模型把一部分新鲜苹果误杀成了腐烂，这对质检场景影响大，因为用户拿着好苹果去识别却被说成坏的。