简介一份面向计算机、人工智能、数据科学等相关专业学生与初学者的边缘检测实践项目基于深度学习完成轮廓提取任务内含HED、PiDiNet等经典模型的Python源码、预训练权重与配套数据集可完整复现训练与推理流程尤其适合毕业设计、课程设计、项目演示及入门进阶。压缩包共三十四个文件体积约8.72MB主要类型包括源码py、模型权重pth、样本图片jpg/png、标注文件xml和使用说明md/txt图片与标注配合可构建训练样本模型权重可直接加载推理说明文档辅助环境搭建与参数调整。目前已有六十七人浏览学习。项目附带运行说明与调试输出参考图便于对照验证检测效果环境配置完成后可快速运行也可基于现有结构改造网络或迁移至其他图像检测场景二次开发空间较大。1. 能直接跑起来的深度学习边缘检测这份 HEDPiDiNet 源码包解决什么问题边缘检测在视觉项目里卡过不少人Canny 在干净背景上还行一旦进到纹理复杂或者弱边缘场景不是断线就是噪声一起冒出来。如果你正在做毕设、课程设计或者要把边缘提取接进分割、检测流程最省力的方式不是从头复现一篇论文而是先拿一份能直接跑的源码来拆。这份《基于深度学习的边缘检测模型python源码模型数据集.zip》给的是 HED 和 PiDiNet 两个主流模型的 Python 实现附带.pth权重、示例图片、测试列表和 debug 输出解压之后改改路径就能把普通图片转成边缘概率图。适合刚入手深度学习的学生也适合想用边缘作先验信息的从业者。2. 先看算法选型HED 和 PiDiNet 为什么能站稳边缘检测2.1 传统边缘检测为什么经常翻车Sobel、Prewitt、Canny 这些传统算子本质上是固定卷积核的局部梯度响应。Sobel 和 Prewitt 对亮度变化敏感但没有语义理解能力Canny 加了高斯平滑和双阈值看起来更聪明可阈值一设不合适边缘要么断裂要么把纹理细节全部当成边缘输出。我在实际项目里测过布料表面边缘Canny 对条纹边缘响应很强但它分不清“布料上的花纹”和“布料边界”哪个才是需要的边缘。真正要检测物体轮廓时需要模型理解图像内容。这就是深度学习边缘检测的价值通过标注数据学习“边缘到底是什么”用上下文信息补全局部梯度缺失因此在弱边缘、阴影、遮挡场景下表现比传统算子稳定得多。2.2 HED 的整体嵌套设计到底做了什么HED 全称 Holistically-Nested Edge Detection整体嵌套边缘检测。它用 VGG16 作为骨干网络从不同卷积阶段引出多个侧输出每个侧输出都对应一个尺度的边缘概率图最后用一个融合层把多尺度结果汇总。HED 最重要的设计是深度监督。传统 CNN 只需要最后一层计算 lossHED 在每个侧输出层都单独计算 loss等于同时让网络浅层学习局部细节、深层学习整体轮廓。训练时侧输出权重可以调节推理时既可以用单独的侧输出也可以用融合后的结果。我拆这份源码的时候特意看了hed_edge.py里的模型结构骨干部分基本遵循 VGG16 的分类层截断但后面接的侧输出层是卷积加反卷积的组合。由于感知野不同浅层侧输出边缘细但噪声多深层侧输出边缘连续但定位粗融合以后两者互补。这是 HED 能扛住复杂场景的根本原因。2.3 PiDiNet 用像素差卷积换轻量化PiDiNet 是另一个方向的代表不追求大网络而是用像素差卷积 PDCPixel Difference Convolution去模拟传统梯度算子。它的核心思想是把像素差值当作卷积特征的一部分既保留了局部灰度变化信息又不需要像 VGG 那样堆叠超大卷积核。和 HED 相比PiDiNet 参数量小得多推理速度快很适合做实时应用。发行版里table5_pidinet.pth这个命名方式应该是参考论文 Table 5 的配置训练出来的权重这类权重对输入分辨率有一定要求跑之前最好先看一眼模型定义里的输入尺寸避免直接拿高分辨率图硬塞。这份资源同时给出 HED 和 PiDiNet等于给你两条路线一条重精度、一条重速度。毕设里做算法对比时这两个模型正好可以形成对照。2.4 解压后文件结构对应哪些环节解压后第一眼会看到一堆.py、.pth、图片和debug_out开头的 png。理解文件归属比直接运行更重要。按常见命名习惯拆一下文件/目录在项目里承担的角色hed_edge.pyHED 模型定义和推理入口重点看侧输出层定义pidinet.pyPiDiNet 模型定义包含 PDC 模块image2edge-main可能是图像转边缘的流程目录放工具脚本或补充代码pipNet/pipNet_mainPiDiNet 相关工程目录常见于多任务版本table5_pidinet.pthPiDiNet 预训练权重加载时注意键名匹配test.txt测试图片路径列表批量推理时按行读路径debug_out_*.png模型输出的边缘概率图可直接用于效果对比若干.jpg/.jpeg/.png示例输入图片也是最快的冒烟测试素材项目里的.idea是 PyCharm 工程配置通常可以忽略white_pic_*.jpg这类名字从内容看是带白色背景的样本图用来测边缘模型会不会把背景纹理误判为边缘。3. 把环境跑通PyTorch 版本、模型权重与第一张边缘图3.1 PyTorch 和依赖版本怎么选才不折腾这类项目大多数基于 PyTorch 1.x 写最新版 PyTorch 2.x 在多数情况下也能跑但要注意.pth文件的加载兼容性。如果碰到算子兼容问题我一般建议直接装 PyTorch 1.8 到 1.13 之间的稳定版本避免torchvision和模型源码里的旧接口不匹配。conda create -n edge_det python3.8 conda activate edge_det pip install torch1.13.1 torchvision0.14.1 -i https://pypi.tuna.tsinghua.edu.cn/simple pip install numpy opencv-python pillow matplotlib tqdm这里-i指定的是清华 PyPI 镜像国内下载快一些。要注意 Python 3.8 是兼容性比较好的选择3.10 以上跑旧代码可能出现np.float被移除之类的问题。安装完成后先用一句话确认环境python -c import torch; print(torch.__version__)能打印出版本号说明 PyTorch 本体没问题。如果导入时报DLL load failed大概率是 CUDA 和 torch 版本不匹配换成 CPU 版 torch 最省心。3.2 用 hed_edge.py 跑单张图片hed_edge.py如果支持命令行参数一般逻辑是--input 图片路径 --output 结果路径如果不支持就需要把文件末尾的调用部分改成下面这种形式。import torch import cv2 from PIL import Image import numpy as np from torchvision import transforms # 假设 hed_edge.py 里已经定义好 HED 模型构建函数 build_hed() from hed_edge import build_hed model build_hed() state_dict torch.load(hed.pth, map_locationcpu) model.load_state_dict(state_dict) model.eval() img Image.open(7.jpeg).convert(RGB) tensor transforms.ToTensor()(img).unsqueeze(0) with torch.no_grad(): outputs model(tensor) if isinstance(outputs, (list, tuple)): # HED 有多个侧输出取融合结果通常是最后一个 edge_map outputs[-1].squeeze().numpy() else: edge_map outputs.squeeze().numpy() edge_map (edge_map - edge_map.min()) / (edge_map.max() - edge_map.min() 1e-8) edge_map (edge_map * 255).astype(np.uint8) cv2.imwrite(hed_output.png, edge_map)逻辑说明模型前向输出如果包含多个侧输出HED 的融合层一般放在列表最后所以取outputs[-1]。这一步容易被忽略如果取第一个侧输出会得到一张噪声偏大的边缘图。后处理部分做了 min-max 归一化防止因为激活值分布不均匀导致保存出来的边缘图过暗。参数说明map_locationcpu表示即使用没有 CUDA 的机器也能加载权重unsqueeze(0)给单张图片补 batch 维度。如果你的权重是用 GPU 保存的又没有改map_location加载时会报CUDA not available的错。3.3 用 pidinet.py 批量测试test.txt 的写法test.txt在这个包里通常保存测试图片路径每行一个。批量推理时可以先读取这个文件再逐一喂给模型。python pidinet.py --test list --file test.txt --save_dir ./results --threshold 0.5有的版本不走命令行而是直接在脚本里写死文件路径。如果脚本里写死的是test.txt就和包里的目录结构对上了。with open(test.txt, r, encodingutf-8) as f: paths [line.strip() for line in f if line.strip()] for idx, p in enumerate(paths): img cv2.imread(p) # 这里略去 resize 和归一化实际使用时需要按模型的输入尺寸处理 out model_inference(img) cv2.imwrite(f./results/{idx:05d}.png, out)逻辑说明按行读路径过滤空行是为了避免最后一行的换行符导致路径拼错。批量推理时每张图写一个独立文件文件名按序号补齐到五位方便后续和标签文件对齐。参数说明threshold这类阈值参数在边缘检测里很关键。0.5 作为二值化阈值只适合轮廓清楚的高质量图片如果边缘图偏灰建议先看概率分布再决定阈值不要上来就硬切。3.4 第一次跑通后先看哪几张输出包里自带7.jpeg、5.jpg、1.jpeg这些图片我建议先用它们冒烟测试不用急着拿自己的数据。成功输出debug_out_1.png、debug_out_7.png这种文件后打开看一眼边缘是否连续、背景有没有大量白色噪点。第一次跑通的目标不是效果完美而是确认模型前向、后处理和保存环节都没问题。如果输出全黑第一反应是归一化出了问题或者是把概率图当成边缘图直接保存后面避坑章里详细说。4. 从推理到训练数据集格式、损失函数与参数调整4.1 数据集目录怎么设计才能跟现有代码兼容这份包里没有显式给出train.py但hed_edge.py和pidinet.py已经有完整模型定义自己补一个训练脚本并不难。关键是先把数据集目录约定好。边缘检测的数据集一般是“原图 边缘标注图”成对出现标注图通常是单通道白色表示边缘黑色表示背景。常见组织方式dataset/ images/ 000001.jpg 000002.jpg edges/ 000001.png 000002.png train.txt val.txttrain.txt和test.txt类似每行可以写一对路径也可以用固定拼接方式省去标签路径。最简单是每行写images/000001.jpg edges/000001.png中间用空格分开。def read_pair_file(path): pairs [] with open(path, r) as f: for line in f: line line.strip() if not line: continue img_path, gt_path line.split() pairs.append((img_path, gt_path)) return pairs逻辑说明这种一对一的路径列表是最稳的做法后续无论做随机裁剪还是数据增强都能保证原图和标注同步操作。不要只在train.txt里写图片路径然后让程序去猜边缘文件位置很容易因为命名不一致翻车。参数说明图片和边缘标注的尺寸默认应该一致。如果原图是 512x512标注也是 512x512训练时直接随机裁剪到 256x256 即可。标注图用 PNG 保存是为了避免 JPEG 压缩带来的伪边缘。4.2 基于 hed_edge.py 改一个最小训练循环HED 的训练和普通分类不一样它有多个侧输出每个侧输出都会算一个损失最后所有损失加权求和。下面这个例子是“直接在现有模型定义上补训练逻辑”的常见写法。import torch import torch.nn as nn from torch.utils.data import Dataset, DataLoader from hed_edge import build_hed class EdgeDataset(Dataset): def __init__(self, pair_list, crop_size256): self.pairs pair_list self.crop_size crop_size def __len__(self): return len(self.pairs) def __getitem__(self, idx): img_path, gt_path self.pairs[idx] img cv2.imread(img_path, cv2.IMREAD_COLOR) gt cv2.imread(gt_path, cv2.IMREAD_GRAYSCALE) # 这里简化处理完整版需要加随机裁剪和数据增强 img cv2.resize(img, (self.crop_size, self.crop_size)) gt cv2.resize(gt, (self.crop_size, self.crop_size)) img torch.from_numpy(img).float().permute(2, 0, 1) / 255.0 gt torch.from_numpy(gt).float().unsqueeze(0) / 255.0 return img, gt model build_hed() optimizer torch.optim.Adam(model.parameters(), lr1e-4) criterion nn.BCEWithLogitsLoss() for epoch in range(30): for i, (img, gt) in enumerate(train_loader): optimizer.zero_grad() outputs model(img) # HED 会返回多个侧输出 loss 0 for out in outputs: loss criterion(out, gt) loss / len(outputs) loss.backward() optimizer.step() if i % 50 0: print(fepoch {epoch}, step {i}, loss {loss.item():.4f})逻辑说明HED 的多个侧输出都要和同一张边缘标注计算损失所以循环里累加每个侧输出的损失再求平均。BCEWithLogitsLoss比手动在输出层接 Sigmoid 再算二值交叉熵更稳定因为它在内部做了数值处理。参数说明学习率1e-4是边缘检测里比较常用的起点Adam 优化器对初始值不那么敏感。crop_size256控制原图和标注同时裁剪的尺寸越小训练越快但边缘定位越粗糙。4.3 训练参数怎么给学习率、batch size 与关键注意点用一张参数表说明我平时调训练时的配置参数推荐值说明batch size4~8边检测模型输入分辨率高太大容易显存溢出初始学习率1e-4模型骨干是预训练权重时用 1e-4 比较稳权重衰减5e-4防止过拟合尤其数据集只有几百张时随机裁剪尺寸256x256兼顾感受野和显存太大训练很慢增强方式翻转、旋转、色彩抖动边缘模型对亮度敏感色彩抖动能提升泛化性训练轮数30~50边缘检测收敛比分类慢30 轮通常刚进入稳定区有一个容易被忽略的点边缘标注里边缘像素远少于背景像素直接算交叉熵会让模型倾向把全部预测成背景。HED 论文里用了类别平衡策略也就是在 loss 计算时给边缘像素更高的权重。常见的简化做法是在BCEWithLogitsLoss里传入pos_weight。criterion nn.BCEWithLogitsLoss(pos_weighttorch.tensor([5.0]))这个pos_weight表示边缘正样本的权重是背景的 5 倍。具体数值要看你的数据集边缘占比如果边缘图特别细可以调到 8 到 10如果边缘很粗2 到 3 就够。这也是很多人训练边缘检测模型时 loss 正常但效果全黑的原因。4.4 从 debug 输出判断训练是不是真的在学训练过程中最怕 loss 一直在降输出却看不到纹理。这份包里留了多张debug_out_*.png其实就是给了一个很好的验证习惯每个 epoch 结束固定拿 3 到 5 张验证图跑一次推理把结果存下来看。如果 loss 下降但输出越来越黑基本是类别不平衡压过了模型让模型学会输出全零图。这时候先别调网络结构把pos_weight调大一点或者把边缘标注的白色区域膨胀几个像素增加正样本占比往往比换模型更有效。如果输出边缘很粗但位置准说明浅层侧输出占主导可以试着在融合层加大深层侧输出的权重。这属于结构级别的调参改完要重新训练但能明显改善边缘定位精度。5. 边缘检测模型常见的坑从环境到结果全排查5.1 现象加载.pth报错提示 key 不匹配加载table5_pidinet.pth时控制台打印一堆Missing key(s)和size mismatch模型参数没有被正确载入。原因多半是模型定义和权重来源不对应。table5_pidinet.pth可能是针对 PiDiNet 的某个特定配置训练出来的而pidinet.py默认初始化的网络结构和它不一致比如卷积核数量不同、输入头不同。另外有些权重保存时带了module.前缀和单卡模型的键名不一致也会导致 key 不匹配。解决先打印权重的 key 和模型 state_dict 的 key 做对比。命令很简单state torch.load(table5_pidinet.pth, map_locationcpu) print(list(state.keys())[:5]) print(list(model.state_dict().keys())[:5])如果发现所有键名都多了module.前缀说明权重是在 DataParallel 下保存的加载前去掉前缀即可new_state {k.replace(module., ): v for k, v in state.items()} model.load_state_dict(new_state)5.2 现象输出边缘图全黑或者全白跑完推理保存的图片是一张纯黑色图或者所有像素都是同一个值。原因最常见的是后处理直接把模型输出当成了图像保存没有做归一到 0-255 的操作。模型最后输出的概率值在 0 到 1 之间如果直接cv2.imwrite会把 0.5 以下全部存成 0看起来就全黑。全白一般是激活函数用错或者把tanh的输出不加归一化直接存。解决保存前强制做一次线性拉伸out torch.sigmoid(outputs[0]) if with_logits else outputs[0] out out.squeeze().numpy() out (out - out.min()) / (out.max() - out.min() 1e-8) out (out * 255).astype(np.uint8)另外检查图片读取是不是读成 BGR 后混了通道。如果cv2.imread之后又做了np.transpose通道顺序错乱会导致特征分布异常输出自然不正常。5.3 现象训练时 loss 一直在高位震荡降不下去训练流程能跑但 loss 在 0.8 到 1.0 之间波动几十个 epoch 也不下降。原因边缘检测的标注大多是细线结构如果 batch size 太大并且没有做随机裁剪大量图片里边缘像素占比很低梯度被背景像素主导。另一个常见原因是BCEWithLogitsLoss里没有处理类别不平衡导致模型永远预测背景。解决先给 loss 加pos_weight同时把输入图像做随机裁剪和翻转增加边缘在局部窗口里的占比。再一个经验是不要用 ImageNet 预训练权重直接全量微调而是先在骨干网络冻结状态下训 10 个 epoch再放开全模型训练。这样能避免模型前期被边缘的强噪声带偏。5.4 现象中文路径下图片读不出来但英文路径一切正常把项目放在D:\毕设\代码这样的目录下运行cv2.imread返回None或者PIL报编码错误。摘要里强调不要用中文路径这确实是血泪。原因OpenCV 的imread对中文路径支持很差PyTorch 的数据加载器在读取路径时也可能因为编码不一致失败。更隐蔽的是某些第三方库内部调用 C 接口读取文件路径只要含中文就失败。解决最省事是解压后直接把项目目录改成英文名比如hed_pidinet_code路径里不要有任何中文。如果临时要用中文路径换PIL读取并先开numpyfrom PIL import Image img np.array(Image.open(img_path).convert(RGB))这种方式在 Windows 下对中文路径的容忍度高得多但治标不治本。项目根目录保持英文是最推荐的。5.5 现象边缘图断续严重要么噪声点多到一个头两个大模型能出边缘但线条断续明显或者背景有大量细碎白色块。原因如果用的是 HED可能只取了某个侧输出而不是融合结果如果用的是 PiDiNet可能是输入分辨率太低导致小结构丢失。另一个常见原因是后处理二值化阈值太低导致背景噪声也过了阈值。解决先把输出从outputs[-1]换成outputs[3]或者中间层对比哪一层更接近你想要的边缘再对概率图做中值滤波和形态学闭运算可以显著改善连续度。比如import cv2 prob (out * 255).astype(np.uint8) prob cv2.medianBlur(prob, 5) _, bin_edge cv2.threshold(prob, 128, 255, cv2.THRESH_BINARY) bin_edge cv2.morphologyEx(bin_edge, cv2.MORPH_CLOSE, np.ones((3, 3), np.uint8))这里的medianBlur去掉单像素噪声MORPH_CLOSE把断裂的小缝隙补上。注意不要用大核否则边缘位置会偏移形态学核的大小视图像分辨率而定一般 3x3 到 5x5 就够了。6. 进阶技巧把 HED 边缘输出做一次清洗再喂给下游分割6.1 快速批量生成边缘图假设你已经跑通了单张推理下一步一定是批处理。原始包里有图片也有test.txt我习惯于先写一个小脚本把自己验证集里的图全部转成边缘图。import os import cv2 import numpy as np import torch def infer_batch(model, img_paths, save_dir): os.makedirs(save_dir, exist_okTrue) for idx, p in enumerate(img_paths): img cv2.imread(p, cv2.IMREAD_COLOR) img cv2.resize(img, (512, 512)) inp torch.from_numpy(img).float().permute(2, 0, 1).unsqueeze(0) / 255.0 with torch.no_grad(): out model(inp) edge torch.sigmoid(out[-1]).squeeze().cpu().numpy() edge (edge * 255).astype(np.uint8) cv2.imwrite(os.path.join(save_dir, f{idx:05d}.png), edge)逻辑说明批量脚本里没有放置归一化到 0-255 的步骤是因为输出概率图直接乘 255 再存就行合成到下游任务前再做量化。6.2 清洗边缘图的三个步骤模型输出的边缘图不是拿来就能用的尤其做数据集时会发现边缘有断点和背景噪声。我的清洗流程固定三步中值滤波、闭运算、按连通域面积过滤。def clean_edge(prob_path): img cv2.imread(prob_path, cv2.IMREAD_GRAYSCALE) img cv2.medianBlur(img, 5) kernel np.ones((3, 3), np.uint8) img cv2.morphologyEx(img, cv2.MORPH_CLOSE, kernel) num_labels, labels, stats, _ cv2.connectedComponentsWithStats(img, connectivity8) out np.zeros_like(img) for i in range(1, num_labels): if stats[i, cv2.CC_STAT_AREA] 100: continue # 去掉面积小于100的小噪声块 out[labels i] 255 return out参数说明面积阈值 100 在 512x512 分辨率下基本能去掉零散噪声点如果边缘图更宽可以增大到 200阈值太小会把细碎边缘也一并删除需要根据自己图像的边缘密集程度调整。6.3 清洗后的边缘图怎么接下游任务边缘概率图可以直接作为分割模型的额外输入通道也可以用来做目标检测的 ROI 裁剪先验。我习惯把清洗后的二值边缘图和原图拼成四通道输入再喂给分割网络这样网络既能看原始纹理又能看到明确的结构先验。这个技巧在医疗图像和工业缺陷检测里效果更明显因为边缘往往是缺陷的边界。从那以后我每跑完一个边缘检测模型都会先做一次批处理边缘图清洗再看效果决定要不要往后面接。这一步看着不起眼省下的调试时间却不少。希望帮到你。本文还有配套的精品资源点击获取