文档矫正的本质,是一个 3D 重建 问题。
你需要从一张 2D 的照片中,推断出纸张在 3D 空间中的弯曲形状,然后把它“展平”。
在 2019 年之前,这是传统 CV(计算机视觉)的噩梦。
现在,我们有两把新武器:
- 几何矫正 (Geometric Unwarping):利用 Transformer 或 CNN 预测像素的位移场(Flow Field)。
- 光照处理 (Illumination Correction):利用 GAN 去除阴影和反光。
1. 核心原理:后向映射 (Backward Mapping)
深度学习模型(如 DocUNet, DewarpNet, DocTr)不会直接生成一张新图。它生成的是一个 映射网格 (Grid / Flow Map)。
- 输入:一张弯曲的图 $I_{src}$ (256×256)。
- 模型输出:一个坐标映射表 $M$ (256x256x2)。
- 即:对于输出图上的坐标 $(x, y)$,它在原图上的位置是 $(u, v)$。
有了这个 $M$,我们就可以用 grid_sample 插值算法,把原图的像素“搬运”到新的位置,从而实现展平。
2. 架构演进:从 UNet 到 Transformer
- 1.0 时代:DocUNet
- 使用类似 UNet 的结构,直接回归 3D 坐标。
- 缺点:边缘处理不好,容易产生锯齿。
- 2.0 时代:DocTr (Document Transformer)
- 这是目前的 SOTA(State of the Art)方向。
- 它引入了 Transformer 来捕捉长距离依赖。因为纸张的弯曲是连续的,左上角的弯曲程度和右下角是有物理关联的。Transformer 的 Self-Attention 机制能完美建模这种物理约束。
3. 代码实战:实现一个简单的 Dewarping 推理
在工程落地中,最常用的并不是沉重的 DocTr,而是轻量级的 DewarpNet 或基于 TPS (Thin Plate Spline) 的方案。
这里演示如何利用 PyTorch 的 grid_sample 实现矫正逻辑。假设模型已经预测出了 flow_field。
Python
import torch
import torch.nn.functional as F
import cv2
import numpy as np
def unwarp_image(distorted_img, flow_field):
"""
distorted_img: 原始弯曲图片 [1, C, H, W]
flow_field: 模型预测出的位移场 [1, H, W, 2], 值域归一化在 [-1, 1] 之间
"""
# 核心魔法:grid_sample
# 它会根据 flow_field 指定的坐标,去 distorted_img 里采样像素
# align_corners=True 是关键,否则边缘会对不齐
rectified_img = F.grid_sample(distorted_img, flow_field, mode='bilinear', padding_mode='zeros', align_corners=True)
return rectified_img
# 伪代码:完整的推理流程
def pipeline(image_path, model):
# 1. 预处理
img_cv = cv2.imread(image_path)
h, w = img_cv.shape[:2]
# 关键工程技巧:低分预测,高分采样
# 模型跑在 256x256 上(为了快)
img_resized = cv2.resize(img_cv, (256, 256))
input_tensor = preprocess(img_resized) # 转 Tensor, Normalize
# 2. 模型预测 Backward Mapping
with torch.no_grad():
# 输出通常是 [1, 256, 256, 2]
predicted_grid = model(input_tensor)
# 3. 上采样 Grid (Upsample)
# 将 256x256 的 Grid 放大回原图尺寸 (比如 1920x1080)
# 这一步比直接在原图上跑模型快 100 倍
full_res_grid = F.interpolate(
predicted_grid.permute(0, 3, 1, 2),
size=(h, w),
mode='bilinear'
).permute(0, 2, 3, 1)
# 4. 应用矫正
full_res_img_tensor = torch.from_numpy(img_cv).permute(2, 0, 1).unsqueeze(0).float() / 255.0
result = unwarp_image(full_res_img_tensor, full_res_grid)
return result
4. 解决光照:GAN 的魔法
把纸展平了,但纸上还有拍照时的阴影(Shadow)和褶皱的黑印。这时候 OCR 还是会挂。
你需要 二值化(Binarization),但 cv2.adaptiveThreshold 太蠢了,它会把阴影里的字也抹掉。
现在的做法是使用 GAN (Generative Adversarial Networks),比如 DocEnTr。
- Generator:负责把“有阴影的图”生成“干净的白纸黑字图”。
- Discriminator:负责判断生成的图是不是像真的扫描件。
在移动端(iOS/Android),我们通常使用轻量级的 GAN 变体(如 Pix2Pix 的剪枝版),专门训练去除手机摄像头的中心光斑和边缘阴影。
5. 工程落地的坑与优化
- 边缘检测是生门:
- 所有的 Dewarp 模型都极度依赖 文档边缘(Document Boundary)。如果背景太乱,或者边缘被遮挡,模型预测的 Grid 就会飞出去。
- 解法:先用一个轻量级语义分割模型(如 DeepLabV3-Mobile)把文档抠出来(Crop),并在周围填充黑色 Padding,再送入 Dewarp 模型。
- 分辨率陷阱:
- 不要试图训练一个能直接处理 4K 图片的 Dewarp 模型,显存会爆炸。
- 解法:始终坚持 “低分预测 Flow,高分应用 Sample” 的策略。Flow Field 是平滑的,放大后插值误差很小。
- 数据哪里来?
- 去哪找成对的“弯曲图”和“平整图”?实拍几乎不可能对齐。
- 解法:使用 Doc3D 数据集。这是一个用 Blender 渲染出来的合成数据集,包含了 10 万张各种弯曲形态的文档,且带有完美的 Ground Truth 坐标。
总结
如果你的 App 是做 “扫描全能王” 竞品的,或者是做 试卷录入 的:
- 放弃 OpenCV 的透视变换,它处理不了书页弯曲。
- 寻找基于 Grid Regression 的模型(关键词:DocUNet, DewarpNet)。
- 部署策略:在端侧(手机 NPU)跑一个小的分割模型 + Flow 预测模型(256×256),然后利用 GPU shader 或 NPU 进行高分辨率的
grid_sample。
这就是为什么扫描全能王能把你拍得像咸菜一样的纸,变成像 PDF 一样平整的魔法。不是魔法,是 Transformer + Grid Sampling。