ARTICLE DETAIL

资讯详情

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

基于PyTorch的CNN玉米粒品质检测:从数据增强到PyQt界面全流程

基于PyTorch的CNN玉米粒品质检测:从数据增强到PyQt界面全流程 简介本资源是一套基于Python与PyTorch实现的CNN深度学习玉米粒品质检测项目面向具备一定深度学习基础、希望掌握图像分类完整流程的学生与开发者可用于课程设计、毕业设计或农产品质检场景的实战练习。压缩包共256个文件以250张jpg图像样本为主辅以3个txt说明与3个py脚本整体约5.58MB涵盖数据集、训练代码与界面程序。项目对图片做了短边补灰、旋转与翻转等预处理以增强数据依次运行数据集文本生成、模型训练与PyQt界面脚本即可完成从数据读取、训练到可视化推理的全流程训练日志会记录每个epoch的验证损失与准确率便于观察收敛情况。目前已有127人学习下载适合想快速跑通CNN图像分类并理解数据增强与模型评估细节的读者参考。1. 玉米粒品质检测这套 CNN 代码到底能跑出什么结果分级车间里最耗人的环节不是称重是肉眼挑霉变和破损粒。一簸箕玉米倒进传送带质检员盯一下午眼睛先扛不住标准还随人漂移。这套基于 Python PyTorch 的 CNN 玉米粒品质检测代码解决的就是把「好粒 / 坏粒」二分类从人眼搬到模型上数据集目录里已经按类别分好文件夹脚本自动读取路径和标签训练完在本地落一个权重文件最后用 PyQt 拉一个能点按钮的界面做单张推理。它适合两类人——刚学完卷积神经网络想找一个完整闭环练手的以及手头有类似农产品分选需求、想拿现成流程改自己数据的。整包是「数据预处理 训练 可视化界面」三段式不是只丢一个模型定义让你自己拼这点对新手比较友好。下面按我实际拆包的顺序讲清楚每一步在干什么、参数怎么改、哪里容易翻车。2. 环境与数据管线从 requirement.txt 到增强后的图片2.1 环境安装与依赖确认拿到压缩包先别急着点运行第一步是把环境对齐。包里有个 requirement.txt这是作者给的依赖清单但深度学习环境光靠 pip 装 requirements 经常不够因为 PyTorch 的版本和 CUDA 驱动是绑死的。我一般会先确认三件事Python 版本、有没有可用的 GPU、torch 能不能正常 import。# 先看 Python 版本建议 3.8 ~ 3.10太新或太旧都可能踩依赖坑 python --version # 建一个独立虚拟环境别污染全局 python -m venv corn_env # Windows 激活 corn_env\Scripts\activate # Linux / macOS 激活 source corn_env/bin/activate # 装依赖requirement.txt 里通常含 torch、torchvision、numpy、Pillow、PyQt5 pip install -r requirement.txt # 验证 torch 是否可用以及是否吃到 GPU python -c import torch; print(torch.__version__, torch.cuda.is_available())这段逻辑很直白虚拟环境隔离依赖避免和你机器上其他项目的 numpy、torch 版本打架最后一行打印版本和cuda.is_available()返回 True 说明 GPU 可用返回 False 就是纯 CPU 跑。参数上唯一要留意的是 torch 的安装源——如果你直接pip install torch拉到的可能是 CPU 版训练会慢到怀疑人生。常见做法是去 PyTorch 官网按你的 CUDA 版本生成对应安装命令再替换掉 requirement.txt 里那一行。CPU 也不是不能跑这套数据量不大只是每个 epoch 会明显拉长。2.2 数据增强做了什么补灰边 旋转数据集文件夹按类别存放图片从文件名能看出增强痕迹aug_0_bad_data104_rotated45.jpg、aug_0_bad_data104_flip.jpg、aug_0_good_data108_rotated45.jpg这类命名说明作者对原始图做了旋转rotated45 即旋转 45 度和翻转flip并且用aug_0_前缀标记增强样本。更关键的一步是「短边补灰边变正方形」——玉米粒图往往不是正方形直接 resize 会把长宽比压变形颗粒被拉扁纹理特征就废了。作者的做法是在较短边补灰边把图撑成正方形再送进网络这样既不丢比例又满足卷积输入对尺寸的要求。from PIL import Image, ImageOps def pad_to_square(img_path, fill(128, 128, 128)): img Image.open(img_path).convert(RGB) w, h img.size if w h: return img # 本来就是正方形不动 # 以长边为目标尺寸短边两侧补灰边 target max(w, h) pad_w target - w pad_h target - h # 左右/上下各补一半 padding (pad_w // 2, pad_h // 2, pad_w - pad_w // 2, pad_h - pad_h // 2) return ImageOps.expand(img, borderpadding, fillfill)逻辑说明ImageOps.expand的 border 参数顺序是左、上、右、下所以要把总差值拆成两半分别补到两侧奇数差值时用pad_w - pad_w // 2保证总数对得上。参数fill(128,128,128)是中灰选它是因为灰色对卷积核的响应接近中性不会像纯黑或纯白那样引入强边缘干扰。这一步和旋转、翻转叠加就是作者扩增数据集的全部手段——没有用 torchvision 的 transforms 在线增强而是离线生成新图存盘好处是训练时读取快、可复现代价是磁盘占用翻几倍。2.3 01 脚本把图片路径和标签写成 txt01数据集文本生成制作.py干的事就是遍历数据集目录把每张图的路径和对应标签写进文本文件供训练脚本读取。这是典型的「先落盘再训练」思路好处是训练阶段不用再扫目录直接按行读。import os data_root dataset # 数据集根目录下面按类别分子文件夹 classes [bad, good] # 类别名顺序决定标签编号 out_file dataset.txt with open(out_file, w, encodingutf-8) as f: for label, cls in enumerate(classes): cls_dir os.path.join(data_root, cls) for name in os.listdir(cls_dir): if name.lower().endswith((.jpg, .png, .jpeg)): path os.path.join(cls_dir, name) # 每行格式图片路径 标签 f.write(f{path} {label}\n) print(done:, out_file)逻辑说明enumerate(classes)把类别映射成 0/1bad 是 0、good 是 1这个顺序一旦定了就不能在训练脚本里改否则标签全反。参数上要注意data_root和classes必须和你实际解压出来的目录名一致——很多人解压后文件夹叫「坏粒」「好粒」或者带中文脚本里写的是英文直接跑就报路径不存在。常见做法是先把目录名改成脚本里约定的名字或者反过来改脚本里的 classes 列表。输出文件每行「路径 空格 标签」训练脚本按空格切分即可。3. 训练脚本拆解02 脚本怎么读数据、怎么存模型3.1 数据集划分与 Dataset 封装02深度学习模型训练.py是核心。它先读 01 生成的 txt然后按比例切训练集和验证集再包成 PyTorch 的 Dataset。这里有个容易被忽略的点划分必须在「同一类别内」按比例切不能简单按行号切否则如果 txt 是按类别顺序写的前 80% 可能全是 bad验证集就失去意义。import torch from torch.utils.data import Dataset, DataLoader from torchvision import transforms from PIL import Image import random class CornDataset(Dataset): def __init__(self, samples, transformNone): self.samples samples # [(path, label), ...] self.transform transform def __len__(self): return len(self.samples) def __getitem__(self, idx): path, label self.samples[idx] img Image.open(path).convert(RGB) if self.transform: img self.transform(img) return img, label # 读 txt samples [] with open(dataset.txt, encodingutf-8) as f: for line in f: p, l line.strip().rsplit( , 1) samples.append((p, int(l))) # 按类别分组后再切保证验证集类别均衡 random.seed(42) by_cls {} for p, l in samples: by_cls.setdefault(l, []).append((p, l)) train, val [], [] for l, items in by_cls.items(): random.shuffle(items) cut int(len(items) * 0.8) train items[:cut] val items[cut:] # 统一尺寸 转 tensor 归一化 tf transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean[0.5, 0.5, 0.5], std[0.5, 0.5, 0.5]), ]) train_loader DataLoader(CornDataset(train, tf), batch_size16, shuffleTrue) val_loader DataLoader(CornDataset(val, tf), batch_size16, shuffleFalse)逻辑说明random.seed(42)固定随机种子保证每次划分一致方便复现。按类别分组再切是这套代码能不能训出有意义结果的关键cut int(len(items) * 0.8)就是 8:2 划分。Resize((224,224))是因为前面已经补成正方形这里直接缩到网络输入尺寸不会变形。Normalize用 0.5 均值和方差是常见做法把像素从 [0,1] 拉到 [-1,1]。参数上batch_size16是显存和速度的折中显存小就降到 8显存富裕可以上 32但学习率要相应调整。3.2 模型结构与训练循环模型部分通常是几层卷积 池化 全连接最后输出 2 类。训练循环里每个 epoch 跑训练集和验证集记录 loss 和准确率写进 log。import torch.nn as nn import torch.optim as optim class SimpleCNN(nn.Module): def __init__(self, num_classes2): super().__init__() self.features nn.Sequential( nn.Conv2d(3, 16, 3, padding1), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(16, 32, 3, padding1), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(32, 64, 3, padding1), nn.ReLU(), nn.MaxPool2d(2), ) self.classifier nn.Sequential( nn.Flatten(), nn.Linear(64 * 28 * 28, 128), nn.ReLU(), nn.Linear(128, num_classes), ) def forward(self, x): return self.classifier(self.features(x)) device torch.device(cuda if torch.cuda.is_available() else cpu) model SimpleCNN().to(device) criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lr1e-3) for epoch in range(20): model.train() for imgs, labels in train_loader: imgs, labels imgs.to(device), labels.to(device) optimizer.zero_grad() out model(imgs) loss criterion(out, labels) loss.backward() optimizer.step() # 验证 model.eval() correct total 0 with torch.no_grad(): for imgs, labels in val_loader: imgs, labels imgs.to(device), labels.to(device) pred model(imgs).argmax(1) correct (pred labels).sum().item() total labels.size(0) print(fepoch {epoch} val_acc {correct/total:.4f}) torch.save(model.state_dict(), corn_cnn.pth)逻辑说明三层卷积通道数 16→32→64 逐层翻倍是轻量 CNN 的常规配置输入 224 经过三次 MaxPool2d(2) 变成 28×28所以全连接输入是64*28*28。CrossEntropyLoss自带 softmax输出层不用再加激活。优化器用 Adam、学习率 1e-3 是入门默认值。每个 epoch 结束保存一次权重覆盖式保存所以最终留下的是最后一个 epoch 的模型——如果你想要验证集最优的那个得自己加判断只在该 epoch 准确率更高时才torch.save这是这套代码可以改进的地方。log 里记录的每个 epoch 验证损失和准确率就是判断有没有过拟合的依据训练准确率一直涨、验证准确率停滞甚至掉就是过拟合信号。4. 避坑与排查这套代码最容易翻车的五个地方4.1 路径写死导致换机器就跑不起来现象在自己电脑上跑得好好的换台机器或者换个解压目录01 脚本直接报 FileNotFoundError。原因脚本里data_root和输出路径多半是相对路径或写死的绝对路径解压位置一变就对不上。解决把路径统一改成基于脚本所在目录的相对路径用os.path.dirname(os.path.abspath(__file__))拼出根目录再往下接 dataset这样无论解压到哪都能找到。4.2 标签顺序不一致导致准确率永远 50%现象训练 loss 不降验证准确率卡在 0.5 左右。原因01 脚本里 classes 顺序是[bad,good]训练脚本里如果又按别的顺序读或者用了sorted(os.listdir())重新排标签就错位了。解决标签映射只在一处定义训练脚本直接读 txt 里的数字标签不要再自己映射一遍。改类别时两个脚本一起改。4.3 显存不足或 CPU 训练慢到放弃现象报 CUDA out of memory或者没 GPU 时一个 epoch 跑十几分钟。原因batch_size 太大或者图片没提前 resize 就送进网络。解决先把 batch_size 降到 8 甚至 4确认Resize((224,224))在 transform 里生效纯 CPU 环境把 epoch 数调小先验证流程能通再考虑上 GPU 或云平台。4.4 增强图片和原图混在一起造成数据泄漏现象验证准确率高得离谱实际用新图测试却很差。原因旋转、翻转生成的增强图如果和原图被分到了训练集和验证集两边等于验证集里出现了训练集的「近亲」指标虚高。解决划分时以「原始图」为单位分组同一张原图的所有增强版本必须落在同一侧。文件名里的data104这种编号就是原始图标识按编号分组再切。4.5 PyQt 界面加载模型报错现象03 脚本启动后点识别没反应或报 state_dict 不匹配。原因界面脚本里模型结构定义和训练脚本不一致或者加载权重时没加map_location在 CPU 机器上加载 GPU 训出的权重就崩。解决把模型定义抽成一个公共模块训练和界面都 import 同一个类加载时写torch.load(corn_cnn.pth, map_locationdevice)让权重自动落到当前设备。5. 从能跑到好用把模型换成迁移学习并做单张验证这套代码默认的 SimpleCNN 从零训练数据量不大的时候准确率上限有限。我一般会做两件事把它拉到能用的水平。第一件是换迁移学习用 torchvision 里预训练的 ResNet18把最后一层全连接改成 2 类只微调后面几层。玉米粒的纹理和霉变特征在 ImageNet 预训练特征里已经有不错的表达收敛快、对小数据集友好。import torchvision.models as models import torch.nn as nn model models.resnet18(weightsmodels.ResNet18_Weights.DEFAULT) # 冻结前面的卷积层只训练最后的分类头 for p in model.parameters(): p.requires_grad False model.fc nn.Linear(model.fc.in_features, 2) # 只这层可训练 model model.to(device) # 优化器只传可训练参数学习率可以稍大 optimizer optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr1e-3)逻辑说明weightsResNet18_Weights.DEFAULT拉预训练权重requires_gradFalse冻结主干model.fc换成新的二分类头。优化器用filter只把需要梯度的参数传进去避免更新被冻结的层。这样训练轮数可以降到 5~10 轮就有不错效果。注意输入尺寸要改成 224 且归一化参数换成 ImageNet 的mean[0.485,0.456,0.406]、std[0.229,0.224,0.225]否则预训练特征对不上。第二件是单张验证别只看验证集准确率就下结论。写个小脚本对单张图推理把预测类别和置信度打出来拿几张明显霉变、明显完好的图各测一遍看模型是不是真的学到了东西。from PIL import Image import torch def predict(img_path, model, tf, classes[bad, good]): model.eval() img Image.open(img_path).convert(RGB) x tf(img).unsqueeze(0).to(device) # 加 batch 维度 with torch.no_grad(): prob torch.softmax(model(x), dim1)[0] idx prob.argmax().item() return classes[idx], prob[idx].item() print(predict(test_bad.jpg, model, tf))逻辑说明unsqueeze(0)把单张图变成[1,3,224,224]的 batchsoftmax把 logits 转成概率argmax取最大概率类别。置信度低于 0.7 的样本建议人工复核这类边界样本往往是光照或拍摄角度导致的不是模型本身的问题。我踩过的坑是验证集准确率 0.95结果拿手机拍的图一测全错原因是训练图都是固定背景、固定光照模型学的是背景不是玉米粒。从那以后我每次换数据源都强制先跑一遍单张验证再谈部署。希望这套流程帮到你把玉米粒分选这件事真正跑通。本文还有配套的精品资源点击获取
返回列表