Epoching's picture
init
c14d9ad
|
raw
history blame
5.22 kB
# DiT for Object Detection
This folder contains Mask R-CNN Cascade Mask R-CNN running instructions on top of [Detectron2](https://github.com/facebookresearch/detectron2) for PubLayNet and ICDAR 2019 cTDaR.
## Usage
### Inference
The quickest way to try out DiT for document layout analysis is the web demo: [![Hugging Face Spaces](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-blue)](https://huggingface.co/spaces/nielsr/dit-document-layout-analysis).
One can run inference using the `inference.py` script. It can be run as follows (from the root of the unilm repository):
```
python ./dit/object_detection/inference.py \
--image_path ./dit/object_detection/publaynet_example.jpeg \
--output_file_name output.jpg \
--config ./dit/object_detection/publaynet_configs/maskrcnn/maskrcnn_dit_base.yaml \
--opts MODEL.WEIGHTS https://layoutlm.blob.core.windows.net/dit/dit-fts/publaynet_dit-b_mrcnn.pth \
```
Make sure that the configuration file (YAML) and PyTorch checkpoint match. The example above uses DiT-base with the Mask R-CNN framework fine-tuned on PubLayNet.
### Data Preparation
**PubLayNet**
Download the data from this [link](https://dax-cdn.cdn.appdomain.cloud/dax-publaynet/1.0.0/publaynet.tar.gz?_ga=2.218138265.1825957955.1646384196-1495010506.1633610665) (~96GB). Then extract it to `PATH-to-PubLayNet`.
A soft link needs to be created to make the data accessible for the program:`ln -s PATH-to-PubLayNet publaynet_data`.
**ICDAR 2019 cTDaR**
Download the data from this [link](https://github.com/cndplab-founder/ICDAR2019_cTDaR) (~4GB). Assume path to this repository is named as `PATH-to-ICDARrepo`.
Then run `python convert_to_coco_format.py --root_dir=PATH-to-ICDARrepo --target_dir=PATH-toICDAR`. Now the path to processed data is `PATH-to-ICDAR`.
Run the following command to get the adaptively binarized images for archival subset.
```
cp -r PATH-to-ICDAR/trackA_archival PATH-to-ICDAR/at_trackA_archival
python adaptive_binarize.py --root_dir PATH-to-ICDAR/at_trackA_archival
```
The binarized archival subset will be in `PATH-to-ICDAR/at_trackA_archival`.
According to the subset you want to evaluate/fine-tune, a soft link should be created:`ln -s PATH-to-ICDAR/trackA_modern data` or `ln -s PATH-to-ICDAR/at_trackA_archival data`.
### Evaluation
Following commands provide two examples to evaluate the fine-tuned checkpoints.
The config files can be found in `icdar19_configs` and `publaynet_configs`.
1) Evaluate the fine-tuned checkpoint of DiT-Base with Mask R-CNN on PublayNet:
```bash
python train_net.py --config-file publaynet_configs/maskrcnn/maskrcnn_dit_base.yaml --eval-only --num-gpus 8 MODEL.WEIGHTS <finetuned_checkpoint_file_path or link> OUTPUT_DIR <your_output_dir>
```
2) Evaluate the fine-tuned checkpoint of DiT-Large with Cascade Mask R-CNN on ICDAR 2019 cTDaR archival subset (make sure you have created a soft link from `PATH-to-ICDAR/at_trackA_archival` to `data`):
```bash
python train_net.py --config-file icdar19_configs/cascade/cascade_dit_large.yaml --eval-only --num-gpus 8 MODEL.WEIGHTS <finetuned_checkpoint_file_path or link> OUTPUT_DIR <your_output_dir>
```
**Note**: We have fixed the **bug** in the [ICDAR2019 measurement tool](https://github.com/cndplab-founder/ctdar_measurement_tool) during integrating the tool into our code. If you use the tool to get the evaluation score, please modify the [code](https://github.com/cndplab-founder/ctdar_measurement_tool/blob/738456d3164a838ffaeefe7d1b5e64f3a4368a0e/evaluate.py#L146
) as follows:
```bash
...
# print(each_file)
# for file in gt_file_lst:
# if file.split(".") != "xml":
# gt_file_lst.remove(file)
# # print(gt_file_lst)
# Comment the code above and add the code below
for i in range(len(gt_file_lst) - 1, -1, -1):
if gt_file_lst[i].split(".")[-1] != "xml":
del gt_file_lst[i]
if len(gt_file_lst) > 0:
...
```
### Training
The following commands provide two examples to train the Mask R-CNN/Cascade Mask R-CNN with DiT backbone on 8 32GB Nvidia V100 GPUs.
1) Fine-tune DiT-Base with Cascade Mask R-CNN on PublayNet:
```bash
python train_net.py --config-file publaynet_configs/cascade/cascade_dit_base.yaml --num-gpus 8 MODEL.WEIGHTS <DiT-Base_file_path or link> OUTPUT_DIR <your_output_dir>
```
2) Fine-tune DiT-Large with Mask R-CNN on ICDAR 2019 cTDaR modern:
```bash
python train_net.py --config-file icdar19_configs/markrcnn/maskrcnn_dit_large.yaml --num-gpus 8 MODEL.WEIGHTS <DiT-Large_file_path or link> OUTPUT_DIR <your_output_dir>
```
[Detectron2's document](https://detectron2.readthedocs.io/en/latest/tutorials/getting_started.html) may help you for more details.
## Citation
If you find this repository useful, please consider citing our work:
```
@misc{li2022dit,
title={DiT: Self-supervised Pre-training for Document Image Transformer},
author={Junlong Li and Yiheng Xu and Tengchao Lv and Lei Cui and Cha Zhang and Furu Wei},
year={2022},
eprint={2203.02378},
archivePrefix={arXiv},
primaryClass={cs.CV}
}
```
## Acknowledgment
Thanks to [Detectron2](https://github.com/facebookresearch/detectron2) for Mask R-CNN and Cascade Mask R-CNN implementation.