简介:资源包围绕柑橘成熟度识别任务,提供基于PyTorch深度学习框架的卷积神经网络完整工程,适合图像分类初学者及农业智能化项目开发者参考。包内共一百二十个文件,以一百一十三张柑橘成熟度图片为核心,另含三个Python脚本及三个文本文件,整体约十点八兆。脚本依次实现数据集文本生成、卷积神经网络模型训练和界面展示,其中数据集预处理涵盖短边灰边补正方形与随机旋转扩增,可有效提升模型泛化能力。训练完成后模型保存至本地,便于后续调用;环境依赖通过依赖清单说明,上手路径清晰。目前已有二百三十七人学习下载,对于希望快速搭建柑橘成熟度识别方案、了解数据增强与基础训练流程的用户具有直接参考价值。
1. 柑橘成熟度识别:不是调参,是先把一条数据流水线跑通
拿到这个“柑橘成熟度识别”项目时,我第一反应是去看它的数据集目录,而不是训练代码。原因很简单:做图像分类的卷积神经网络项目,90% 的坑都藏在数据怎么整理、怎么预处理、怎么喂进模型里。这个项目恰好把整条流水线拆成了三个脚本——先做数据集文本生成,再做模型训练,最后套一个 PyQt 界面,每个步骤都能单独跑、单独验证。它解决的是“从一堆原始照片到能弹窗识别柑橘成熟度”的完整闭环问题,适合做课程设计、毕业设计,或者第一次想完整跑通一个 PyTorch 图像分类项目的人。复现的价值不在于网络结构多新颖,而在于数据集生成脚本和训练脚本之间的衔接方式值得照着写一遍。
2. 数据集与预处理:灰度边补正方形、旋转扩增与 txt 生成
2.1 数据集结构分析与类别映射
打开压缩包里的数据集文件夹,能看到两类图片:Healthy和Greening,对应的就是健康成熟柑橘和青绿未熟柑橘。文件命名里藏着预处理痕迹,比如Healthy (9_rotated45.jpg是旋转 45 度后的扩增样本,Healthy (9_flip.jpg是水平翻转后的样本,原始图不带这些后缀。这意味着数据集不是“干干净净”的原始照片,而是已经预先做过一次增强的版本。
类别映射很简单,按文件夹名走:
| 文件夹名 | 含义 | 观察特征 |
|---|---|---|
| Healthy | 健康/成熟柑橘 | 表皮偏黄、色泽均匀 |
| Greening | 未熟/返青柑橘 | 表皮带明显绿色 |
有一个需要提前注意的问题:Greening类别的原始样本在文件名列表中只出现了一张。如果训练集和验证集按常规比例切分,验证集里可能只有一两条 Greening 样本,评估指标会很虚。后面避坑章节我会专门说这个问题。
2.2 01 脚本核心逻辑:遍历文件夹、按比例生成训练与验证 txt
01 脚本的名字叫“数据集文本生成制作”,它的作用是把图片路径和标签写成 txt,供训练脚本读取。常见实现思路是这样:
import os import random from glob import glob # 数据集根目录,改成你自己的实际路径 data_root = "./dataset" train_ratio = 0.8 random_seed = 42 class_names = ["Healthy", "Greening"] class_to_idx = {name: i for i, name in enumerate(class_names)} random.seed(random_seed) train_lines, val_lines = [], [] for class_name in class_names: class_dir = os.path.join(data_root, class_name) # 只筛选 jpg/jpeg/png 后缀,避免把缩略图或系统文件混进来 img_paths = glob(os.path.join(class_dir, "*.jpg")) + \ glob(os.path.join(class_dir, "*.jpeg")) + \ glob(os.path.join(class_dir, "*.png")) img_paths = sorted(img_paths) random.shuffle(img_paths) split_idx = int(len(img_paths) * train_ratio) for p in img_paths[:split_idx]: train_lines.append(f"{p} {class_to_idx[class_name]}\n") for p in img_paths[split_idx:]: val_lines.append(f"{p} {class_to_idx[class_name]}\n") with open("train.txt", "w", encoding="utf-8") as f: f.writelines(train_lines) with open("val.txt", "w", encoding="utf-8") as f: f.writelines(val_lines) print(f"训练集样本数: {len(train_lines)}, 验证集样本数: {len(val_lines)}")这段代码做的事情很直接:遍历每个类别目录,收集图片路径,打乱后按 8:2 切分训练集和验证集,写成两行一组的路径 标签格式。train_ratio控制切分比例,小数据集建议 0.8,图片多的时候可以提高到 0.85 或 0.9。random_seed保证每次运行切分结果一致,这样复现实验时不会因为随机性导致结果对不上。
值得留意的点是后缀筛选。你实际运行时系统可能生成Thumbs.db或.DS_Store,不显式限制后缀的话,这些文件会被 glob 收集进列表,训练时读图直接报错。这属于常见翻车点,后面避坑章我会展开。
2.3 灰度边补正方形的原理与 OpenCV 实现
摘要里明确提到预处理方式是“通过在较短边增加灰边,使图片变为正方形”。这一步的目的是让输入图片尺寸统一。CNN 的池化层和全连接层通常要求固定输入尺寸,直接用resize拉伸会把柑橘压成椭圆,破坏形状特征。补边则能保留原始比例,只扩展背景区域。
补黑边还是灰边?我建议不要补黑边。黑色像素的 RGB 值全是 0,经过归一化后会变成很大的负数,在网络前几层产生强烈的虚假边缘响应。补灰边是更稳的做法,YOLO 系列里填充默认值 114 也是同样的思路。常见代码是这样:
import cv2 import numpy as np def letterbox_square(img, size=224, fill_value=114): h, w = img.shape[:2] # 短边补到和长边一致,凑成正方形 if h > w: top = 0 bottom = 0 pad = (h - w) // 2 left = pad right = h - w - pad else: left = 0 right = 0 pad = (w - h) // 2 top = pad bottom = w - h - pad # 用常数填充,112 或 114 都可以,关键是固定 img_padded = cv2.copyMakeBorder( img, top, bottom, left, right, borderType=cv2.BORDER_CONSTANT, value=(fill_value, fill_value, fill_value) ) img_resized = cv2.resize(img_padded, (size, size)) return img_resized这里cv2.copyMakeBorder在短边两侧对称补边,value是填充灰度值。选 224 是因为它是 ResNet、VGG 这类常见网络的默认输入尺寸,能直接套用 ImageNet 预训练权重。如果你后面要换 EfficientNet,可能要改成 240 或 260,这个数字要和模型输入层匹配。
2.4 旋转与翻转扩增的标注一致性
扩增的逻辑是“旋转角度 + 翻转让数据集增大”。注意旋转 45 度后,图片四个角会露出黑区域。如果你的训练脚本没有对扩增图重新做描边处理,模型就会看到大量黑色角块,这会在特征图里形成固定模式,干扰识别。
我的习惯是:旋转前先补灰边成正方形,旋转后再用cv2.BORDER_REPLICATE或继续常量填充补一次角。扩增只对训练集生效,验证集永远用原始图和原始尺寸做同样预处理,不额外扩增。一句话总结,数据增强的目标是让模型在每个 epoch 看到不同的图,而不是让验证集看起来更“华丽”。
3. 训练流程:PyTorch DataLoader、迁移学习与关键参数
3.1 自定义 Dataset 读取 txt
01 脚本生成train.txt和val.txt后,02 脚本就要负责读它。PyTorch 里通过自定义 Dataset 把文本内容转为可迭代的训练样本:
import torch from torch.utils.data import Dataset from PIL import Image class CitrusDataset(Dataset): def __init__(self, txt_path, transform=None): self.samples = [] self.transform = transform with open(txt_path, "r", encoding="utf-8") as f: for line in f: line = line.strip() if not line: continue img_path, label = line.split() self.samples.append((img_path, int(label))) def __len__(self): return len(self.samples) def __getitem__(self, idx): img_path, label = self.samples[idx] img = Image.open(img_path).convert("RGB") if self.transform: img = self.transform(img) return img, label读 txt 的好处是训练和验证的划分逻辑在 01 脚本里已经定死,02 脚本不用关心数据集目录结构,换数据集时只要重新生成 txt 就行。需要特别说明的是.convert("RGB"),如果数据集中混入了灰度图或 RGBA 图,不加这一行会在torch.stack时报通道数不一致错误。这是实际复现时最常见的翻车原因之一。
DataLoader 的参数也有讲究。batch_size小数据集建议 8 或 16,过大容易显存溢出;num_workers在 Windows 上设置为 0 更稳,设置为大于 0 可能触发多进程读取的 RuntimeError。
3.2 模型选型:自定义 CNN 还是迁移学习
02 脚本里的模型结构以自带代码为准,但如果你在复现后觉得训练收敛太慢,我的建议是直接用 PyTorch 内置的 ResNet18 做迁移学习:
import torchvision.models as models def build_model(num_classes=2): model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1) # 冻结前四层,只微调后面的层 for name, param in model.named_parameters(): if "layer4" not in name and "fc" not in name: param.requires_grad = False # 替换最后一层全连接 model.fc = torch.nn.Linear(model.fc.in_features, num_classes) return model为什么这么做?这个数据集的样本量撑不起从零训练一个深层 CNN。ImageNet 预训练权重已经在海量自然图片上学过基础纹理、颜色和形状特征,柑橘表面的斑点、颜色渐变对网络来说不是陌生概念,所以只需微调最后几层。如果你的机器没有 GPU,可以把weights参数去掉,但训练时间会变长,且准确率大概率上不去。num_classes=2对应 Healthy 和 Greening 两类,后面改成多分类时这里要同步调整。
3.3 训练循环与模型保存策略
训练主体循环是标准的 PyTorch 写法,关键在 checkpoint 保存策略和优化器选择:
import torch.optim as optim from torch.optim.lr_scheduler import ReduceLROnPlateau device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") model = build_model().to(device) criterion = torch.nn.CrossEntropyLoss() optimizer = optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-5) scheduler = ReduceLROnPlateau(optimizer, mode="min", patience=5, factor=0.5) best_val_acc = 0.0 for epoch in range(50): model.train() running_loss = 0.0 for imgs, labels in train_loader: imgs, labels = imgs.to(device), labels.to(device) optimizer.zero_grad() outputs = model(imgs) loss = criterion(outputs, labels) loss.backward() optimizer.step() running_loss += loss.item() * imgs.size(0) # 每个 epoch 结束后在验证集上评估 model.eval() correct = 0 total = 0 with torch.no_grad(): for imgs, labels in val_loader: imgs, labels = imgs.to(device), labels.to(device) outputs = model(imgs) _, preds = torch.max(outputs, 1) correct += (preds == labels).sum().item() total += labels.size(0) val_acc = correct / total # 只保存验证集准确率最高的权重 if val_acc > best_val_acc: best_val_acc = val_acc torch.save(model.state_dict(), "best_model.pth") scheduler.step(val_acc) print(f"Epoch {epoch+1}: loss={running_loss/len(train_loader.dataset):.4f}, val_acc={val_acc:.4f}")ReduceLROnPlateau会在验证集准确率连续 5 个 epoch 不升时把学习率减半,这是小数据集上最实用的调度策略。保存最优模型而非最后一个模型,能避免最后的随机波动毁掉整个训练成果。代码里用with torch.no_grad()包裹验证过程,不计算梯度能省大量显存。如果你不想在训练时盯日志,可以在循环里加一个early_stop计数器,连续 10 个 epoch 验证集没有提升就提前终止。
4. 复现避坑:预处理不一致、类别不平衡与显存不足
4.1 训练 loss 震荡,验证集准确率忽高忽低
现象:loss 整体在下降,但验证集准确率每个 epoch 波动超过 5%,甚至出现不升反降。
原因:一是旋转扩增后的图片无规律地同时出现在训练集和验证集,导致验证集分布不稳定;二是在 01 脚本里随机切分时没有按类别分层采样,Greening 类样本本来就少,验证集可能分不到或只分到一条。
解决:01 脚本切分时要按类别分别做 8:2 切分,而不是把所有图片混在一起随机分。扩增图只放训练集,验证集只用原始图。具体做法是在split_idx计算前,先过滤出不带rotated、flip后缀的样本作为验证集候选。
4.2 模型精度高,但看热力图发现它学的是灰边不是果实
现象:训练时验证集准确率能到 90% 以上,但把模型部署到 PyQt 界面后,随便拍一张带背景的照片准确率骤降。
原因:训练前预处理固定用短边补灰边,图片中柑橘周围全是均匀灰色块。网络很可能学到“边缘灰色区域大 = Healthy”这类伪特征,而不是果实本身的颜色纹理。
解决:训练和推理必须走完全一致的前处理链路。更稳的做法是在补灰边后加一步随机裁剪,让柑橘在画面中的位置有轻微变化,打破灰边和类别的强关联。我复现时是先把图片从 224 放大到 256,补灰边后随机裁剪回 224,效果比纯补边好很多。
4.3 01 脚本跑完,训练时报图片读取失败
现象:train.txt生成正常,但 02 脚本训练到一半抛PIL.UnidentifiedImageError或FileNotFoundError。
原因:txt 里写的是相对路径,当前工作目录和生成 txt 时的目录不一致;或者 glob 把Thumbs.db、.DS_Store混进去了。
解决:01 脚本里改成os.path.abspath绝对路径写入 txt;后缀筛选必须同时排除*.db和*.store。如果数据集是从网盘下载解压的,还要检查是不是有嵌套目录,有的压缩包会多出一层父文件夹,导致路径前缀对不上。
4.4 GPU 显存不足,训练直接崩
现象:运行 02 脚本时出现CUDA out of memory,代码当场终止。
原因:batch_size设置过大,或者num_workers开太多导致内存碎片化。小显存显卡跑 ResNet18 时,batch_size=16加input_size=224就可能满。
解决:先把batch_size降到 8,num_workers设 0,去掉pin_memory=True。要是还崩,就把输入尺寸从 224 缩到 160,这会牺牲一点精度但能稳定跑通。跑通之后再逐步调大,明确边界之后才不慌。
4.5 Greening 类只有一张原始图,模型偏向 Healthy
现象:训练结束后打印混淆矩阵,Healthy 类 F1 很好,Greening 类查全率很低,模型总是把青果判成健康果。
原因:类别严重不平衡,Greening 原始样本只有 1 张,即使旋转翻转扩增,信息量也远小于 Healthy。单靠增强解决不了,因为增强只是把同一张图变换视角,并没有增加真实的绿色果皮纹理变化。
解决:对这个类别的扩增不能只用旋转翻转,建议叠加亮度扰动、对比度扰动和 HSV 色相微调,模拟不同光照条件下的青果外观。如果还是不行,就去网上补 10 到 20 张未熟柑橘的图片,放到 Greening 文件夹重新跑 01 脚本。数据增强救不了信息缺失,补样本才是根治办法。
5. PyQt 界面推理:加载权重、统一前处理与结果展示
5.1 界面结构:选择图片、模型推理、置信度显示
03 脚本是 PyQt 界面,基本组成是三块:一个显示图片的 QLabel、一个“选择图片”按钮、一个显示推理结果的文本框。界面逻辑不复杂,核心是把训练好的best_model.pth加载进来,对用户选的图片走一遍前处理,然后输出类别和置信度。
这里要复用训练时的前处理代码。不要重新写一套 resize 逻辑,否则很容易出现训练用 224、推理用 256 这种低级错配。
5.2 推理前处理函数与置信度输出
下面是一个可以和 PyQt 槽函数直接对接的推理函数:
import torch import cv2 import numpy as np from torchvision import transforms device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") def preprocess_for_inference(img_bgr, size=224, fill_value=114): # 和前文 letterbox_square 完全一致,保持训练推理前处理统一 img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB) img_square = letterbox_square(img_rgb, size=size, fill_value=fill_value) img_normalized = img_square.astype(np.float32) / 255.0 # 用 ImageNet 均值和标准差归一化 mean = np.array([0.485, 0.456, 0.406]) std = np.array([0.229, 0.224, 0.225]) img_normalized = (img_normalized - mean) / std tensor = torch.from_numpy(img_normalized.transpose(2, 0, 1)).unsqueeze(0) return tensor.to(device) def predict(image_path, model, class_names): img_bgr = cv2.imread(image_path) if img_bgr is None: return "无法读取图片,检查路径或文件格式", 0.0 input_tensor = preprocess_for_inference(img_bgr) model.eval() with torch.no_grad(): outputs = model(input_tensor) probs = torch.softmax(outputs, dim=1)[0] conf, pred_idx = torch.max(probs, 0) return class_names[pred_idx.item()], conf.item()归一化用的mean和std是 ImageNet 的标准值。如果你的训练脚本里面没有这段标准化,而是在 Dataset 里用transforms.ToTensor()直接归一化到 0 到 1,那推理时也要保持一致,否则输入分布错位,置信度会变得毫无意义。一个排查技巧是:跑一张训练集里的图走推理,如果输出的置信度在 0.9 以上,说明前处理链路通了;如果连训练集图片都判错,第一件事检查前处理是否和训练时一致,而不是怀疑模型权重有问题。
5.3 置信度阈值与实际使用
界面里通常要加一个置信度判断:当conf低于某个阈值时,显示“无法判断,请重新拍摄”。常见做法是阈值设在 0.6 到 0.7 之间,低于阈值就提示用户换一张更清晰、果实居中的照片。这个阈值不能写死在模型里,要在界面上留一个可调的输入框,因为不同光照条件下模型输出的校准程度不一样。加了这一层,界面的容错率会明显提升,不会出现“拿着一张糊图也硬给结论”的情况。
6. 扩展验证:多分类改造、Grad-CAM 与结果核对
如果你复现完这个二分类项目,想让它更有说服力,可以做三个方向的扩展。第一个是把二分类改成多分类,例如把 Healthy 细分为“未熟”“半熟”“成熟”。操作步骤不复杂:在数据集文件夹里新建三个子目录并把图片归类,01 脚本里把class_names改成三个类别,03 脚本里把model.fc的输出维度改成 3,重新跑一遍训练即可。注意多分类时训练集数量要尽量均衡,否则模型会偏向样本多的类别。
第二个方向是验证模型到底在看什么。用 Grad-CAM 或 PyTorch 自带的torchcam库,把最后卷积层的梯度响应叠到原图上,能直观看到模型分类时关注的区域。如果高亮区域集中在果皮纹理而不是灰边背景,说明训练是健康的;如果高亮区全在边缘填充区,就回到第 4 章讲的问题,重新审视预处理。这一步在答辩或写报告的时候非常加分,比贴准确率曲线更有说服力。
第三个方向是跨目录泛化测试。把训练好的best_model.pth放到另一台电脑,随便找几张没参与训练的真实柑橘照片跑一遍,记录准确率和置信度。这能暴露过拟合问题,特别是 Greening 类如果只依赖增强图,真实照片大概率会翻车。核对的指标不只看准确率,还要看两类各自的召回率。
从那以后,我每次拿到一个新的图像分类项目,都强制先走一遍完整的数据流:先跑 01 脚本看生成的 txt 内容,确认路径无缺失、无混入文件、类别平衡;再跑训练看 val_acc 是否合理,不急着调模型结构;最后才碰界面。数据流是这条链路的底盘,底盘不稳,后面调什么都像玄学。希望这个项目的复现过程能帮你把 PyTorch 图像分类的整套流程真正跑通。
本文还有配套的精品资源,点击获取