DocRes 是 CVPR 2024 论文提出的统一文档图像复原通用模型，一个 Restormer 网络即可完成去弯（Dewarping）、去阴影（Deshadowing）、外观增强（Appearance）、去模糊（Deblurring）和二值化（Binarization）五大任务。本文面向新手，手把手讲清DocRes 文档图像复原训练的完整流程：如何准备 Doc3D、TDD、FSDSRD、RealDAE 等数据集、如何编辑datasets_setting、如何用 8 卡分布式启动训练，最后用eval.py复现论文中的官方评测结果。

一、模型与训练代码速览

在开始训练前，先了解 DocRes 的核心结构，有助于理解后面的数据准备要求：

训练输入不是"裸图"：每个任务都会在线计算一个 3 通道提示（如 Sobel 高频图、阴影背景图、Sauvola 二值化结果），与输入图拼接后送入网络。提示的计算逻辑在 loaders/docres_loader.py 中。

二、环境准备：安装 PyTorch 与依赖

建议新建 Python 3.8 虚拟环境，安装项目声明的依赖（锁定 torch 1.11 + CUDA 11.3）：

# 克隆仓库（如尚未克隆） git clone https://gitcode.com/gh_mirrors/do/DocRes cd DocRes # 安装依赖 pip install -r requirements.txt

requirements.txt 中包含torch==1.11.0+cu113、opencv-python-headless、scikit-image、einops等。若显卡 CUDA 版本不同，只需替换为对应版本的 PyTorch，其余依赖无需改动。

三、数据集准备：Doc3D、TDD 等数据组织

训练数据目录必须遵循 data/README.md 规定的结构：

data/ ├── eval/ # 评测集：dir300 / kligler / jung / osr / realdae / docunet_docaligner / tdd / dibco18 │ └── realdae/1_in.png + 1_gt.png └── train/ ├── dewarping/doc3d/ # 去弯：Doc3D ├── deshadowing/{fsdsrd,tdd}/ # 去阴影：FSDSRD + TDD ├── appearance/{clean_pdfs,realdae}/ # 外观增强 ├── deblurring/tdd/ # 去模糊：TDD └── binarization/{bickly,dibco,noise_office,phibd,msi}/ # 二值化

每个任务还需要一份 JSON 清单，描述输入图、真值图（以及提示图）的相对路径。

3.1 去弯：Doc3D 的 mask 与背景准备

Doc3D 提供带 UV 数据的弯曲书页扫描图，训练时要做两件事：

1. 提取 mask：从 Doc3D 的 uv 数据中为每张图提取页面 mask；
2. 准备背景图：用于在 mask 外填充随机背景，增强模型对"页面以外区域"的鲁棒性，需在 loaders/docres_loader.py 中把self.background_paths指向背景文件夹。

Doc3D 中的页面 mask 示意（白区域为页面）：

JSON 需要同时给出in_path、mask_path、gt_path（gt 为 .npy 格式的基准映射）：

[ { "in_path": "dewarping/doc3d/img/1/102_1-pp_Page_048-xov0001.png", "mask_path": "dewarping/doc3d/mask/1/102_1-pp_Page_048-xov0001.png", "gt_path": "dewarping/doc3d/bm/1/102_1-pp_Page_048-xov0001.npy" } ]

3.2 去阴影 / 去模糊：TDD、FSDSRD 与 RDD

• 去阴影：使用 FSDSRD 与 RDD 数据集，JSON 只需in_path和gt_path；
• 去模糊：使用 TDD 数据集，格式与去阴影相同。

带阴影的文档训练样例（DocRes 去阴影任务的典型输入）：

3.3 外观增强：Doc3DShade 影子提取

外观增强任务的核心是 Doc3DShade 合成数据，流程为：

1. 收集网络 PDF 并转成图像（appearance/clean_pdfs/），作为干净的 gt；
2. 运行 data/preprocess/shadow_extraction.py 提取"影子图"：先利用 mask 把拍摄图（cap_im）与 PDF 渲染图（alb_im）分别去弯，再由两者相除得到影子；
3. 在 loaders/docres_loader.py 中将self.shadow_paths指向影子图文件夹。

原始弯曲拍摄书页与对应 PDF 渲染图：

训练时 loader 会从影子库随机取一张影子乘到干净 PDF 上，合成出带阴影的输入图（in_path字段只写 PDF 路径即可）；RealDAE 真实数据则同时写in_path与gt_path。

3.4 二值化：5 个数据集 + Sauvola 预处理

二值化任务聚合了 Bickly、DIBCO、Noise Office、PHIDB、MSI 五个数据集。由于在线计算 DTPrompt（Sauvola 二值化 + 阈值图 + 梯度图）开销大，官方选择离线预计算：

python data/preprocess/sauvola_binarize.py

该脚本基于 utils.py 中的SauvolaModBinarization生成_bin.png、_thr.png、_gradient.png三张提示图。一份带噪点的二值化训练输入样例：

JSON 需列出 5 个字段：

[ { "in_path": "binarization/noise_office/imgs/1.png", "gt_path": "binarization/noise_office/gt_imgs/1.png", "bin_path": "binarization/noise_office/imgs/1_bin.png", "thr_path": "binarization/noise_office/imgs/1_thr.png", "gradient_path": "binarization/noise_office/imgs/1_gradient.png" } ]

3.5 评测数据集准备

评测集包括 DIR300、Kligler、Jung、OSR、RealDAE、DocUNet_DocAligner、TDD、DIBCO18。下载后把输入图与真值图重命名为xxx_in.png/xxx_gt.png，放入data/eval/<数据集名>/即可，eval.py会自动匹配。

四、编辑 datasets_setting：接入你的数据

数据就位后，修改 train.py 中的datasets_setting，把路径换成你自己的数据目录：

datasets_setting = [ {'task':'deblurring','ratio':1,'im_path':'<你的路径>/train/deblurring/','json_paths':['.../deblurring/tdd/train.json']}, {'task':'dewarping','ratio':1,'im_path':'<你的路径>/train/dewarping/','json_paths':['.../dewarping/doc3d/train_1_19.json']}, {'task':'binarization','ratio':1,'im_path':'...','json_paths':['.../binarization/train.json']}, {'task':'deshadowing','ratio':1,'im_path':'...','json_paths':['.../deshadowing/train.json']}, {'task':'appearance','ratio':1,'im_path':'...','json_paths':['.../appearance/trainv2.json']} ]

• ratio控制该任务被随机采样的比例，官方全部设为 1（五任务均匀采样）；
• 另外务必把 loaders/docres_loader.py 里硬编码的background_paths（DTD 背景图）与shadow_paths（Doc3DShade 影子图）改成实际路径，否则去弯/外观增强任务会取不到数据。

五、启动训练：start_train.sh 一键 8 卡分布式

训练基于 PyTorch DDP，start_train.sh 用 8 张 GPU 启动：

CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python -m torch.distributed.launch --nproc_per_node 8 --master_port 26413 train.py

GPU 数量不同时，改CUDA_VISIBLE_DEVICES与--nproc_per_node即可。关键超参数（train.py 中均有默认值）：

训练过程中的行为：

• 每10 步在终端和logdir/<实验名>/log.txt打印 5 个任务的 loss 与剩余时间；
• 每5000 步在 rank0 保存{iter}.pkl（含模型 + 优化器状态）；
• 中断后可用--resume <ckpt路径>续训。

六、官方实验复现：eval.py 评测流程

1. 把 MBD 去弯辅助模型权重mbd.pkl放入data/MBD/checkpoint/，把 DocRes 权重docres.pkl放入./checkpoints/（下载地址见仓库 README 的 Inference/Evaluation 部分）；
2. 按 3.5 节放好评测集后运行：

python eval.py --dataset realdae

--dataset可选dir300、kligler、jung、osr、docunet_docaligner、realdae、tdd、dibco18。eval.py 会自动对*_in.*图片批量推理，结果存为*_docres.*，并输出论文同款指标：

• 去阴影 / 外观 / 去模糊：PSNR、SSIM（scikit-image 计算）；
• 二值化：F-measure、pF-measure、PSNR（utils.py 中bin_metric实现）；
• 去弯（DIR300 等）：只生成复原图，指标按论文流程人工/脚本核对。

七、常见坑与实用建议

1. 路径没改全：datasets_setting、loader 中的background_paths、shadow_paths三处路径都要替换，这是新手最常被卡住的点；
2. 二值化提示必须预生成：漏跑sauvola_binarize.py会导致bin_path等文件缺失直接报错；
3. 显存不够：先降--batch_size（如 5），或减少--nproc_per_node；im_size保持 256 以对齐论文设置；
4. 确定性复现：train.py 提供seed_torch()固定随机种子，被注释掉了，需要严格复现时取消注释；
5. 观察训练：加--tboard参数后可用tensorboard --logdir ./checkpoints/查看五任务 loss 曲线，正常情况下各任务 loss 应同步缓慢下降。

准备好以上数据与配置后，bash start_train.sh即可开启 DocRes 文档图像复原模型的统一训练；训练收敛后再用eval.py跑通 8 个基准数据集，即完成官方实验的完整复现 🎉

附：关键文件速查

• 数据集结构与 JSON 规范：data/README.md
• 影子提取预处理：data/preprocess/shadow_extraction.py
• Sauvola 二值化预处理：data/preprocess/sauvola_binarize.py
• 训练主脚本：train.py｜启动脚本：start_train.sh
• 网络定义：models/restormer_arch.py
• 评测脚本：eval.py｜工具函数（Sauvola、bin_metric）：utils.py