/PDF-Extract-Kit

A Comprehensive Toolkit for High-Quality PDF Content Extraction

Primary LanguagePythonApache License 2.0Apache-2.0

English | 简体中文

Overview

PDF documents contain a wealth of knowledge, yet extracting high-quality content from PDFs is not an easy task. To address this, we have broken down the task of PDF content extraction into several components:

  • Layout Detection: Using the LayoutLMv3 model for region detection, such as images, tables, titles, text, etc.;
  • Formula Detection: Using YOLOv8 for detecting formulas, including inline formulas and isolated formulas;
  • Formula Recognition: Using UniMERNet for formula recognition;
  • Optical Character Recognition: Using PaddleOCR for text recognition;

Note: Due to the diversity of document types, existing open-source layout and formula detection models struggle with diverse PDF documents. Therefore, we have collected diverse data for annotation and training to achieve precise detection effects on various types of documents. For details, refer to the sections on Layout Detection and Formula Detection. For formula recognition, the UniMERNet method rivals commercial software in quality across various types of formulas. For OCR, we use PaddleOCR, which performs well for both Chinese and English.

The PDF content extraction framework is illustrated below:

PDF-Extract-Kit Output Format 
{
    "layout_dets": [    # Elements on the page
        {
            "category_id": 0, # Category ID, 0~9, 13~15
            "poly": [
                136.0, # Coordinates are in image format, need to convert back to PDF coordinates, order is top-left, top-right, bottom-right, bottom-left x,y coordinates
                781.0,
                340.0,
                781.0,
                340.0,
                806.0,
                136.0,
                806.0
            ],
            "score": 0.69,   # Confidence score
            "latex": ''      # Formula recognition result, only categories 13, 14 have content, others are empty, additionally 15 is the OCR result, this key will be replaced with text
        },
        ...
    ],
    "page_info": {         # Page information: resolution size when extracting bounding boxes, alignment can be based on this information if scaling is involved
        "page_no": 0,      # Page number
        "height": 1684,    # Page height
        "width": 1200      # Page width
    }
}

The types included in category_id are as follows:

{0: 'title',              # Title
 1: 'plain text',         # Text
 2: 'abandon',            # Includes headers, footers, page numbers, and page annotations
 3: 'figure',             # Image
 4: 'figure_caption',     # Image caption
 5: 'table',              # Table
 6: 'table_caption',      # Table caption
 7: 'table_footnote',     # Table footnote
 8: 'isolate_formula',    # Display formula (this is a layout display formula, lower priority than 14)
 9: 'formula_caption',    # Display formula label

 13: 'inline_formula',    # Inline formula
 14: 'isolated_formula',  # Display formula
 15: 'ocr_text'}          # OCR result

Visualization of Results

By annotating a variety of PDF documents, we have trained robust models for layout detection and formula detection. Our pipeline achieves accurate extraction results on diverse types of PDF documents such as academic papers, textbooks, research reports, and financial statements, and is highly robust even in cases of scanned blurriness or watermarks.

Evaluation Metrics

Existing open-source models are often trained on data from Arxiv papers and fall short when facing diverse PDF documents. In contrast, our models, trained on diverse data, are capable of adapting to various document types for extraction.

Layout Detection

We have compared our model with existing open-source layout detection models, including DocXchain, Surya, and two models from 360LayoutAnalysis. The model present as LayoutLMv3-SFT in the table refers to the checkpoint we further trained with our SFT data on LayoutLMv3-base-chinese pre-trained model. The validation set for academic papers consists of 402 pages, while the textbook validation set is composed of 587 pages from various sources of textbooks.

Model Academic papers val Textbook val
mAP AP50 AR50 mAP AP50 AR50
DocXchain 52.8 69.5 77.3 34.9 50.1 63.5
Surya 24.2 39.4 66.1 13.9 23.3 49.9
360LayoutAnalysis-Paper 37.7 53.6 59.8 20.7 31.3 43.6
360LayoutAnalysis-Report 35.1 46.9 55.9 25.4 33.7 45.1
LayoutLMv3-SFT 77.6 93.3 95.5 67.9 82.7 87.9

Formula Detection

We have compared our model with the open-source formula detection model Pix2Text-MFD. Additionally, the YOLOv8-Trained is the weight obtained after we performed training on the basis of the YOLOv8l model. The paper's validation set is composed of 255 academic paper pages, and the multi-source validation set consists of 789 pages from various sources, including textbooks and books.

Model Academic papers val Multi-source val
AP50 AR50 AP50 AR50
Pix2Text-MFD 60.1 64.6 58.9 62.8
YOLOv8-Trained 87.7 89.9 82.4 87.3

Formula Recognition

The formula recognition we used is based on the weights downloaded from Unimernet, without any further SFT training, and the accuracy validation results can be obtained on its GitHub page.

Installation Guide

conda create -n pipeline python=3.10

pip install -r requirements.txt

pip install --extra-index-url https://miropsota.github.io/torch_packages_builder detectron2==0.6+pt2.3.1cu121

After installation, you may encounter some version conflicts leading to version changes. If you encounter version-related errors, you can try the following commands to reinstall specific versions of the libraries.

pip install pillow==8.4.0

In addition to version conflicts, you may also encounter errors where torch cannot be invoked. First, uninstall the following library and then reinstall cuda12 and cudnn.

pip uninstall nvidia-cusparse-cu12

Refer to Model Download to download the required model weights.

Running on Windows

If you intend to run this project on Windows, please refer to Using PDF-Extract-Kit on Windows.

Running on macOS

If you intend to run this project on macOS, please refer to Using PDF-Extract-Kit on macOS.

Run Extraction Script

python pdf_extract.py --pdf data/pdfs/ocr_1.pdf

Parameter explanations:

  • --pdf: PDF file to be processed; if a folder is passed, all PDF files in the folder will be processed.
  • --output: Path where the results are saved, default is "output".
  • --vis: Whether to visualize the results; if yes, detection results including bounding boxes and categories will be visualized.
  • --render: Whether to render the recognized results, including LaTeX code for formulas and plain text, which will be rendered and placed in the detection boxes. Note: This process is very time-consuming, and also requires prior installation of xelatex and imagemagic.

License

This repository is licensed under the Apache-2.0 License.

Please follow the model licenses to use the corresponding model weights: LayoutLMv3 / UniMERNet / YOLOv8 / PaddleOCR.

Acknowledgement