DocRes: A Generalist Model Toward Unifying Document Image Restoration Tasks

This is the official implementation of our paper DocRes: A Generalist Model Toward Unifying Document Image Restoration Tasks.

News

🔥 A comprehensive Recommendation for Document Image Processing is available.

Inference

1. Put MBD model weights mbd.pkl to ./data/MBD/checkpoint/
2. Put DocRes model weights docres.pkl to ./checkpoints/
3. Run the following script and the results will be saved in ./restorted/. We have provided some distorted examples in ./input/.

python inference.py --im_path ./input/for_dewarping.png --task dewarping --save_dtsprompt 1

• --im_path: the path of input document image
• --task: task that need to be executed, it must be one of dewarping, deshadowing, appearance, deblurring, binarization, or end2end
• --save_dtsprompt: whether to save the DTSPrompt

Evaluation

1. Dataset preparation, see dataset instruction
2. Put MBD model weights mbd.pkl to data/MBD/checkpoint/
3. Put DocRes model weights docres.pkl to ./checkpoints/
4. Run the following script

python eval.py --dataset realdae

• --dataset: dataset that need to be evaluated, it can be set as dir300, kligler, jung, osr, docunet_docaligner, realdae, tdd, and dibco18.

Training

1. Dataset preparation, see dataset instruction
2. Specify the datasets_setting within train.py based on your dataset path and experimental setting.
3. Run the following script

bash start_train.sh

Citation:

@inproceedings{zhangdocres2024, 
Author = {Jiaxin Zhang, Dezhi Peng, Chongyu Liu , Peirong Zhang and Lianwen Jin}, 
Booktitle = {In Proceedings of the IEEE/CV Conference on Computer Vision and Pattern Recognition}, 
Title = {DocRes: A Generalist Model Toward Unifying Document Image Restoration Tasks}, 
Year = {2024}}   

⭐ Star Rising