概述
文档图像的边缘校正是图像处理中的一项重要任务，尤其在文档数字化和自动化扫描过程中，确保文档图像的几何形状和内容准确性具有重要意义。传统的文档图像校正方法通常依赖于手动选择或简单的几何变换。然而，随着深度学习和计算机视觉技术的发展，语义分割与直线检测被广泛应用于文档边缘校正任务中，极大提高了处理的效率和精度。
语义分割技术通过将图像中的每个像素进行分类，可以精准地提取出文档的边缘信息。在文档图像中，文档区域通常会被标记为一个特定的类别，而背景或其他元素则被分离出来。直线检测通常用于检测文档的直线边缘，尤其是四个角的直线，这些直线在文档边缘校正中起到了关键作用。通过使用霍夫变换等算法，可以在文档图像中准确提取出边缘的直线，并且利用这些直线信息计算出文档边缘的四个角，从而进行文档的透视校正，使得拍摄的文档图像变得规整，避免了因拍摄角度偏差导致的几何畸变。
1. 为什么使用深度学习
在之前博客中，有总结过使用传统数字图像处理方法实现拍照文档边缘校正。用传统方法实现文档边缘校存在许多局限性。传统方法通常依赖于边缘检测（如Canny算子）和霍夫变换检测直线，通过这些步骤定位文档的边缘。这种方法的精度和效果依赖于图像质量、光照条件和噪声等因素，因此在不同场景中需要不断地调整参数，而泛化能力有限，难以应对复杂背景或不规则形状的文档。
相比之下，基于深度学习的方法虽然训练模型时需要大量带标注的数据和较长的时间，但在部署后其精度高、适用性强，并且能够处理多样化的场景，包括各种角度、光照和背景的复杂文档。这种方法的鲁棒性和适应性在实际应用中能明显提升用户体验。
2. 选择模型
ENet（Efficient Neural Network）是一种专为实时语义分割任务设计的深度神经网络架构。ENet 的网络结构设计简洁，注重减少模型的参数数量和计算复杂度。虽然 ENet 和 SegNet 都采用编码-解码结构，但 ENet 采用了一个大的编码器（Encoder）和一个小的解码器（Decoder）的不对称架构，大大减少了模型的参数量。这种设计哲学基于一个观点：解码器的主要作用是对编码器的输出进行上采样，并微调细节，而这并不需要一个复杂的结构。而 SegNet 的架构相对对称，包含大量的参数。
ENet 还通过早期下采样、分解卷积核等策略，有效降低了计算复杂度，减少了浮点运算量，从而实现更快的推理速度。此外，ENet 在下采样过程中使用了空洞卷积，实现在不降低特征图分辨率的同时扩大图像目标的感受野，获取更广泛的上下文信息，更好地理解图像中物体与周围环境的关系，从而提高分割精度，尤其适用于复杂场景下对相似类别物体的区分。
现有语义分割模型（例如 SegNet 和全卷积网络（FCN）），大多数都基于 VGG16 架构，而 VGG16 是为多类别分类任务设计的超大模型。这些网络拥有巨大的参数量和较长的推理时间，在实际应用中无法满足许多移动或电池供电设备的需求，这些应用要求图像处理速率高于 10帧/秒（fps）。而 ENet 的速度提高了 18 倍，浮点运算量减少了 75%，参数量减少了 79%，同时在准确率上提供了相似或更好的表现。这使得 ENet 能够在存储和计算资源有限的设备上更易于部署和运行，实现快速的推理和实时的语义分割。因此，本文选择 ENet 网络分割文档。
论文地址：https://arxiv.org/abs/1606.02147
2.1 网络架构
ENet 的网络结构如下表格。
initial 阶段包含一个模块，如下图所示。第 1 阶段包含 5 个 bottleneck 模块，而第 2 阶段和第 3 阶段的结构相同，但第3阶段在开始时不进行下采样（因此省略了第 2 阶段的第0个 bottleneck 模块）。这前三个阶段组成了编码器。第 4 阶段和第 5 阶段则属于解码器。
ENet 的网络结构借鉴了 ResNet ，描述为一个主分支和一个带有卷积核的附加分支，这些附加分支从主分支分离，然后最后通过逐像素相加与主分支进行融合。每个 bottleneck 模块由三个卷积层组成：1 × 1 卷积层，用于降低维度；主卷积层（卷积核的尺寸为 3 × 3。有时，会将其替换为非对称卷积，即一系列 5 × 1 和 1 × 5 卷积的组合。）；1 × 1 卷积层，用于恢复维度。在所有卷积层之间加入了批量归一化和 PReLU 激活函数 。如果 bottleneck 模块执行下采样操作，则在主分支中会额外添加一个最大池化层。此外，将第一个 1 × 1 卷积层替换成一个卷积核为 2 × 2 、步幅为 2 的卷积层，对激活值进行零填充（zero padding），以匹配特征图的数量。
在解码器中，最大池化（MaxPooling）被替换为最大反池化（Max Unpooling），填充（Padding）则被替换为无偏置的空间卷积。在最后的上采样模块中也就是第 5 阶段，没有使用池化索引（因为 initial 模块操作的是输入图像的 3 个通道，而最终的输出是具有 C个特征图（即目标类别的数量）。此外，出于性能考虑，将全卷积作为网络的最后一个模块，这个模块本身占据了解码器处理时间的较大部分。
3. 数据集准备
笔者使用 labelme 进行文档标注。
3.1 labelme安装
conda create --name labelme python=3.8
conda activate labelme
pip install pyqt5 pillow
pip install labelme
安转完成之后输入labelme命令即可打开labelme，界面如图所示。
3.2 数据标注
标注完成后，得到的对应的标签信息是以JSON格式保存的，里面保存着物体的位置信息。
json文件里面的内容：
3.3 数据处理
根据图像和JSON文件，需要生成标签图像，像素值是类别的索引（例如，本数据集背景是0，目标是1）。虽然看起来标签图像是黑色的，但其实并不是纯黑，0和1之间有微小差距。其次，需要生成生成train.txt和val.txt，分别表示训练集和验证集。
import os
import sys
import glob
import json
import math
import uuid
import random
import numpy as np
import PIL.Image
import PIL.ImageDraw
from tqdm import tqdm
def shape_to_mask(img_shape, points, shape_type=None,
line_width=10, point_size=5):
mask = np.zeros(img_shape[:2], dtype=np.uint8)
mask = PIL.Image.fromarray(mask)
draw = PIL.ImageDraw.Draw(mask)
xy = [tuple(point) for point in points]
if shape_type == 'circle':
assert len(xy) == 2, 'Shape of shape_type=circle must have 2 points'
(cx, cy), (px, py) = xy
d = math.sqrt((cx - px) ** 2 + (cy - py) ** 2)
draw.ellipse([cx - d, cy - d, cx + d, cy + d], outline=1, fill=1)
elif shape_type == 'rectangle':
assert len(xy) == 2, 'Shape of shape_type=rectangle must have 2 points'
draw.rectangle(xy, outline=1, fill=1)
elif shape_type == 'line':
assert len(xy) == 2, 'Shape of shape_type=line must have 2 points'
draw.line(xy=xy, fill=1, width=line_width)
elif shape_type == 'linestrip':
draw.line(xy=xy, fill=1, width=line_width)
elif shape_type == 'point':
assert len(xy) == 1, 'Shape of shape_type=point must have 1 points'
cx, cy = xy[0]
r = point_size
draw.ellipse([cx - r, cy - r, cx + r, cy + r], outline=1, fill=1)
else:
assert len(xy) > 2, 'Polygon must have points more than 2'
draw.polygon(xy=xy, outline=1, fill=1)
mask = np.array(mask, dtype=bool)
return mask
def shapes_to_label(img_shape, shapes, label_name_to_value):
cls = np.zeros(img_shape[:2], dtype=np.int32)
ins = np.zeros_like(cls)
instances = []
for shape in shapes:
points = shape['points']
label = shape['label']
group_id = shape.get('group_id')
if group_id is None:
group_id = uuid.uuid1()
shape_type = shape.get('shape_type', None)
cls_name = label
instance = (cls_name, group_id)
if instance not in instances:
instances.append(instance)
ins_id = instances.index(instance) + 1
cls_id = 1
# cls_id = label_name_to_value[cls_name]
mask = shape_to_mask(img_shape[:2], points, shape_type)
cls[mask] = cls_id
ins[mask] = ins_id
return cls, ins
def lblsave(filename, lbl):
if os.path.splitext(filename)[1] != '.png':
filename += '.png'
# Assume label ranses [-1, 254] for int32,
# and [0, 255] for uint8 as VOC.
if lbl.min() >= 0 and lbl.max() <= 255:
lbl_pil = PIL.Image.fromarray(lbl.astype(np.uint8), mode='L')
lbl_pil.save(filename)
else:
raise ValueError(
'[%s] Cannot save the pixel-wise class label as PNG. '
'Please consider using the .npy format.' % filename
)
if __name__ == '__main__':
data_path = sys.argv[1]
out_path = sys.argv[2]
if not os.path.exists(out_path):
os.makedirs(out_path)
label_name_to_value = {
'_background_': 0,
'a': 1,
}
json_fns = glob.glob(os.path.join(data_path, '**/*.json'), recursive=True)
out_lst = []
for json_fn in tqdm(json_fns):
with open(json_fn, 'r') as f:
data = json.load(f)
img_shape = (data['imageHeight'], data['imageWidth'])
lbl, _ = shapes_to_label(img_shape, data['shapes'], label_name_to_value)
image_fn = json_fn.replace('.json', '.jpg')
label_fn = json_fn.replace('.json', '_label.png')
if not os.path.exists(image_fn):
print(image_fn + ' not exists')
continue
# else:
# img = PIL.Image.open(image_fn)
# mask = PIL.Image.open(label_fn)
# if img.size != mask.size:
# print(image_fn, img.size, mask.size)
# continue
lblsave(label_fn, lbl)
out_lst.append(image_fn + ',' + label_fn)
random.shuffle(out_lst)
trn_num = int(len(out_lst) * 0.9)
with open(os.path.join(out_path, 'train.txt'), 'w') as f:
f.write('\n'.join(out_lst[:trn_num]))
with open(os.path.join(out_path, 'val.txt'), 'w') as f:
f.write('\n'.join(out_lst[trn_num:]))
运行上述代码：
python generate_label.py E:/data/ E:/data/dataset
参数说明：第一个参数为图像和JSON目录，第二个参数为输出文件夹
4. 模型训练
代码地址：https://github.com/davidtvs/PyTorch-ENet
4.1环境安装
conda create --name enet python=3.8
conda activate enet
conda install pytorch torchvision cudatoolkit=10.2 -c pytorch
pip install cython matplotlib tqdm opencv-python scipy pillow
或
conda create --name enet python=3.8
conda activate enet
conda install pytorch torchvision cudatoolkit=10.2 -c pytorch
pip install -r requirements.txt
4.2 模型训练
python main.py -m train --save-dir save/ENet_document --name ENet --dataset document --dataset-dir E:/data/dataset --epochs 100 --height 512 --width 512 --print-step
参数说明：
- – save-dir：训练模型保存路径
- –dataset：训练数据类型
- –dataset-dir：train.txt和val.txt所在路径
- –epochs：训练总epoch数
- –height：输入高度
- –width：输入宽度
- –print-step：是否打印每个step的loss
注意：如果需要训练自己的数据集，需要改一下代码。例如，本文将数据集划分为训练集和验证集，因此需要修改自制的数据类型（document），在 args.py 中 mode 的修改，在 main.py 中有些需要把 test 改成 val 。官方代码，mode 模式只有 train、test。
4.3 模型转换
使用以下代码将训练得到的模型转换为ONNX模型。
python convert_to_onnx.py --input save/ENet_document/ENet --output save/ENet_document/ENet.onnx
import os
import argparse
from addict import Dict
import yaml
import torch
import torch.optim as optim
from models.enet import ENet
import utils
def get_args():
parser = argparse.ArgumentParser(description='convert pytorch weights to onnx')
parser.add_argument('--input', type=str,required=True, help="input pytorch model")
parser.add_argument('--output', type=str, required=True, help="output onnx model")
args = parser.parse_args()
return args
if __name__ == '__main__':
args = get_args()
model = ENet(2, encoder_relu=True).to('cpu')
save_dir = os.path.dirname(args.input)
name = os.path.basename(args.input)
optimizer = optim.Adam(model.parameters())
model = utils.load_checkpoint(model, optimizer, save_dir, name)[0]
inputs = torch.FloatTensor(1, 3, 512, 512)
torch.onnx.export(model, inputs, args.output, opset_version=9)
5. 模型推理
5.1 NCNN
NCNN 是由腾讯开源的一个高性能的神经网络前向推理框架，主要用于在资源受限的移动端和嵌入式设备（如安卓、iOS，以及边缘设备等）上高效地运行深度学习模型。
NCNN 的模型由两个文件组成：.param 和.bin。.param 文件包含模型结构，包括层类型、层名称、层参数和输入/输出形状。这个文件通常是一个轻量级的文本文件，可以很容易地被 NCNN 框架解析和加载。.bin 文件包含二进制格式的训练模型权重。这些权重是在训练过程中学习的，并用于在推理过程中进行预测。.bin 文件通常比 .param 文件大得多，因为它包含学习参数的实际数值。
NCNN 支持多种移动操作系统，包括 Android 和 iOS。这使得开发者可以在不同的移动平台上使用相同的深度学习模型，而无需对模型进行大量的修改。例如，一个在 Android 设备上训练好的图像分类模型，可以很容易地通过 NCNN 移植到 iOS 设备上进行应用。它支持将主流深度学习框架（如 PyTorch、TensorFlow、ONNX）训练的模型转换为 ncnn 格式。可以使用 ncnn 提供的工具将训练好的模型（例如 .onnx 格式）转换为 ncnn 支持的 param 和 bin 文件。
针对移动设备的计算资源（如 CPU、GPU 等）进行了优化。它采用了许多优化策略来减少计算量和内存占用。它还支持模型量化，通过将模型参数从高精度（如 32 位浮点数）转换为低精度（如 8 位整数），可以在不显著降低模型精度的情况下，大幅减少模型的存储空间和计算量。可以针对不同硬件设备调优参数，例如线程数、是否开启 FP16 加速等。对 CPU 进行了优化，支持 SIMD 指令集，如 NEON (ARM) 和 SSE/AVX (x86)。它还支持 Vulkan 加速，可以利用 GPU 提升推理速度。
5.2 环境配置
5.2.1 NCNN下载
可以从这个地址https://github.com/Tencent/ncnn 获取源码进行编译，也可以下载官方编译好的 lib。本文直接下载编译好的 ncnn。笔者的系统是 Windows11，VS 版本 2019，cuda 版本 11.7，cudnn 版本 8.5。下载的文件包含 ncnn 和 glslang 。
在 NCNN 框架中使用 Glslang 主要是为了在 GPU 上执行高效的着色器计算，尤其是通过 Vulkan 后端进行模型推理。NCNN 使用 Vulkan 作为 GPU 后端来实现高性能的模型推理，而使用 SPIR-V 才能够最大限度发挥 Vulkan 的并行计算能力，从而在 GPU 上加速神经网络模型的推理过程。在 NCNN 中，部分算子（如卷积、激活、矩阵乘法等）是通过 GLSL 着色器代码实现的。Glslang 工具将 GLSL 着色器代码编译成 SPIR-V，确保能够在 Vulkan 设备上正确执行。
下载完成后，解压即可。
5.2.2 Vulkan下载
NCNN 使用 Vulkan 作为 GPU 后端来加速神经网络模型的推理过程。Vulkan 可以在这个地址https://vulkan.lunarg.com/sdk/home下载对应的版本。笔者下载的是1.2.141.2。
5.2.3 OpenCV下载
笔者下载的opencv 版本为 4.7.0 ，下载链接为 https://opencv.org/releases/。
5.2.4 属性配置
安装完成之后在项目工程中配置 ncnn、vulkan、opencv。
并将三者的 bin 文件中的 dll文件复制到工程的 release 文件下。
5.3 模型转换
方式1：ONNX2NCNN一键转换工具，地址：https://convertmodel.com。这是一个在线的模型转换网站。建议先尝试方式1，失败再尝试方式2。
方式2：可以从https://github.com/Tencent/ncnn 获取源码进行编译，也可以下载官方编译好的lib进行转换。笔者直接下载官方编译好的进行转换。
onnx2ncnn E:/py/ENet/save/ENet_Card/ENet.onnx ENet.param ENet.bin
注意：转换的 ncnn 模型，需要与属性配置的 ncnn 版本一致。
5.4 推理流程
头文件：
#include <ncnn/net.h> // ncnn 推理网络的核心头文件
#include <opencv2/opencv.hpp> // 如果需要预处理图像，使用 OpenCV
（1）初始化网络并加载模型
ncnn 使用 .param 和 .bin 文件来加载网络结构和权重。
// Step 1: 初始化 ncnn 网络
ncnn::Net net;
// 可选：设置网络优化选项
if (has_gpu)
{
net.opt.use_vulkan_compute = true; // 如果设备支持 Vulkan，可以开启 GPU 加速，false则使用 CPU 进行推理。
std::cout << "Vulkan GPU compute enabled." << std::endl;
}
net.opt.use_fp16_arithmetic = true; // 使用 FP16 进行加速计算
// Step 2: 加载网络结构和权重
int ret_param = net.load_param("ENet.param");
int ret_model = net.load_model("ENet.bin");
if (ret_param != 0 || ret_model != 0)
{
std::cerr << "Error: Failed to load model files." << std::endl;
return -1;
}
如果使用 GPU 进行推理加速，可以先检查设备是否存在支持 Vulkan 的 GPU。可以使用以下代码进行检查。
// Step 1: 初始化 GPU 实例
ncnn::create_gpu_instance();//初始化 GPU 计算实例，必须在使用 GPU 功能前调用。
// Step 2: 检查 GPU 是否可用
bool has_gpu = ncnn::get_gpu_count() > 0;
if (has_gpu)
{
std::cout << "GPU is available! GPU count: " << ncnn::get_gpu_count() << std::endl;
}
else
{
std::cerr << "No GPU found! Falling back to CPU." << std::endl;
}
（2）图像预处理
输入图像的形状必须与训练时的输入形状一致。可以使用 OpenCV 读取图片并进行归一化、尺寸缩放等预处理。
// 读取输入图像并预处理
cv::Mat image = cv::imread("input.jpg");
if (image.empty())
{
std::cerr << "Error: Could not load input image." << std::endl;
return -1;
}
// 调整输入图像尺寸，例如网络输入是 512x512
cv::Mat resized_image;
cv::resize(image, resized_image, cv::Size(512, 512));
// 将输入图像数据转换为 ncnn 所需的格式并调整图像大小
ncnn::Mat input = ncnn::Mat::from_pixels_resize(
resized_image.data, ncnn::Mat::PIXEL_BGR,
resized_image.cols, resized_image.rows, input_width, input_height);
// 归一化 [0,1]
const float norm_vals[3] = { 1 / 255.f, 1 / 255.f, 1 / 255.f }; // 归一化因子
input.substract_mean_normalize(0, norm_vals);
ncnn::Mat::from_pixels_resize 用于将输入图像数据转换为 ncnn 所需的格式，并在过程中调整图像大小。
参数说明：
- 第一个参数，指向输入图像数据的指针。通常为 OpenCV 中 cv::Mat 的 .data 成员。图像数据需按照行优先存储。
- 第二个参数，输入图像的通道类型。
- ncnn::Mat::PIXEL_RGB：3 通道 RGB 图像。
- ncnn::Mat::PIXEL_BGR：3 通道 BGR 图像（OpenCV 默认）。
- ncnn::Mat::PIXEL_GRAY：1 通道灰度图像。
- 第三四，输入图像的原始宽度和高粗。
- 第五六，输出图像的目标宽度和高度。
（3）执行前向推理
使用 Extractor 提取网络输出。
// Step 1: 创建 Extractor 并设置输入
ncnn::Extractor extractor = net.create_extractor();
extractor.set_light_mode(true); // 启用轻量模式
extractor.set_num_threads(4); // 设置线程数，提升性能
// 设置输入（名称和网络输入绑定）
extractor.input("input", input);// "input" 是网络输入的节点名
// Step 2: 执行前向推理并获取输出
ncnn::Mat output;
int ret = extractor.extract("output", output); // "output" 是网络输出的节点名
if (ret != 0)
{
std::cerr << "Error: Failed to extract output." << std::endl;
return -1;
}
注意：“input” 和 “output” 是网络的输入和输出节点名。
如何获取这些名称？？？
- 在 ONNX 导出的模型中，可以用工具（例如 Netron）查看输入和输出节点的名称。
- 在 ncnn 网络的 .param 文件中，可以直接看到输入和输出的节点名称。
（4）解析输出结果
ncnn 输出的结果是一个 ncnn::Mat 对象。
cv::Mat cv_seg = cv::Mat::zeros(cv::Size(output.w, output.h), CV_8UC1);
for (int i = 0; i < output.h; ++i)
{
for (int j = 0; j < output.w; ++j)
{
const float* bg = output.channel(0);// 背景通道的概率
const float* fg = output.channel(1);// 前景通道的概率
if (bg[i * output.w + j] < fg[i * output.w + j])//比较背景与前景的概率，若前景概率更高，将该像素标记为 255（白色）
{
cv_seg.data[i * output.w + j] = 255;
}
}
}
cv::resize(cv_seg, cv_enet, cv::Size(cv_src.cols, cv_src.rows), cv::INTER_LINEAR);//将分割图像 cv_seg（掩码） 重新 resize 到原始图像大小
// 将掩码图像与原始图像结合，提取前景区域
cv::Mat result;
image.copyTo(result, cv_enet); // 将 cv_image 中掩码为 255 的区域复制到 result 中
（5）释放资源
ncnn 会自动清理资源，但可以显式地释放网络对象。
net.clear();//释放网络对象
ncnn::destroy_gpu_instance();//释放 GPU 实例
5.5 推理测试
掩码图像：
分割图像：
6. 项目实现
（1）图像预处理：对拍摄的文档图像调整图像大小，满足模型所需要的输入大小。
（2）语义分割：使用 NCNN 对训练得到的ENet模型进行推理，识别出文档的轮廓。
（3）绘制轮廓：对掩码图像进行轮廓提取，然后计算每个轮廓的面积，并且在掩码图像上绘制面积最大的轮廓。
（3）直线检测：利用霍夫变换对轮廓图像进行直线检测。如果检测到的直线小于文档的边缘线数目，则对掩码图像进行边缘检测、轮廓提取、直线检测等操作。将检测到的直线划分为水平直线、垂直直线，并对线段进行筛选，最终得到文档边缘的四条线。
（4）计算交点：计算文档边缘四条线的交点，并将四个角点映射到原始图像。
（5）几何变换：根据检测到的文档边缘直线，应用透视变换（如仿射变换、投影变换等），将文档校正为标准矩形形状。通过计算长边和短边，确定目标矩形的宽和高，保证变换后的比例正常。
