简介面向表面缺陷视觉检测场景提供基于轻量化重构网络的完整Python项目实现涵盖源码、预训练权重与文档说明适合计算机视觉、人工智能等方向的在校学生、老师及企业开发者用于课程设计、毕业设计或初期项目演示。压缩包共562个文件包含400张图像样本、17个Python源码文件、5个PTH模型权重、8个XML配置及58个BAT运行脚本整体大小约62.08MB目录结构清晰便于按需检索。目前已有64人浏览学习代码均经过测试运行成功功能可靠可放心使用。资源附带README说明文档结合图像样本与脚本可快速复现检测流程理解轻量化重构网络在表面缺陷识别中的实际应用同时可在该代码基础上修改扩展适配其他视觉检测任务适合不同基础的学习者进阶参考。1. 轻量化重构网络做表面缺陷视觉检测先搞清这份源码解决什么问题表面缺陷视觉检测在工业现场最头疼的不是算法选型而是缺陷样本根本凑不齐。划痕、凹坑、麻点、氧化斑种类多到没法用分类网络一个个列出来标数据的人工成本比训练本身还贵。这份 Python PyTorch 源码换了个思路只用正常表面的图像训练一个轻量化重构网络让它学会还原正常表面长什么样推理时拿真实输入图和重构结果做差残差明显的地方就是缺陷。模型结构走深度可分离卷积加瓶颈残差块显存占用比常规 U-Net 小一个量级普通消费级显卡就能训练。资源包里有完整项目源码、prepare.bat 和 show_logs.bat 两个脚本外加文档说明从环境准备到训练再到日志查看顺序跑完就能看到曲线和检测结果。适合做毕设、课设也适合想快速验证重构式异常检测思路的从业者新手照着 README 能落地熟手可以直接改网络和损失函数做二次开发。2. 重构网络为什么能检测缺陷重构误差、轻量化设计与目录结构2.1 从分类到重构缺陷样本不够时的务实选型传统视觉检测思路是先收集缺陷样本标好类别再训练分类网络或者目标检测模型。这个路线在品类单一的产线上没问题但表面缺陷恰恰是长尾分布某条产线上划痕占八成凹坑占一成半剩下的是说不清名字的异常。少样本类别喂不饱分类器硬凑数据增强也只是缓解。重构网络Reconstruction Network直接把任务换掉——识别缺陷变成找出和正常不一样的地方。训练时只喂正常样本模型要把正常表面的纹理、光照、周期性结构都编码进权重里。推理时输入一张图经过编码器压缩、解码器还原得到重构图。正常表面重构出来与输入几乎一致缺陷区域的模式编码器没见过重构不出来残差就在缺陷位置出现明显峰值。这个思路在 MVTec AD、NEU-DET 这类工业表面数据集上都有成熟验证核心优势是训练阶段不需要正负样本均衡缺陷检测问题被简化成正常样本建模 残差判定。我一般把整个流程拆成四步特征压缩、特征重建、残差计算、阈值分割。这份源码的推理流程也是这个顺序只是阈值分割被放在后处理脚本里方便单独调整不用每次动网络结构。2.2 轻量化具体轻在哪深度可分离卷积与瓶颈残差块轻量化重构网络的核心不是简单把网络做小而是在保证重构质量的前提下砍参数量和计算量。常见做法是用深度可分离卷积Depthwise Separable Convolution替换标准卷积标准卷积同时做空间滤波和通道融合深度可分离卷积把它拆成逐通道卷积加逐点卷积两步。一个标准 3×3 卷积的参数量是 3×3×C_in×C_out拆开后变成 3×3×C_in 加 1×1×C_in×C_out。C_in 和 C_out 都是 128 时参数量大约从 14.7 万降到 1.7 万差了一个数量级。import torch import torch.nn as nn class DepthwiseSeparableConv(nn.Module): 深度可分离卷积块逐通道卷积 逐点卷积 可选残差连接 def __init__(self, in_channels, out_channels, stride1): super().__init__() self.dw nn.Conv2d(in_channels, in_channels, kernel_size3, stridestride, padding1, groupsin_channels) self.pw nn.Conv2d(in_channels, out_channels, kernel_size1) self.bn1 nn.BatchNorm2d(in_channels) self.bn2 nn.BatchNorm2d(out_channels) self.act nn.ReLU(inplaceTrue) self.use_residual (stride 1 and in_channels out_channels) def forward(self, x): identity x x self.act(self.bn1(self.dw(x))) x self.bn2(self.pw(x)) if self.use_residual: x x identity return self.act(x)注意几个参数groups 设成 in_channels 就是逐通道卷积每个通道独立卷积不跨通道融合stride1 且通道数不变时才允许残差相加stride2 的下采样层通道数会翻倍残差接线对不上直接放弃。这和 MobileNetV2 的瓶颈结构思路一致源码里编码器前几层就是这种块堆出来的。整体结构是编码器逐级下采样把 256×256 的输入压到 32×32 左右的瓶颈特征图解码器用最近邻上采样加深度可分离卷积还原回原尺寸。这里有一个很关键的设计决定跳跃连接的处理。纯 U-Net 层层加跳跃连接是为了保住分割细节但重构网络里跳跃连接太密解码器可以直接抄编码器的输出网络偷懒把输入原样搬出来缺陷区域也就跟着被重构出来了检测直接失效。源码里只在最浅层保留一条跳跃连接中间和深层全靠瓶颈特征重建强迫网络真的学会理解正常表面而不是做图像复制。2.3 先读目录再跑代码五块内容一眼分清拿到源码先别急着点 prepare.bat花五分钟把目录结构过一遍后面排错会快很多。这份资源目录大致分五块路径作用说明models/网络结构定义轻量化重构网络、损失函数datasets/数据加载与预处理正常样本读取、增强、归一化tools/训练、测试、推理脚本prepare_data.py、train.py、infer.pyconfigs/超参与路径配置图片尺寸、学习率、数据集根目录logs/训练日志与权重TensorBoard 事件文件、.pth 权重数据流是datasets 读图 → 缩放到统一尺寸 → 归一化 → 进模型前向计算 loss → 反向传播 → 每 N 个 epoch 存一次权重同时把 loss 和验证重构误差写进 TensorBoard 日志。推理是另一条线读权重、加载模型、逐张图算残差、阈值分割、输出结果图。这里有一个常见误解有人打开 datasets 发现训练目录只有一类正常样本以为是资源不全。重构网络的训练数据本来就应该是纯正常的混入缺陷样本反而会让模型学会重构缺陷检测就废了。官方数据集通常已经把训练集和测试集按这个规则切好自己准备数据时务必把缺陷图从训练目录里单独拿走。3. 从 prepare.bat 到训练曲线环境配置与一键运行的完整链路3.1 运行环境Python、虚拟环境与依赖安装这份源码是标准 Python PyTorch 工程环境要求不苛刻。Python 3.8 以上就行PyTorch 1.8 到 2.x 都能跑依赖主要是 torch、torchvision、opencv-python、numpy、tensorboard、tqdm。有 N 卡就装 CUDA 版 PyTorch纯 CPU 也能跑只是 256×256 的图一个 epoch 会慢不少。先把虚拟环境和依赖装齐python -m venv .venv # Windows 激活 .venv\Scripts\activate # Linux/macOS 激活 source .venv/bin/activate pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118 pip install opencv-python numpy tensorboard tqdm说明第一行创建虚拟环境避免依赖污染系统 Python第二、三行按平台激活环境第一条 pip 指定了 CUDA 11.8 的 PyTorch 安装地址如果机器 CUDA 版本不同把 cu118 换成 cu121或者去掉 --index-url 装默认版本。后面一行装的是图像处理、数值计算、日志可视化和进度条四个基础库。装不上时优先检查两件事。一是 Python 版本是不是 3.12 以上太新的版本和部分旧版 torch 包有兼容问题我一般直接用 3.10 最省事。二是 pip 下载超时把 pip 索引地址换成国内可用的加速源再重装。环境这步是整个项目的地基在这里卡住后面全白搭所以我把这一步写在最前面而不是只看 README 就跳过去。3.2 prepare.bat 做了什么目录、依赖与数据预检这个脚本名字叫 prepare作用是在训练前把目录建好、检查依赖、扫描数据。它不会替你下载几个 G 的数据集数据文件需要自己放进 datasets/ 对应目录脚本只负责校验和建索引。脚本内容大致是这个骨架echo off chcp 65001 nul echo [1/4] 创建目录结构... if not exist logs mkdir logs if not exist weights mkdir weights if not exist datasets\train mkdir datasets\train if not exist datasets\test mkdir datasets\test echo [2/4] 检查 Python 环境... python -c import torch, torchvision; print(torch, torch.__version__) if errorlevel 1 ( echo [ERROR] 请先安装依赖参考 README pause exit /b 1 ) echo [3/4] 检查数据集... python tools/prepare_data.py --data_root datasets if errorlevel 1 ( echo [ERROR] 数据集目录为空或格式不对 pause exit /b 1 ) echo [4/4] 预检完成可以开始训练 pause逐行解释chcp 65001 是防止 Windows 控制台中文乱码if not exist 是幂等建目录重复运行不会报错python -c 那一行用 import 测试 torch 是否能正常加载能打印版本号说明环境没问题errorlevel 是批处理里的退出码判断import 失败就提示装依赖并暂停。tools/prepare_data.py 做的事是扫描数据目录、按文件名过滤正常样本、生成训练索引清单顺便统计图片尺寸是否统一。一个常见误操作是看到弹窗就狂点鼠标。prepare 卡在 pause 上不是死机是批处理在等人按任意键确认。想看完整输出就老老实实在项目根目录打开 cmd 窗口手动执行 prepare.bat所有日志都能滚动回看比双击跑完什么都看不见强得多。3.3 训练入口与关键参数动手前先把超参摸清训练入口是 tools/train.py参数用 argparse 管理不想改代码时直接在命令行覆盖。这是最常用的参数组合python tools/train.py \ --data_root datasets \ --img_size 256 \ --batch_size 16 \ --epochs 100 \ --lr 1e-3 \ --backbone light_mobile \ --loss msessim \ --save_dir weights参数逐个说明img_size 控制训练分辨率256 是速度和效果比较平衡的档位想更精细可以试 384显存会跟着涨batch_size 按显存调8G 显存跑 256 分辨率建议取 8 到 16lr 用 1e-3 起步配合 cosine 衰减或每 30 个 epoch 降为原来的十分之一都行backbone 指定轻量化网络变体light_mobile 对应第四章讲的结构light_unet 是带更多跳跃连接的版本数据量大时效果略好但参数量也大。loss 选 msessim 是这份源码里的主流组合mse 负责像素级逼近ssim 负责结构相似度两者加权相加默认权重前者 0.7、后者 0.3。训练开始后控制台每隔一个 iter 打印一次当前 loss每个 epoch 结束打印平均 loss 和验证重构误差。显存不够就把 batch_size 降到 4epochs 相应加到 150分辨率不变的前提下效果差距不大。这套参数是我调过几次后觉得最不容易翻车的一档新手不用再动。3.4 show_logs.bat 与 TensorBoard训练过程的可视化入口训练过程中所有指标都会写进 logs/ 目录下的 TensorBoard 事件文件。show_logs.bat 本质就是帮你启动 TensorBoard 服务的包装脚本echo off chcp 65001 nul call .venv\Scripts\activate.bat tensorboard --logdir logs --port 6006 --bind_all start http://localhost:6006解释一下call 激活虚拟环境tensorboard 默认只绑本机回环地址--bind_all 是让同一局域网的其他机器也能访问方便你在另一台电脑上看曲线--port 6006 是常用端口被占用就改成 6007。运行完脚本会尝试自动打开浏览器没弹出来就手动访问 http://localhost:6006。在 TensorBoard 里主要盯三条曲线train_loss 是否平稳下降、eval_recon_error 验证重构误差是否收敛、learning_rate 学习率变化。一个实用习惯是每跑 20 个 epoch 看一次验证重构误差如果它在后期开始反弹说明过拟合了回退到上一个 epoch 的权重文件重新调阈值比从头重训省钱。4. 核心模块拆解重构前向、损失函数与缺陷定位后处理4.1 重构网络前向编码器、瓶颈与浅层跳跃连接网络结构在设计上有几个矛盾点下采样太狠会丢细节重构出来的表面发糊太浅又记不住正常纹理。这份源码用三级下采样加瓶颈256 输入压到 32×32 特征图解码器逐级上采样还原整体参数量大概只有常规 U-Net 的四分之一到三分之一。import torch.nn as nn class LightReconNet(nn.Module): 轻量化重构网络深度可分离卷积 瓶颈 浅层跳跃连接 def __init__(self, in_channels3, base_channels32): super().__init__() self.head nn.Conv2d(in_channels, base_channels, 3, padding1) # 编码器三级下采样通道翻倍 self.enc1 self._make_stage(base_channels, base_channels * 2, stride2) self.enc2 self._make_stage(base_channels * 2, base_channels * 4, stride2) self.enc3 self._make_stage(base_channels * 4, base_channels * 8, stride2) # 瓶颈不再下采样只做特征压缩与融合 self.bottleneck self._make_stage(base_channels * 8, base_channels * 8, stride1) # 解码器最近邻上采样 深度可分离卷积 self.dec3 self._make_stage(base_channels * 8, base_channels * 4, stride1) self.dec2 self._make_stage(base_channels * 4, base_channels * 2, stride1) self.dec1 self._make_stage(base_channels * 2, base_channels, stride1) self.tail nn.Conv2d(base_channels, in_channels, 3, padding1) def _make_stage(self, in_c, out_c, stride): return DepthwiseSeparableConv(in_c, out_c, stridestride) def forward(self, x): h self.head(x) e1 self.enc1(h) # 256 - 128 e2 self.enc2(e1) # 128 - 64 e3 self.enc3(e2) # 64 - 32 b self.bottleneck(e3) # 32 - 32 d3 self.dec3(self.upsample(b)) # 32 - 64 d2 self.dec2(self.upsample(d3)) # 64 - 128 d1 self.dec1(self.upsample(d2) e1) # 128 - 256浅层跳跃连接 out self.tail(d1) return out代码里的 upsample 我直接在类里定义成 nn.Upsample(scale_factor2, modenearest)取最近邻插值而不是双线性原因是最近邻不引入额外平滑重构图边缘更锐利残差图里缺陷边界不会被糊掉。浅层跳跃连接加在 d1 和 e1 之间这两级特征图空间分辨率高携带的是细粒度纹理加一条跳跃能让解码器把真实表面质感带回来。中间两级刻意不加连接所有语义信息必须经过瓶颈特征重建这是重构误差能暴露缺陷的关键。如果显存还有余量可以把 base_channels 从 32 提到 48重构质量会有可感知的提升但训练时间大概增加一半需要自己权衡。4.2 损失函数MSE、SSIM 与加权组合损失函数直接决定网络在优化什么。只用 MSE网络倾向于输出平均化的正常表面纹理细节被磨平缺陷区域的残差和正常区域的纹理残差混在一起阈值很难切。只用 SSIM结构相似度对亮度偏差不敏感灰度均匀的缺陷类型会漏。所以这份源码用加权组合。损失特点适合场景MSE逐像素 L2对灰度差异敏感但容易模糊边缘均匀纹理表面如钢板、铝板L1比 MSE 温和异常点带来的梯度更小噪声较多的粗糙表面SSIM关注局部结构相似能保留边缘锐度纹理、周期性图案表面import torch import torch.nn.functional as F try: from pytorch_msssim import ssim except ImportError: ssim None class ReconLoss(nn.Module): 重构损失MSE SSIM 加权组合 def __init__(self, mse_w0.7, ssim_w0.3, data_range1.0): super().__init__() self.mse_w mse_w self.ssim_w ssim_w self.data_range data_range def forward(self, pred, target): mse F.mse_loss(pred, target) if ssim is not None: # SSIM 越大越相似所以用 1 - ssim 作为惩罚项 s ssim(pred, target, data_rangeself.data_range) else: s 1.0 return self.mse_w * mse self.ssim_w * (1 - s)两个权重参数是调参重点。纹理表面可以把 ssim_w 拉到 0.4让边缘结构更锐利表面灰度很均匀的mse_w 拉到 0.8 效果更直接。pytorch_msssim 是一个独立库没装的话 import 会失败代码里做了降级处理——没有 SSIM 时只跑 MSE不会直接崩。建议还是装上重构类任务里 SSIM 的收益非常明显。4.3 推理与缺陷定位残差图、自适应阈值与后处理推理阶段的任务是把重构误差转成缺陷区域。代码流程是这样import cv2 import numpy as np import torch def infer_one(model, img_tensor, device, fallback_threshold0.1): 推理重构 残差计算 阈值分割返回重构图、残差图、掩码 model.eval() with torch.no_grad(): recon model(img_tensor.unsqueeze(0).to(device)) # 残差图逐像素绝对差跨通道取均值 res torch.abs(recon.squeeze(0) - img_tensor).mean(dim0) res res.cpu().numpy() # 归一化到 0~1避免不同批次的绝对数值不可比 res (res - res.min()) / (res.max() - res.min() 1e-8) # 先尝试大津法自适应阈值 gray (res * 255).astype(np.uint8) th, mask cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY cv2.THRESH_OTSU) # 大津法算出的阈值过低时回退到固定阈值 if th fallback_threshold * 255: th fallback_threshold * 255 _, mask cv2.threshold(gray, th, 255, cv2.THRESH_BINARY) return recon, res, mask这段代码里最容易忽视的是归一化。残差图绝对数值随图像亮度变化同一模型在不同光照下残差均值差几倍都有可能所以先做 min-max 归一化再进阈值。大津法Otsu是自适应阈值正常图没有缺陷时残差分布接近单峰Otsu 会把噪声均值附近强行切成两类产生一片伪缺陷所以加了一道回退逻辑——Otsu 算出的阈值低于固定下限时就改用固定阈值。实际工程里我一般再叠加一次形态学操作开运算去掉小于 9 平方像素的噪点闭运算把断裂的缺陷区域连起来这一步能直接决定最终检测框的完整度。5. 避坑指南换数据、调超参和工程化的五类常见问题5.1 训练 loss 一路下降测试却发现所有图都被判成有缺陷现象train loss 降到 0.01 以下验证集的整体重构误差也很低但推理时几乎每张残差图都整片发白阈值分割后全是缺陷区域。原因最常见的是训练和推理的预处理不一致。训练时用了 mean/std 归一化推理脚本里没做或换了一组不同的 mean/std输入分布直接错位。另一种情况是模型过拟合到了训练集的光照和色调分布换一张光照不同的正常图进来重构误差整体被抬高。解决把训练和推理的预处理完全对齐——同一份 normalize 参数同一个 resize 插值方式这一步写在代码里做成公共函数谁也不要单独改。另外验证时不要只看 loss 数值抽样打印几张正常样本的残差均值如果正常样本和缺陷样本的残差均值差距小于 0.02说明模型在死记硬背需要去掉更深层的跳跃连接或加大数据增强强度。5.2 显存不足batch_size16 直接 CUDA out of memory现象按 README 默认参数跑刚启动几十个 iter 就报 CUDA out of memory进程直接退出。原因不一定是 batch_size 的锅。256 分辨率下的三级下采样虽然把特征图压缩了但解码器上采样叠加训练的中间变量仍然吃显存另外验证阶段如果没关梯度计算图被保留显存占用直接翻倍。解决第一反应是把 batch_size 降到 4 确认能跑然后每轮翻倍输入分辨率降到 192 也能显著缓解。验证阶段务必包在 torch.no_grad() 里并且确认每个 batch 结束调了 optimizer.zero_grad()这两个操作能省掉至少两成显存。真的还差一口气就把网络里的 BatchNorm 换成 GroupNorm显存占用会再降一截。5.3 prepare.bat 双击闪退或提示python 不是内部或外部命令现象双击 prepare.bat黑窗一闪就消失手动执行则报python 不是内部或外部命令。原因python 没加进系统 PATH命令行里根本找不到这个命令。闪退往往是 chcp 65001 之后脚本里出现编码异常或者批处理走到 exit /b 直接退出没给人看报错的机会。解决在项目根目录打开 cmd手动执行 prepare.bat先让报错停住用 where python 确认 python 解释器路径系统装了多个 Python 版本的就用 py -3.10 显式指定。用了 3.1 的虚拟环境方案后激活状态下 bat 里调用的 python 来自 .venv大部分 PATH 问题都能绕开。另外注意 bat 文件保存编码选 ANSI别用带 BOM 的 UTF-8否则第一行注释在 chcp 之后可能变成乱码解析失败。5.4 换成自己的数据集效果明显崩掉现象官方示例数据集训练效果不错换到现场拍的几百张正常图重构误差高得离谱该检出的缺陷大面积漏检。原因工业现场数据普遍存在光照不均、镜头畸变、背景干扰分布和官方数据集差很远。更隐蔽的问题是自己收集的正常图里大概率混着带缺陷的图模型学进去之后缺陷被当成了正常模式重构误差自然压不住。解决第一步过滤数据把缺陷图人工挑出去宁可少放几十张也不能混入一张缺陷图第二步固定采集条件光源、相机高度、角度都锁定让训练数据和测试数据来自同一成像环境第三步如果光照差异实在压不住在预处理里加直方图均衡化或灰度归一化。重构网络对分布漂移很敏感训练集和测试集来自不同产线时效果必然下降这不是模型结构的问题是数据一致性的问题。5.5 TensorBoard 打开了却看不到曲线现象show_logs.bat 启动成功浏览器能访问 6006 端口但 SCALARS 面板一片空白。原因事件文件目录不对。训练脚本把日志写在 logs/exp001/ 这类子目录而 tensorboard --logdir logs 扫描的是 logs 根目录旧版本 TensorBoard 对多层目录支持不友好。另一种情况是训练进程刚启动第一条事件还没写入文件是空的。解决把 --logdir 指到具体子目录例如 tensorboard --logdir logs/exp001 再刷新页面确认训练脚本里用的是 SummaryWriter(log_dirlogs/exp001)而不是默认的 runs/ 目录最后看 logs 下有没有事件文件文件名以 event 开头、后缀 .tfevents 且文件大小大于 0才说明有数据写入。我排查这类问题固定三步走目录路径、文件大小、端口占用十分钟内能定位。6. 进阶用法残差热力图可视化与缺陷面积统计6.1 残差热力图可视化直观看到网络注意哪里调试重构网络最直观的手段是把残差图叠加到原图上。缺陷区域会以高亮度呈现整个表面的纹理噪声则被摊平在低灰度区。import cv2 import numpy as np def save_heatmap(img_bgr, res_map, save_path): 将归一化的残差图转成热力图并叠加到原图 heat cv2.applyColorMap((res_map * 255).astype(np.uint8), cv2.COLORMAP_JET) overlay cv2.addWeighted(img_bgr, 0.6, heat, 0.4, 0) cv2.imwrite(save_path, overlay)6.2 用连通域统计缺陷面积与坐标线上检测光给一张热力图不够下游要的是缺陷面积和位置。用连通域分析把掩码转成统计量def defect_stats(mask): 统计每个缺陷连通域的面积和中心坐标返回排序后的列表 n, labels, stats, _ cv2.connectedComponentsWithStats(mask, 8) results [] for i in range(1, n): # 0 是背景 area stats[i, cv2.CC_STAT_AREA] if area 25: # 小于 5x5 的噪点过滤 continue results.append({ area_px: int(area), center: (int(stats[i, cv2.CC_STAT_LEFT] stats[i, cv2.CC_STAT_WIDTH] / 2), int(stats[i, cv2.CC_STAT_TOP] stats[i, cv2.CC_STAT_HEIGHT] / 2)), }) return sorted(results, keylambda r: r[area_px], reverseTrue)面积单位是像素转成实际平方毫米需要预先标定像素与物理尺寸的比例标定方法是用一张已知长度的工件图像做像素当量换算。热力图和面积统计加在一起才算把检测闭环成可用的检测结果。这套流程我前后跑了三轮最大的教训是重构网络的调参没有银弹每个数据集都要重新过一遍数据纯净度、归一化一致性和阈值标定。从那以后我每次拿到新表面数据都强制先抽 20 张正常图跑一遍残差分布摸底确认正常样本的残差峰值远低于缺陷样本再进训练循环。这个习惯帮我挡掉过好多次无效训练希望帮到你。本文还有配套的精品资源点击获取