ARTICLE DETAIL

资讯详情

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

基于CNN的鞋面缺陷识别:PyTorch实战与工业部署指南

基于CNN的鞋面缺陷识别:PyTorch实战与工业部署指南 简介这份资源面向具备一定Python基础、希望入门工业视觉缺陷检测的开发者与高校学生提供了一套基于PyTorch框架的CNN鞋面缺陷识别完整方案。压缩包共353个文件以347张jpg缺陷样本图片为核心数据辅以3个txt说明文本与3个py脚本整体约27.55MB体积轻便便于本地快速跑通。数据集已做预处理通过短边补灰边统一为正方形并施加旋转角度来扩增样本覆盖多类鞋面缺陷形态。代码按流程拆分为数据集文本生成、深度学习模型训练与PyQt可视化界面三个环节训练完成后模型会保存至本地方便后续推理与演示。目前已有152人学习适合想快速搭建缺陷分类基线、理解数据增强与CNN训练全流程的读者参考借鉴。1. 鞋面缺陷识别这套 CNN 源码为什么我建议先跑通再改鞋面缺陷识别在产线上是个挺典型的视觉检测场景布料纹理本身有周期性缺陷又分破洞、抽纱、污渍、色差好几类传统阈值分割一遇到换批次就崩。这套「基于 CNN 深度学习 python 的鞋面缺陷识别-含数据集.zip」把 PyTorch 训练脚本、推理脚本和一份已经分好类的鞋面图像数据集打包在一起拿到手就能直接训练不用自己从零标注。它适合两类人一类是想找一个真实工业缺陷数据集练手 CNN 的深度学习入门者另一类是工厂里做质检自动化、想快速验证方案可行性的工程师。我拆包之后第一件事不是看模型结构而是先把训练跑起来看 loss 曲线因为数据集质量比网络结构更能决定这套东西能不能用。2. 拆包先看数据鞋面缺陷数据集的目录结构与类别分布2.1 数据集长什么样先数一遍再决定怎么读拿到压缩包解压后常见做法是先别急着写 Dataset用几行脚本把目录结构和每类样本数摸清楚。鞋面缺陷这种工业数据集类别不平衡是常态——破洞可能几百张色差可能只有几十张不先看清楚训练时 accuracy 会骗你。import os from collections import Counter root dataset # 解压后的数据集根目录按实际路径改 for split in [train, val, test]: split_dir os.path.join(root, split) if not os.path.isdir(split_dir): continue counter Counter() for cls in os.listdir(split_dir): cls_dir os.path.join(split_dir, cls) if os.path.isdir(cls_dir): n len([f for f in os.listdir(cls_dir) if f.lower().endswith((.jpg, .png, .jpeg, .bmp))]) counter[cls] n print(split, dict(counter), 总计, sum(counter.values()))这段脚本做三件事遍历 train/val/test 三个划分统计每个类别下的图片数量最后打印总数。参数上唯一要改的是root指向你解压出来的数据集根目录。如果打印出来发现某个类在 train 里有 500 张、在 val 里只有 5 张那验证集指标波动会非常大后面调参时别被单次结果带偏。图片后缀我列了 jpg/png/jpeg/bmp 四种工业相机导出的图常见就这几类如果你的数据是 tif自己补一个后缀。2.2 类别不平衡时ImageFolder 直接读会有什么问题PyTorch 的torchvision.datasets.ImageFolder按子文件夹名当标签读起来最省事但它默认不做任何重采样。鞋面缺陷里如果「正常」样本远多于「破洞」模型会倾向于全预测成正常accuracy 看着有 90%实际漏检一堆。from torchvision import datasets, transforms from torch.utils.data import DataLoader, WeightedRandomSampler import torch train_tf transforms.Compose([ transforms.Resize((224, 224)), # 统一到 CNN 常用输入尺寸 transforms.RandomHorizontalFlip(), # 鞋面左右翻转不改变缺陷性质 transforms.RandomRotation(10), # 小角度旋转增强模拟摆放偏差 transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]), # ImageNet 统计量 ]) train_set datasets.ImageFolder(dataset/train, transformtrain_tf) # 按类别频次的反比给每个样本算权重缓解不平衡 targets [s[1] for s in train_set.samples] class_count torch.bincount(torch.tensor(targets)) class_weight 1.0 / class_count.float() sample_weight class_weight[torch.tensor(targets)] sampler WeightedRandomSampler(sample_weight, num_sampleslen(sample_weight), replacementTrue) train_loader DataLoader(train_set, batch_size32, samplersampler, num_workers4)逻辑上WeightedRandomSampler让少样本类别在每个 epoch 里被抽到的概率更高等价于做了重采样。参数说明Resize的 224 是给 ResNet 这类骨干用的如果你自己搭的小 CNN可以降到 128 省显存RandomRotation(10)的 10 度是经验值鞋面缺陷方向性不强可以再大点但破洞这种有明确形状的别转太狠Normalize用的 ImageNet 均值方差是因为后面大概率要加载预训练权重保持一致才不会让第一层输入分布错位。num_workers在 Windows 上如果报错就设成 0这是血泪经验多进程 DataLoader 在 Windows 下容易卡死。3. 模型与训练从零搭 CNN 还是拿预训练骨干微调3.1 两种路线的取舍别一上来就 ResNet50这套资源里如果带了自定义 CNN 结构通常是三到四个卷积块加全连接参数量小、训练快适合先验证数据集本身有没有问题。但鞋面缺陷的类间差异有时候很细微从零训练的小网络容易欠拟合。我的建议是先用小 CNN 跑通全流程确认数据管道没问题再换预训练骨干微调。import torch.nn as nn import torchvision.models as models def build_model(num_classes, backboneresnet18, pretrainedTrue): if backbone resnet18: model models.resnet18(weightsmodels.ResNet18_Weights.DEFAULT if pretrained else None) # 替换最后的全连接层输出改成你的缺陷类别数 model.fc nn.Linear(model.fc.in_features, num_classes) elif backbone custom: model nn.Sequential( nn.Conv2d(3, 32, 3, padding1), nn.BatchNorm2d(32), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(32, 64, 3, padding1), nn.BatchNorm2d(64), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(64, 128, 3, padding1), nn.BatchNorm2d(128), nn.ReLU(), nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(128, num_classes), ) return model model build_model(num_classes4, backboneresnet18, pretrainedTrue)build_model里num_classes必须和你数据集的实际类别数一致数错了训练时 loss 会直接报维度不匹配。pretrainedTrue会下载 ImageNet 权重第一次跑需要联网下不下来就改成 False 走从零训练。custom分支里我用了AdaptiveAvgPool2d(1)把任意尺寸特征图压成 1x1这样输入图片尺寸不用严格固定比直接 Flatten 更稳。选 ResNet18 而不是 50是因为鞋面缺陷数据集通常就几千张大骨干参数量远超样本量过拟合风险高18 层在精度和速度之间更平衡。3.2 训练循环里必须盯的三个量训练脚本本身不复杂但有几个量不看就会翻车训练 loss、验证 loss、验证集上每个类别的召回。只看总 accuracy遇到不平衡数据等于闭眼开车。import torch.optim as optim from torch.optim.lr_scheduler import CosineAnnealingLR device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device) criterion nn.CrossEntropyLoss() optimizer optim.AdamW(model.parameters(), lr1e-3, weight_decay1e-4) scheduler CosineAnnealingLR(optimizer, T_max30) for epoch in range(30): 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) scheduler.step() print(fepoch {epoch} train_loss {running_loss / len(train_set):.4f} lr {scheduler.get_last_lr()[0]:.6f})AdamW的weight_decay1e-4是给权重加 L2 正则抑制过拟合CosineAnnealingLR让学习率按余弦曲线从 1e-3 降到接近 0比固定学习率更容易收敛到好的局部解。T_max30要和总 epoch 数一致不一致的话学习率曲线会提前走完。每个 epoch 打印 train_loss 和当前学习率如果 loss 在前几个 epoch 就卡住不降先查数据标签有没有错、学习率是不是太大如果 train_loss 一直降但验证指标不涨那就是过拟合该加数据增强或者减模型容量了。验证部分我一般单独写一个函数用torch.no_grad()包起来按类别算混淆矩阵这样能看出模型到底把哪类缺陷认成了哪类。4. 推理与部署把训练好的权重接到实际图片上4.1 单张图片推理脚本怎么写才不踩预处理坑训练时用了 Resize、Normalize推理时必须一模一样地走一遍否则输入分布对不上预测结果会莫名其妙。这是最常见的翻车点训练准确率 95%拿单张图一测全是错的八成是预处理不一致。from PIL import Image import torch.nn.functional as F def predict(image_path, model, class_names, device): model.eval() 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]), ]) img Image.open(image_path).convert(RGB) # 强制三通道防止灰度图报错 tensor tf(img).unsqueeze(0).to(device) # 加 batch 维度 with torch.no_grad(): logits model(tensor) probs F.softmax(logits, dim1) conf, pred probs.max(dim1) return class_names[pred.item()], conf.item() class_names [normal, hole, stain, color_diff] # 顺序必须和训练时一致 label, conf predict(test.jpg, model, class_names, device) print(label, round(conf, 4))关键点有三个convert(RGB)防止灰度图或带 alpha 通道的图直接喂进去报维度错unsqueeze(0)补上 batch 维度因为模型期望输入是[N, C, H, W]class_names的顺序必须和ImageFolder按文件夹名排序后的顺序完全一致顺序错了标签就全乱了。softmax之后的置信度只用来做参考工业场景里低于某个阈值比如 0.7的样本最好转人工复核别硬信模型输出。4.2 批量推理和结果落盘产线上不可能一张一张手动跑常见做法是写个批量脚本遍历文件夹把结果写进 CSV方便后续统计和追溯。import csv, os def batch_predict(folder, model, class_names, device, out_csvresult.csv): rows [] for fname in os.listdir(folder): if not fname.lower().endswith((.jpg, .png, .jpeg, .bmp)): continue path os.path.join(folder, fname) label, conf predict(path, model, class_names, device) rows.append([fname, label, round(conf, 4)]) with open(out_csv, w, newline, encodingutf-8) as f: writer csv.writer(f) writer.writerow([filename, predicted_label, confidence]) writer.writerows(rows) print(f写入 {len(rows)} 条到 {out_csv}) batch_predict(test_images, model, class_names, device)out_csv里保留了文件名、预测标签和置信度三列产线追溯时按文件名就能定位到原图。置信度这一列别丢后面做阈值筛选、挑出低置信样本人工复检都靠它。如果图片量很大os.listdir一次性读进来内存吃紧可以换成os.scandir迭代但一般几千张以内没必要优化。5. 避坑与排查这套鞋面缺陷代码最容易翻车的五个地方5.1 现象训练 loss 正常下降验证准确率始终在随机水平原因通常是标签和图片没对上。ImageFolder按文件夹名映射标签如果文件夹名有中文、空格或者顺序和你的class_names不一致标签就错位了。另一个可能是验证集和训练集用了不同的 transform验证集忘了 Normalize。解决先打印train_set.classes和train_set.class_to_idx确认映射关系再检查验证集的 transform 是否和训练集一致除了数据增强部分。我一般会在训练前抽几张图连同标签可视化一遍肉眼确认没对错。5.2 现象Windows 上 DataLoader 报 BrokenPipeError 或直接卡死原因num_workers 0时 PyTorch 在 Windows 下用 spawn 启动子进程如果训练代码没有放在if __name__ __main__:保护块里子进程会重复执行主模块导致崩溃。解决把训练入口包进if __name__ __main__:或者干脆把num_workers设成 0。设 0 会慢一些但调试阶段稳定优先等跑通了再往上加。5.3 现象显存不够报 CUDA out of memory原因batch_size太大或者输入尺寸 224 对某些小显存卡来说偏高又或者验证阶段忘了加torch.no_grad()导致计算图一直累积。解决先把batch_size降到 16 或 8还不行就把输入尺寸降到 128验证和推理代码务必用with torch.no_grad():包住。另外loss.item()要取标量别把带梯度的 tensor 存进列表那也会占显存。5.4 现象模型把所有样本都预测成样本最多的那一类原因类别不平衡加上没用重采样或加权 loss模型发现全猜多数类就能拿到不错的 accuracy于是躺平。解决用第 2 章里的WeightedRandomSampler或者给CrossEntropyLoss传weight参数按类别频次反比设置。同时把评估指标从 accuracy 换成每类召回率和 F1这样少数类的表现才看得见。5.5 现象推理时单张图预测结果和训练时验证结果差很多原因预处理不一致最常见的是推理时忘了 Normalize或者Resize的插值方式不同又或者图片通道数不对灰度图没转 RGB。解决把训练和推理的 transform 抽成同一个函数或同一个变量别两处各写一遍。推理前统一convert(RGB)。如果还是对不上把推理时的 tensor 反归一化后存成图片和原图对比看是不是预处理阶段就变形了。6. 把鞋面缺陷模型用起来置信度阈值与误检复核的实操技巧训练跑通只是第一步真正上线要解决的是「模型说有问题到底信不信」。鞋面缺陷检测里漏检把缺陷判成正常和误检把正常判成缺陷的代价不一样通常漏检更致命所以阈值不能一刀切。我的做法是先把验证集所有样本的预测置信度导出来按类别画一下分布找到正常类和缺陷类置信度重叠的那段区间把阈值定在重叠区偏缺陷一侧宁可多报一点让人工复核。import numpy as np # 假设 val_probs 是验证集所有样本的 softmax 输出, val_labels 是真实标签 # 找出缺陷类非 normal的置信度分布 defect_idx [i for i, c in enumerate(class_names) if c ! normal] defect_conf val_probs[:, defect_idx].max(axis1) normal_conf val_probs[:, class_names.index(normal)] for th in np.arange(0.5, 0.95, 0.05): # 缺陷置信度高于阈值就判缺陷否则判正常 pred (defect_conf th).astype(int) true (val_labels ! class_names.index(normal)).astype(int) tp ((pred 1) (true 1)).sum() fn ((pred 0) (true 1)).sum() fp ((pred 1) (true 0)).sum() recall tp / (tp fn 1e-6) precision tp / (tp fp 1e-6) print(fth{th:.2f} recall{recall:.3f} precision{precision:.3f})这段代码遍历 0.5 到 0.9 的阈值每个阈值下算缺陷类的召回率和精确率。召回率高意味着漏检少精确率低意味着误检多、人工复核量大。产线上根据人力配置选一个平衡点比如要求召回率不低于 0.95那就选满足这个条件里精确率最高的阈值。1e-6是防止除零验证集里如果没有缺陷样本分母会是 0。另一个实用技巧是保存误检样本。每次批量推理后把置信度在阈值附近比如阈值上下 0.1 区间的图片单独拷到一个文件夹定期看这些图你会发现模型反复在某一类纹理或某种光照条件下出错。把这些图补进训练集重新微调比盲目加数据增强有效得多。我一般会保留一个hard_cases文件夹每次迭代往里加模型版本更新时优先用这批数据验证有没有改善。从那以后我每次拿到新的缺陷数据集都强制先跑一遍类别分布统计和预处理一致性检查再动模型结构。这套鞋面缺陷识别的代码包把数据、训练、推理都串好了省掉的是搭框架的时间但数据本身的问题还是得自己盯。希望帮到你。本文还有配套的精品资源点击获取
返回列表