PyramidTabNet: Transformer-Based Table Recognition in Image-Based Documents
Muhammad Umer, Muhammad Ahmed Mohsin, Adnan Ul-Hasan, and Faisal Shafait
Presented at ICDAR 2023: International Conference on Document Analysis and Recognition
Springer Link

In this paper, we introduce PyramidTabNet (PTN), a method that builds upon Convolution-less Pyramid Vision Transformer to detect tables in document images. Furthermore, we present a tabular image generative augmentation technique to effectively train the architecture. The proposed augmentation process consists of three steps, namely, clustering, fusion, and patching, for the generation of new document images containing tables. Our proposed pipeline demonstrates significant performance improvements for table detection on several standard datasets. Additionally, it achieves performance comparable to the state-of-the-art methods for structure recognition tasks.

It is recommended to create a new virtual environment so that updates/downgrades of packages do not break other projects.

• Environment characteristics
  python = 3.9.12 torch = 1.11.0 cuda = 11.3 torchvision = 0.12.0

  conda install pytorch==1.11.0 torchvision==0.12.0 torchaudio==0.11.0 cudatoolkit=11.3 -c pytorch
  

• This repo uses toolboxes provided by OpenMMLab to train and test models. Head over to the official documentation of MMDetection for installation instructions if you want to train your own model.

• Alternatively, if all you want to do is to test the model, you can install mmdet as a third-party package. Run:

  pip install -r requirements.txt

• After all the packages has been successfully installed, install mmcv by executing the following commands:

  pip install -U openmim
  mim install mmcv-full==1.6.0

• Alternatively, you can install mmcv using pip as:

  pip install mmcv-full==1.6.0 -f https://download.openmmlab.com/mmcv/dist/cu113/torch1.11/index.html
  

We provide the test set of cTDaR - TRACK A in COCO JSON format by default (for evaluation). You can access the full cTDaR dataset from the following publicly available GitHub repo: cTDaR - All Tracks. Other public datasets can be downloaded and placed in data directory for training/evaluation.

Refer to augmentation directory for instructions on how to use the scripts to generate new document images.

Following sections provide instructions to evaluate and/or train PyramidTabNet on your own data.
Note: It is recommended to execute the scripts from the project root in order to utilize the relative paths to the test set.

• Refer to Data Augmentation to generate additional training samples to improve model performance. ❤️

• Before firing up the train.py script, make sure to configure the data keys in the config file

• Refer to MMDetection documentation for more details on how to modify the keys.

  python model/train.py path/to/config/file --gpu-id 0

• Alternatively, you can launch training on multiple GPUs using the following script:

  bash model/dist_train.sh ${CONFIG_FILE} \
                          ${GPU_NUM} \
                          [optional args]

• Download link of fine-tuned weights are available in this section..

• Execute test.py with the appropriate command line arguments. Example usage:

  python model/test.py --config-file path/to/config/file \
                      --input path/to/directory \
                      --weights path/to/finetuned/checkpoint \
                      --device "cuda"

• To perform end-to-end table analysis (visualize detections) on a single image/test directory, execute main.py. Download the weights from Weights & Metrics and place them in the weights/ directory. Example usage:

  python main.py --config-file path/to/config/file \
                --input path/to/input/image or directory \
                --weights-dir path/to/weights/directory \
                --device "cuda"

• To perform table detection on a single image/test directory, execute td.py. Example usage:

  python model/td.py --config-file path/to/config/file \
                    --input path/to/input/image or directory \
                    --weights path/to/detection/weights \
                    --device "cuda" \
                    --save

• To perform stucture recognition on a single image/test directory, execute tsr.py. Example usage:

  python model/tsr.py --config-file path/to/config/file \
                      --input path/to/input/image or directory \
                      --structure-weights path/to/structure/weights \
                      --cell-weights path/to/cell/weights \
                      --device "cuda" \
                      --save

Evaluation metrics are displayed in the following tables. Note: End-user should place the downloaded weights in the weights/ directory for a streamlined evaluation of scripts.

• To download all the weights, execute:

  bash weights/get_weights.sh
  bash weights/fine_tuned.sh

Note: FinTabNet fine-tuned model is for cell-detection.

• Machines running variants of Microsoft Windows encounter issues with mmcv imports. Follow the installation guide on the official MMCV documentation to resolve such issues. Example:

  ModuleNotFoundError: No module named 'mmcv._ext'

• For table detection, if you get an error of the following form:

  Error(s) in loading state_dict for TDModel; Missing key(s) in state_dict

Resolve it by passing in the correct command line argument for --config-file.

If you find this work useful for your research, please cite our paper:

@inproceedings{umer2023pyramidtabnet,
  title={PyramidTabNet: Transformer-Based Table Recognition in Image-Based Documents},
  author={Umer, Muhammad and Mohsin, Muhammad Ahmed and Ul-Hasan, Adnan and Shafait, Faisal},
  booktitle={International Conference on Document Analysis and Recognition},
  pages={420--437},
  year={2023},
  organization={Springer}
}

Special thanks to the following contributors without which this repo would not be possible:

1. The MMDetection team for creating their amazing framework to push the state of the art computer vision research and enabling us to experiment and build various models very easily.

2. The authors of Pyramid Vision Transformer (PVT v2) for their wonderful contribution to enhance advancements in computer vision.

3. The authors of Craft Text Detector for their awesome repository for text detection.

4. The author of mAP Repo for providing a straightforward script to evaluate deep learning models for object detection metrics.

5. Google Colaboratory for providing free-high end GPU resources for research and development. All of the code base was developed using their platform and could not be possible without it.