Doc2Graph

This library is the implementation of the paper Doc2Graph: a Task Agnostic Document Understanding Framework based on Graph Neural Networks, accepted at TiE @ ECCV 2022.

The model and pipeline aim at being task-agnostic on the domain of Document Understanding: key-value extraction on forms, business documents layout analysis and table detection.

Quick Start

Get up and running with Doc2Graph in minutes:

# Clone and install
git clone https://github.com/andreagemelli/doc2graph.git
cd doc2graph
uv sync
uv pip install -e .

# Initialize the project (downloads datasets)
uv run python -m doc2graph.main --init

# Run inference on a document
uv run python -m doc2graph.main \
   -addG -addT -addE -addV \
   --weights e2e-funsd-best.pt \
   --inference \
   --docs /path/to/your/image.png

Check out the tutorial notebook for a complete walkthrough!

Index:

• Doc2Graph
  • News!
  • Environment Setup
  • Training
  • Testing
  • Cite this project

---

News!

• 🔒 v0.3.1-beta — security patch. All dependencies upgraded to versions with no known vulnerabilities (Pillow 13, urllib3, tornado, Jupyter stack, ...). Development tools (pytest, ruff, pip-audit, jupyter) moved to a dev dependency group (installed by default with uv sync, skip them with uv sync --no-dev); unused dependencies (tensorboard, matplotlib, fasttext, pydantic, itsdangerous) removed. Python 3.11 is now required. A CI now runs lint, dependency audit, secret scanning and CPU smoke tests on every PR. The pretrained FUNSD checkpoints give the same results as v0.3.0-beta.

• 🚀 v0.3.0-beta — migration from DGL to PyTorch Geometric. Doc2Graph now builds on PyTorch Geometric instead of DGL. DGL is no longer actively maintained (no release in about two years), it does not support recent PyTorch versions, and its open security advisory (GHSA-3x5x-fw77-g54c) has no patched release. Moving to PyG lets us follow current PyTorch releases and keep dependencies free of known vulnerabilities. Requirements are now PyTorch >=2.9.1 and Python >=3.10.
  What was tested (CPU only, no GPU available):

  • old DGL layers vs. new PyG layers with the same weights on random batched graphs: identical outputs (max difference 0.0);
  • pretrained e2e-funsd-best.pt on the FUNSD test set, DGL vs. PyG: identical metrics (entity labeling F1 0.8266, entity linking F1 0.5713, edge AUC 0.7721), so the shipped checkpoints load and behave exactly as before;
  • short from-scratch training runs on FUNSD (train, validation, test and result saving) complete and converge.

  
  Not yet validated: full from-scratch training to the paper's numbers, GPU runs, and the PAU pipeline. This is a beta, please report any issue.
  Need DGL? For backward compatibility use release v0.2.0-beta or version 0.2.1 (last DGL-based commit: git checkout f602316).

• 🔥 Added inference method: you can now use Doc2Graph directly on your documents simply passing a path to them!
  This call will output an image with the connected entities and a json / dictionary with all the useful information you need! 🤗

uv run python -m doc2graph.main \
   -addG -addT -addE -addV \
   --weights e2e-funsd-best.pt \
   --inference \
   --docs /path/to/your/image.png

• 🔥 Added tutorial folder: get to know how to use Doc2Graph from the tutorial notebooks!

Environment Setup

Prerequisites

• Python 3.11
• uv package manager

Installation

1. Install uv (if not already installed):

pip install uv

1. Clone and setup the project:

git clone https://github.com/andreagemelli/doc2graph.git
cd doc2graph

1. Install the package in development mode:

uv sync
uv pip install -e .

This will:

• Install all dependencies with compatible versions
• Install the doc2graph package in development mode
• Set up the virtual environment with Python 3.11

Additional Setup (Optional)

For GPU acceleration, you may need to install a CUDA-specific version of PyTorch (PyG is pure Python and needs nothing extra for the ops used here):

# For CUDA 11.8 (adjust version as needed)
uv add torch==1.13.1+cu118 torchvision==0.14.1+cu118 --index-url https://download.pytorch.org/whl/cu118

Note: For different OS installations or CUDA versions, refer to PyTorch documentation.

Finally, create the project folder structure and download data:

uv run python -m doc2graph.main --init

The script will download and setup:

• FUNSD and the 'adjusted_annotations' for FUNSD[^1] are given by the work of[^3].
• The yolo detection bbox described in the paper (If you want to use YOLOv5-small to detect entities, script in notebooks/YOLO.ipynb, refer to their github for the installation. Clone the repository into doc2graph/models/yolov5).
• The Pau Riba's[^2] dataset with our train / test split.

[^1]: G. Jaume et al., FUNSD: A Dataset for Form Understanding in Noisy Scanned Documents, ICDARW 2019 [^2]: P. Riba et al, Table Detection in Invoice Documents by Graph Neural Networks, ICDAR 2019 [^3]: Hieu M. Vu et al., REVISING FUNSD DATASET FOR KEY-VALUE DETECTION IN DOCUMENT IMAGES, arXiv preprint 2020

Checkpoints You can download our model checkpoints here. Place them into doc2graph/models/checkpoints.

---

Training

1. To train our Doc2Graph model (using CPU) use:

uv run python -m doc2graph.main [SETTINGS]

1. Instead, to test a trained Doc2Graph model (using GPU) [weights can be one or more file]:

uv run python -m doc2graph.main [SETTINGS] --gpu 0 --test --weights *.pt

The project can be customized either changing directly configs/base.yaml file or providing these flags when calling doc2graph.main.

Features

• --add-geom: bool (to add positional features to graph nodes)
• --add-embs: bool (to add textual features to graph nodes)
• --add-hist: bool (to add visual features to graph nodes)
• --add-visual: bool (to add visual features to graph nodes)
• --add-eweights: bool (to add polar relative coordinates between nodes to graph edges)

Data

• --src-data: string [FUNSD, PAU or CUSTOM] (CUSTOM still under dev)
• --src-type: string [img, pdf] (if src_data is CUSTOM, still under dev)

Graphs

• --edge-type: string [fully, knn] (to change the kind of connectivity)
• --node-granularity: string [gt, yolo, ocr] (choose the granularity of nodes to be used, gt (if given), ocr (words) or yolo (entities))
• --num-polar-bins: int [Default 8] (number of bins into which discretize the space for edge polar features. It must be a power of 2)

Inference (only for KiE)

• --inference: bool (run inference on given document/s path/s)
• --docs: list (list your absolute path to your document)

Change directly configs/train.yaml for training settings or pass these flags to doc2graph.main. To create your own model (changing hyperparams) copy configs/models/*.yaml.

Training/Testing

• --model: string [e2e, edge, gcn] (which model to use, which yaml file to load)
• --gpu: int [Default -1] (which GPU to use. Set -1 to use CPU)
• --test: true / false (skip training if true)
• --weights: strin(s) (provide weight file(s) relative path(s), if testing)

Testing

You can use our pretrained models over the test sets of FUNSD[^1] and Pau Riba's[^2] datasets.

1. On FUNSD we were able to perform both Semantic Entity Labeling and Entity Linking:

E2E-FUNSD-GT:

uv run python -m doc2graph.main \
  -addG -addT -addE -addV \
  --gpu 0 \
  --test \
  --weights e2e-funsd-best.pt

E2E-FUNSD-YOLO:

uv run python -m doc2graph.main \
  -addG -addT -addE -addV \
  --gpu 0 \
  --test \
  --weights e2e-funsd-best.pt \
  --node-granularity yolo

1. on Pau Riba's dataset, we were able to perform both Layout Analysis and Table Detection

E2E-PAU:

uv run python -m doc2graph.main \
  -addG -addT -addE -addV \
   --gpu 0 \
   --test  \
   --weights e2e-pau-best.pt \
   --src-data PAU \
   --edge-type knn

---

Cite this project

If you want to use our code in your project(s), please cite us:

@InProceedings{10.1007/978-3-031-25069-9_22,
author="Gemelli, Andrea
and Biswas, Sanket
and Civitelli, Enrico
and Llad{\'o}s, Josep
and Marinai, Simone",
editor="Karlinsky, Leonid
and Michaeli, Tomer
and Nishino, Ko",
title="Doc2Graph: A Task Agnostic Document Understanding Framework Based on Graph Neural Networks",
booktitle="Computer Vision -- ECCV 2022 Workshops",
year="2023",
publisher="Springer Nature Switzerland",
address="Cham",
pages="329--344",
abstract="Geometric Deep Learning has recently attracted significant interest in a wide range of machine learning fields, including document analysis. The application of Graph Neural Networks (GNNs) has become crucial in various document-related tasks since they can unravel important structural patterns, fundamental in key information extraction processes. Previous works in the literature propose task-driven models and do not take into account the full power of graphs. We propose Doc2Graph, a task-agnostic document understanding framework based on a GNN model, to solve different tasks given different types of documents. We evaluated our approach on two challenging datasets for key information extraction in form understanding, invoice layout analysis and table detection. Our code is freely accessible on https://github.com/andreagemelli/doc2graph.",
isbn="978-3-031-25069-9"
}

---