深度学习边缘检测实战：HED与PiDiNet源码解析及PyTorch部署指南

简介：一份面向计算机、人工智能、数据科学等相关专业学生与初学者的边缘检测实践项目，基于深度学习完成轮廓提取任务，内含HED、PiDiNet等经典模型的Python源码、预训练权重与配套数据集，可完整复现训练与推理流程，尤其适合毕业设计、课程设计、项目演示及入门进阶。压缩包共三十四个文件，体积约8.72MB，主要类型包括源码（py）、模型权重（pth）、样本图片（jpg/png）、标注文件（xml）和使用说明（md/txt）；图片与标注配合可构建训练样本，模型权重可直接加载推理，说明文档辅助环境搭建与参数调整。目前已有六十七人浏览学习。项目附带运行说明与调试输出参考图，便于对照验证检测效果；环境配置完成后可快速运行，也可基于现有结构改造网络或迁移至其他图像检测场景，二次开发空间较大。

1. 能直接跑起来的深度学习边缘检测：这份 HED+PiDiNet 源码包解决什么问题

边缘检测在视觉项目里卡过不少人：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 只需要最后一层计算 loss，HED 在每个侧输出层都单独计算 loss，等于同时让网络浅层学习局部细节、深层学习整体轮廓。训练时侧输出权重可以调节，推理时既可以用单独的侧输出，也可以用融合后的结果。

我拆这份源码的时候特意看了hed_edge.py里的模型结构，骨干部分基本遵循 VGG16 的分类层截断，但后面接的侧输出层是卷积加反卷积的组合。由于感知野不同，浅层侧输出边缘细但噪声多，深层侧输出边缘连续但定位粗，融合以后两者互补。这是 HED 能扛住复杂场景的根本原因。

2.3 PiDiNet 用像素差卷积换轻量化

PiDiNet 是另一个方向的代表：不追求大网络，而是用像素差卷积 PDC（Pixel Difference Convolution）去模拟传统梯度算子。它的核心思想是把像素差值当作卷积特征的一部分，既保留了局部灰度变化信息，又不需要像 VGG 那样堆叠超大卷积核。

和 HED 相比，PiDiNet 参数量小得多，推理速度快，很适合做实时应用。发行版里table5_pidinet.pth这个命名方式应该是参考论文 Table 5 的配置训练出来的权重，这类权重对输入分辨率有一定要求，跑之前最好先看一眼模型定义里的输入尺寸，避免直接拿高分辨率图硬塞。

这份资源同时给出 HED 和 PiDiNet，等于给你两条路线：一条重精度、一条重速度。毕设里做算法对比时，这两个模型正好可以形成对照。

2.4 解压后文件结构对应哪些环节

解压后第一眼会看到一堆.py、.pth、图片和debug_out开头的 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 python=3.8 conda activate edge_det pip install torch==1.13.1 torchvision==0.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_location="cpu") 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_location="cpu"表示即使用没有 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", encoding="utf-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.txt

train.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_size=256): 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(), lr=1e-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(f"epoch {epoch}, step {i}, loss {loss.item():.4f}")

逻辑说明：HED 的多个侧输出都要和同一张边缘标注计算损失，所以循环里累加每个侧输出的损失再求平均。BCEWithLogitsLoss比手动在输出层接 Sigmoid 再算二值交叉熵更稳定，因为它在内部做了数值处理。

参数说明：学习率1e-4是边缘检测里比较常用的起点，Adam 优化器对初始值不那么敏感。crop_size=256控制原图和标注同时裁剪的尺寸，越小训练越快但边缘定位越粗糙。

4.3 训练参数怎么给：学习率、batch size 与关键注意点

用一张参数表说明我平时调训练时的配置：

有一个容易被忽略的点：边缘标注里边缘像素远少于背景像素，直接算交叉熵会让模型倾向把全部预测成背景。HED 论文里用了类别平衡策略，也就是在 loss 计算时给边缘像素更高的权重。常见的简化做法是在BCEWithLogitsLoss里传入pos_weight。

criterion = nn.BCEWithLogitsLoss(pos_weight=torch.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_location="cpu") 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读取并先开numpy：

from 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_ok=True) 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, connectivity=8) 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 裁剪先验。我习惯把清洗后的二值边缘图和原图拼成四通道输入，再喂给分割网络，这样网络既能看原始纹理，又能看到明确的结构先验。这个技巧在医疗图像和工业缺陷检测里效果更明显，因为边缘往往是缺陷的边界。

从那以后，我每跑完一个边缘检测模型，都会先做一次批处理边缘图清洗，再看效果决定要不要往后面接。这一步看着不起眼，省下的调试时间却不少。希望帮到你。

本文还有配套的精品资源，点击获取