Initial commit
This commit is contained in:
Christoph Auer
commit
e2d996753b
442
.gitignore
vendored
Normal file
442
.gitignore
vendored
Normal file
@ -0,0 +1,442 @@
|
|||||||
|
model_artifacts/
|
||||||
|
scratch/
|
||||||
|
ds_convert_models/
|
||||||
|
|
||||||
|
# Created by https://www.toptal.com/developers/gitignore/api/python,macos,virtualenv,pycharm,visualstudiocode,emacs,vim,jupyternotebooks
|
||||||
|
# Edit at https://www.toptal.com/developers/gitignore?templates=python,macos,virtualenv,pycharm,visualstudiocode,emacs,vim,jupyternotebooks
|
||||||
|
|
||||||
|
### Emacs ###
|
||||||
|
# -*- mode: gitignore; -*-
|
||||||
|
*~
|
||||||
|
\#*\#
|
||||||
|
/.emacs.desktop
|
||||||
|
/.emacs.desktop.lock
|
||||||
|
*.elc
|
||||||
|
auto-save-list
|
||||||
|
tramp
|
||||||
|
.\#*
|
||||||
|
|
||||||
|
# Org-mode
|
||||||
|
.org-id-locations
|
||||||
|
*_archive
|
||||||
|
|
||||||
|
# flymake-mode
|
||||||
|
*_flymake.*
|
||||||
|
|
||||||
|
# eshell files
|
||||||
|
/eshell/history
|
||||||
|
/eshell/lastdir
|
||||||
|
|
||||||
|
# elpa packages
|
||||||
|
/elpa/
|
||||||
|
|
||||||
|
# reftex files
|
||||||
|
*.rel
|
||||||
|
|
||||||
|
# AUCTeX auto folder
|
||||||
|
/auto/
|
||||||
|
|
||||||
|
# cask packages
|
||||||
|
.cask/
|
||||||
|
dist/
|
||||||
|
|
||||||
|
# Flycheck
|
||||||
|
flycheck_*.el
|
||||||
|
|
||||||
|
# server auth directory
|
||||||
|
/server/
|
||||||
|
|
||||||
|
# projectiles files
|
||||||
|
.projectile
|
||||||
|
|
||||||
|
# directory configuration
|
||||||
|
.dir-locals.el
|
||||||
|
|
||||||
|
# network security
|
||||||
|
/network-security.data
|
||||||
|
|
||||||
|
|
||||||
|
### JupyterNotebooks ###
|
||||||
|
# gitignore template for Jupyter Notebooks
|
||||||
|
# website: http://jupyter.org/
|
||||||
|
|
||||||
|
.ipynb_checkpoints
|
||||||
|
*/.ipynb_checkpoints/*
|
||||||
|
|
||||||
|
# IPython
|
||||||
|
profile_default/
|
||||||
|
ipython_config.py
|
||||||
|
|
||||||
|
# Remove previous ipynb_checkpoints
|
||||||
|
# git rm -r .ipynb_checkpoints/
|
||||||
|
|
||||||
|
### macOS ###
|
||||||
|
# General
|
||||||
|
.DS_Store
|
||||||
|
.AppleDouble
|
||||||
|
.LSOverride
|
||||||
|
|
||||||
|
# Icon must end with two \r
|
||||||
|
Icon
|
||||||
|
|
||||||
|
|
||||||
|
# Thumbnails
|
||||||
|
._*
|
||||||
|
|
||||||
|
# Files that might appear in the root of a volume
|
||||||
|
.DocumentRevisions-V100
|
||||||
|
.fseventsd
|
||||||
|
.Spotlight-V100
|
||||||
|
.TemporaryItems
|
||||||
|
.Trashes
|
||||||
|
.VolumeIcon.icns
|
||||||
|
.com.apple.timemachine.donotpresent
|
||||||
|
|
||||||
|
# Directories potentially created on remote AFP share
|
||||||
|
.AppleDB
|
||||||
|
.AppleDesktop
|
||||||
|
Network Trash Folder
|
||||||
|
Temporary Items
|
||||||
|
.apdisk
|
||||||
|
|
||||||
|
### macOS Patch ###
|
||||||
|
# iCloud generated files
|
||||||
|
*.icloud
|
||||||
|
|
||||||
|
### PyCharm ###
|
||||||
|
# Covers JetBrains IDEs: IntelliJ, RubyMine, PhpStorm, AppCode, PyCharm, CLion, Android Studio, WebStorm and Rider
|
||||||
|
# Reference: https://intellij-support.jetbrains.com/hc/en-us/articles/206544839
|
||||||
|
|
||||||
|
# User-specific stuff
|
||||||
|
.idea/**/workspace.xml
|
||||||
|
.idea/**/tasks.xml
|
||||||
|
.idea/**/usage.statistics.xml
|
||||||
|
.idea/**/dictionaries
|
||||||
|
.idea/**/shelf
|
||||||
|
|
||||||
|
# AWS User-specific
|
||||||
|
.idea/**/aws.xml
|
||||||
|
|
||||||
|
# Generated files
|
||||||
|
.idea/**/contentModel.xml
|
||||||
|
|
||||||
|
# Sensitive or high-churn files
|
||||||
|
.idea/**/dataSources/
|
||||||
|
.idea/**/dataSources.ids
|
||||||
|
.idea/**/dataSources.local.xml
|
||||||
|
.idea/**/sqlDataSources.xml
|
||||||
|
.idea/**/dynamic.xml
|
||||||
|
.idea/**/uiDesigner.xml
|
||||||
|
.idea/**/dbnavigator.xml
|
||||||
|
|
||||||
|
# Gradle
|
||||||
|
.idea/**/gradle.xml
|
||||||
|
.idea/**/libraries
|
||||||
|
|
||||||
|
# Gradle and Maven with auto-import
|
||||||
|
# When using Gradle or Maven with auto-import, you should exclude module files,
|
||||||
|
# since they will be recreated, and may cause churn. Uncomment if using
|
||||||
|
# auto-import.
|
||||||
|
# .idea/artifacts
|
||||||
|
# .idea/compiler.xml
|
||||||
|
# .idea/jarRepositories.xml
|
||||||
|
# .idea/modules.xml
|
||||||
|
# .idea/*.iml
|
||||||
|
# .idea/modules
|
||||||
|
# *.iml
|
||||||
|
# *.ipr
|
||||||
|
|
||||||
|
# CMake
|
||||||
|
cmake-build-*/
|
||||||
|
|
||||||
|
# Mongo Explorer plugin
|
||||||
|
.idea/**/mongoSettings.xml
|
||||||
|
|
||||||
|
# File-based project format
|
||||||
|
*.iws
|
||||||
|
|
||||||
|
# IntelliJ
|
||||||
|
out/
|
||||||
|
|
||||||
|
# mpeltonen/sbt-idea plugin
|
||||||
|
.idea_modules/
|
||||||
|
|
||||||
|
# JIRA plugin
|
||||||
|
atlassian-ide-plugin.xml
|
||||||
|
|
||||||
|
# Cursive Clojure plugin
|
||||||
|
.idea/replstate.xml
|
||||||
|
|
||||||
|
# SonarLint plugin
|
||||||
|
.idea/sonarlint/
|
||||||
|
|
||||||
|
# Crashlytics plugin (for Android Studio and IntelliJ)
|
||||||
|
com_crashlytics_export_strings.xml
|
||||||
|
crashlytics.properties
|
||||||
|
crashlytics-build.properties
|
||||||
|
fabric.properties
|
||||||
|
|
||||||
|
# Editor-based Rest Client
|
||||||
|
.idea/httpRequests
|
||||||
|
|
||||||
|
# Android studio 3.1+ serialized cache file
|
||||||
|
.idea/caches/build_file_checksums.ser
|
||||||
|
|
||||||
|
### PyCharm Patch ###
|
||||||
|
# Comment Reason: https://github.com/joeblau/gitignore.io/issues/186#issuecomment-215987721
|
||||||
|
|
||||||
|
# *.iml
|
||||||
|
# modules.xml
|
||||||
|
# .idea/misc.xml
|
||||||
|
# *.ipr
|
||||||
|
|
||||||
|
# Sonarlint plugin
|
||||||
|
# https://plugins.jetbrains.com/plugin/7973-sonarlint
|
||||||
|
.idea/**/sonarlint/
|
||||||
|
|
||||||
|
# SonarQube Plugin
|
||||||
|
# https://plugins.jetbrains.com/plugin/7238-sonarqube-community-plugin
|
||||||
|
.idea/**/sonarIssues.xml
|
||||||
|
|
||||||
|
# Markdown Navigator plugin
|
||||||
|
# https://plugins.jetbrains.com/plugin/7896-markdown-navigator-enhanced
|
||||||
|
.idea/**/markdown-navigator.xml
|
||||||
|
.idea/**/markdown-navigator-enh.xml
|
||||||
|
.idea/**/markdown-navigator/
|
||||||
|
|
||||||
|
# Cache file creation bug
|
||||||
|
# See https://youtrack.jetbrains.com/issue/JBR-2257
|
||||||
|
.idea/$CACHE_FILE$
|
||||||
|
|
||||||
|
# CodeStream plugin
|
||||||
|
# https://plugins.jetbrains.com/plugin/12206-codestream
|
||||||
|
.idea/codestream.xml
|
||||||
|
|
||||||
|
# Azure Toolkit for IntelliJ plugin
|
||||||
|
# https://plugins.jetbrains.com/plugin/8053-azure-toolkit-for-intellij
|
||||||
|
.idea/**/azureSettings.xml
|
||||||
|
|
||||||
|
### Python ###
|
||||||
|
# Byte-compiled / optimized / DLL files
|
||||||
|
__pycache__/
|
||||||
|
*.py[cod]
|
||||||
|
*$py.class
|
||||||
|
|
||||||
|
# C extensions
|
||||||
|
*.so
|
||||||
|
|
||||||
|
# Distribution / packaging
|
||||||
|
.Python
|
||||||
|
build/
|
||||||
|
develop-eggs/
|
||||||
|
downloads/
|
||||||
|
eggs/
|
||||||
|
.eggs/
|
||||||
|
lib/
|
||||||
|
lib64/
|
||||||
|
parts/
|
||||||
|
sdist/
|
||||||
|
var/
|
||||||
|
wheels/
|
||||||
|
share/python-wheels/
|
||||||
|
*.egg-info/
|
||||||
|
.installed.cfg
|
||||||
|
*.egg
|
||||||
|
MANIFEST
|
||||||
|
|
||||||
|
# PyInstaller
|
||||||
|
# Usually these files are written by a python script from a template
|
||||||
|
# before PyInstaller builds the exe, so as to inject date/other infos into it.
|
||||||
|
*.manifest
|
||||||
|
*.spec
|
||||||
|
|
||||||
|
# Installer logs
|
||||||
|
pip-log.txt
|
||||||
|
pip-delete-this-directory.txt
|
||||||
|
|
||||||
|
# Unit test / coverage reports
|
||||||
|
htmlcov/
|
||||||
|
.tox/
|
||||||
|
.nox/
|
||||||
|
.coverage
|
||||||
|
.coverage.*
|
||||||
|
.cache
|
||||||
|
nosetests.xml
|
||||||
|
coverage.xml
|
||||||
|
*.cover
|
||||||
|
*.py,cover
|
||||||
|
.hypothesis/
|
||||||
|
.pytest_cache/
|
||||||
|
cover/
|
||||||
|
|
||||||
|
# Translations
|
||||||
|
*.mo
|
||||||
|
*.pot
|
||||||
|
|
||||||
|
# Django stuff:
|
||||||
|
*.log
|
||||||
|
local_settings.py
|
||||||
|
db.sqlite3
|
||||||
|
db.sqlite3-journal
|
||||||
|
|
||||||
|
# Flask stuff:
|
||||||
|
instance/
|
||||||
|
.webassets-cache
|
||||||
|
|
||||||
|
# Scrapy stuff:
|
||||||
|
.scrapy
|
||||||
|
|
||||||
|
# Sphinx documentation
|
||||||
|
docs/_build/
|
||||||
|
|
||||||
|
# PyBuilder
|
||||||
|
.pybuilder/
|
||||||
|
target/
|
||||||
|
|
||||||
|
# Jupyter Notebook
|
||||||
|
|
||||||
|
# IPython
|
||||||
|
|
||||||
|
# pyenv
|
||||||
|
# For a library or package, you might want to ignore these files since the code is
|
||||||
|
# intended to run in multiple environments; otherwise, check them in:
|
||||||
|
# .python-version
|
||||||
|
|
||||||
|
# pipenv
|
||||||
|
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
|
||||||
|
# However, in case of collaboration, if having platform-specific dependencies or dependencies
|
||||||
|
# having no cross-platform support, pipenv may install dependencies that don't work, or not
|
||||||
|
# install all needed dependencies.
|
||||||
|
#Pipfile.lock
|
||||||
|
|
||||||
|
# poetry
|
||||||
|
# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
|
||||||
|
# This is especially recommended for binary packages to ensure reproducibility, and is more
|
||||||
|
# commonly ignored for libraries.
|
||||||
|
# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
|
||||||
|
#poetry.lock
|
||||||
|
|
||||||
|
# pdm
|
||||||
|
# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
|
||||||
|
#pdm.lock
|
||||||
|
# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
|
||||||
|
# in version control.
|
||||||
|
# https://pdm.fming.dev/#use-with-ide
|
||||||
|
.pdm.toml
|
||||||
|
|
||||||
|
# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
|
||||||
|
__pypackages__/
|
||||||
|
|
||||||
|
# Celery stuff
|
||||||
|
celerybeat-schedule
|
||||||
|
celerybeat.pid
|
||||||
|
|
||||||
|
# SageMath parsed files
|
||||||
|
*.sage.py
|
||||||
|
|
||||||
|
# Environments
|
||||||
|
.env
|
||||||
|
.venv
|
||||||
|
env/
|
||||||
|
venv/
|
||||||
|
ENV/
|
||||||
|
env.bak/
|
||||||
|
venv.bak/
|
||||||
|
|
||||||
|
# Spyder project settings
|
||||||
|
.spyderproject
|
||||||
|
.spyproject
|
||||||
|
|
||||||
|
# Rope project settings
|
||||||
|
.ropeproject
|
||||||
|
|
||||||
|
# mkdocs documentation
|
||||||
|
/site
|
||||||
|
|
||||||
|
# mypy
|
||||||
|
.mypy_cache/
|
||||||
|
.dmypy.json
|
||||||
|
dmypy.json
|
||||||
|
|
||||||
|
# Pyre type checker
|
||||||
|
.pyre/
|
||||||
|
|
||||||
|
# pytype static type analyzer
|
||||||
|
.pytype/
|
||||||
|
|
||||||
|
# Cython debug symbols
|
||||||
|
cython_debug/
|
||||||
|
|
||||||
|
# PyCharm
|
||||||
|
# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
|
||||||
|
# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
|
||||||
|
# and can be added to the global gitignore or merged into this file. For a more nuclear
|
||||||
|
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
|
||||||
|
.idea/
|
||||||
|
|
||||||
|
### Python Patch ###
|
||||||
|
# Poetry local configuration file - https://python-poetry.org/docs/configuration/#local-configuration
|
||||||
|
poetry.toml
|
||||||
|
|
||||||
|
# ruff
|
||||||
|
.ruff_cache/
|
||||||
|
|
||||||
|
### Vim ###
|
||||||
|
# Swap
|
||||||
|
[._]*.s[a-v][a-z]
|
||||||
|
!*.svg # comment out if you don't need vector files
|
||||||
|
[._]*.sw[a-p]
|
||||||
|
[._]s[a-rt-v][a-z]
|
||||||
|
[._]ss[a-gi-z]
|
||||||
|
[._]sw[a-p]
|
||||||
|
|
||||||
|
# Session
|
||||||
|
Session.vim
|
||||||
|
Sessionx.vim
|
||||||
|
|
||||||
|
# Temporary
|
||||||
|
.netrwhist
|
||||||
|
# Auto-generated tag files
|
||||||
|
tags
|
||||||
|
# Persistent undo
|
||||||
|
[._]*.un~
|
||||||
|
|
||||||
|
|
||||||
|
### Visual Studio Code ###
|
||||||
|
.vscode/
|
||||||
|
|
||||||
|
### VirtualEnv ###
|
||||||
|
# Virtualenv
|
||||||
|
# http://iamzed.com/2009/05/07/a-primer-on-virtualenv/
|
||||||
|
[Bb]in
|
||||||
|
[Ii]nclude
|
||||||
|
[Ll]ib
|
||||||
|
[Ll]ib64
|
||||||
|
[Ll]ocal
|
||||||
|
[Ss]cripts
|
||||||
|
pyvenv.cfg
|
||||||
|
pip-selfcheck.json
|
||||||
|
|
||||||
|
### VisualStudioCode ###
|
||||||
|
.vscode/*
|
||||||
|
!.vscode/settings.json
|
||||||
|
!.vscode/tasks.json
|
||||||
|
!.vscode/launch.json
|
||||||
|
!.vscode/extensions.json
|
||||||
|
!.vscode/*.code-snippets
|
||||||
|
|
||||||
|
# Local History for Visual Studio Code
|
||||||
|
.history/
|
||||||
|
|
||||||
|
# Built Visual Studio Code Extensions
|
||||||
|
*.vsix
|
||||||
|
|
||||||
|
### VisualStudioCode Patch ###
|
||||||
|
# Ignore all local history of files
|
||||||
|
.history
|
||||||
|
.ionide
|
||||||
|
|
||||||
|
|
||||||
|
# Docs
|
||||||
|
# docs/**/*.png
|
||||||
|
# docs/**/*.svg
|
||||||
34
.pre-commit-config.yaml
Normal file
34
.pre-commit-config.yaml
Normal file
@ -0,0 +1,34 @@
|
|||||||
|
fail_fast: true
|
||||||
|
repos:
|
||||||
|
- repo: local
|
||||||
|
hooks:
|
||||||
|
- id: system
|
||||||
|
name: Black
|
||||||
|
entry: poetry run black docling examples
|
||||||
|
pass_filenames: false
|
||||||
|
language: system
|
||||||
|
files: '\.py$'
|
||||||
|
- repo: local
|
||||||
|
hooks:
|
||||||
|
- id: system
|
||||||
|
name: isort
|
||||||
|
entry: poetry run isort docling examples
|
||||||
|
pass_filenames: false
|
||||||
|
language: system
|
||||||
|
files: '\.py$'
|
||||||
|
# - repo: local
|
||||||
|
# hooks:
|
||||||
|
# - id: system
|
||||||
|
# name: flake8
|
||||||
|
# entry: poetry run flake8 docling
|
||||||
|
# pass_filenames: false
|
||||||
|
# language: system
|
||||||
|
# files: '\.py$'
|
||||||
|
# - repo: local
|
||||||
|
# hooks:
|
||||||
|
# - id: system
|
||||||
|
# name: MyPy
|
||||||
|
# entry: poetry run mypy docling
|
||||||
|
# pass_filenames: false
|
||||||
|
# language: system
|
||||||
|
# files: '\.py$'
|
||||||
129
CODE_OF_CONDUCT.md
Normal file
129
CODE_OF_CONDUCT.md
Normal file
@ -0,0 +1,129 @@
|
|||||||
|
# Contributor Covenant Code of Conduct
|
||||||
|
|
||||||
|
## Our Pledge
|
||||||
|
|
||||||
|
We as members, contributors, and leaders pledge to make participation in our
|
||||||
|
community a harassment-free experience for everyone, regardless of age, body
|
||||||
|
size, visible or invisible disability, ethnicity, sex characteristics, gender
|
||||||
|
identity and expression, level of experience, education, socio-economic status,
|
||||||
|
nationality, personal appearance, race, religion, or sexual identity
|
||||||
|
and orientation.
|
||||||
|
|
||||||
|
We pledge to act and interact in ways that contribute to an open, welcoming,
|
||||||
|
diverse, inclusive, and healthy community.
|
||||||
|
|
||||||
|
## Our Standards
|
||||||
|
|
||||||
|
Examples of behavior that contributes to a positive environment for our
|
||||||
|
community include:
|
||||||
|
|
||||||
|
* Demonstrating empathy and kindness toward other people
|
||||||
|
* Being respectful of differing opinions, viewpoints, and experiences
|
||||||
|
* Giving and gracefully accepting constructive feedback
|
||||||
|
* Accepting responsibility and apologizing to those affected by our mistakes,
|
||||||
|
and learning from the experience
|
||||||
|
* Focusing on what is best not just for us as individuals, but for the
|
||||||
|
overall community
|
||||||
|
|
||||||
|
Examples of unacceptable behavior include:
|
||||||
|
|
||||||
|
* The use of sexualized language or imagery, and sexual attention or
|
||||||
|
advances of any kind
|
||||||
|
* Trolling, insulting or derogatory comments, and personal or political attacks
|
||||||
|
* Public or private harassment
|
||||||
|
* Publishing others' private information, such as a physical or email
|
||||||
|
address, without their explicit permission
|
||||||
|
* Other conduct which could reasonably be considered inappropriate in a
|
||||||
|
professional setting
|
||||||
|
|
||||||
|
## Enforcement Responsibilities
|
||||||
|
|
||||||
|
Community leaders are responsible for clarifying and enforcing our standards of
|
||||||
|
acceptable behavior and will take appropriate and fair corrective action in
|
||||||
|
response to any behavior that they deem inappropriate, threatening, offensive,
|
||||||
|
or harmful.
|
||||||
|
|
||||||
|
Community leaders have the right and responsibility to remove, edit, or reject
|
||||||
|
comments, commits, code, wiki edits, issues, and other contributions that are
|
||||||
|
not aligned to this Code of Conduct, and will communicate reasons for moderation
|
||||||
|
decisions when appropriate.
|
||||||
|
|
||||||
|
## Scope
|
||||||
|
|
||||||
|
This Code of Conduct applies within all community spaces, and also applies when
|
||||||
|
an individual is officially representing the community in public spaces.
|
||||||
|
Examples of representing our community include using an official e-mail address,
|
||||||
|
posting via an official social media account, or acting as an appointed
|
||||||
|
representative at an online or offline event.
|
||||||
|
|
||||||
|
## Enforcement
|
||||||
|
|
||||||
|
Instances of abusive, harassing, or otherwise unacceptable behavior may be
|
||||||
|
reported to the community leaders responsible for enforcement using
|
||||||
|
[deepsearch-core@zurich.ibm.com](mailto:deepsearch-core@zurich.ibm.com).
|
||||||
|
|
||||||
|
All complaints will be reviewed and investigated promptly and fairly.
|
||||||
|
|
||||||
|
All community leaders are obligated to respect the privacy and security of the
|
||||||
|
reporter of any incident.
|
||||||
|
|
||||||
|
## Enforcement Guidelines
|
||||||
|
|
||||||
|
Community leaders will follow these Community Impact Guidelines in determining
|
||||||
|
the consequences for any action they deem in violation of this Code of Conduct:
|
||||||
|
|
||||||
|
### 1. Correction
|
||||||
|
|
||||||
|
**Community Impact**: Use of inappropriate language or other behavior deemed
|
||||||
|
unprofessional or unwelcome in the community.
|
||||||
|
|
||||||
|
**Consequence**: A private, written warning from community leaders, providing
|
||||||
|
clarity around the nature of the violation and an explanation of why the
|
||||||
|
behavior was inappropriate. A public apology may be requested.
|
||||||
|
|
||||||
|
### 2. Warning
|
||||||
|
|
||||||
|
**Community Impact**: A violation through a single incident or series
|
||||||
|
of actions.
|
||||||
|
|
||||||
|
**Consequence**: A warning with consequences for continued behavior. No
|
||||||
|
interaction with the people involved, including unsolicited interaction with
|
||||||
|
those enforcing the Code of Conduct, for a specified period of time. This
|
||||||
|
includes avoiding interactions in community spaces as well as external channels
|
||||||
|
like social media. Violating these terms may lead to a temporary or
|
||||||
|
permanent ban.
|
||||||
|
|
||||||
|
### 3. Temporary Ban
|
||||||
|
|
||||||
|
**Community Impact**: A serious violation of community standards, including
|
||||||
|
sustained inappropriate behavior.
|
||||||
|
|
||||||
|
**Consequence**: A temporary ban from any sort of interaction or public
|
||||||
|
communication with the community for a specified period of time. No public or
|
||||||
|
private interaction with the people involved, including unsolicited interaction
|
||||||
|
with those enforcing the Code of Conduct, is allowed during this period.
|
||||||
|
Violating these terms may lead to a permanent ban.
|
||||||
|
|
||||||
|
### 4. Permanent Ban
|
||||||
|
|
||||||
|
**Community Impact**: Demonstrating a pattern of violation of community
|
||||||
|
standards, including sustained inappropriate behavior, harassment of an
|
||||||
|
individual, or aggression toward or disparagement of classes of individuals.
|
||||||
|
|
||||||
|
**Consequence**: A permanent ban from any sort of public interaction within
|
||||||
|
the community.
|
||||||
|
|
||||||
|
## Attribution
|
||||||
|
|
||||||
|
This Code of Conduct is adapted from the [Contributor Covenant][homepage],
|
||||||
|
version 2.0, available at
|
||||||
|
[https://www.contributor-covenant.org/version/2/0/code_of_conduct.html](https://www.contributor-covenant.org/version/2/0/code_of_conduct.html).
|
||||||
|
|
||||||
|
Community Impact Guidelines were inspired by [Mozilla's code of conduct
|
||||||
|
enforcement ladder](https://github.com/mozilla/diversity).
|
||||||
|
|
||||||
|
Homepage: [https://www.contributor-covenant.org](https://www.contributor-covenant.org)
|
||||||
|
|
||||||
|
For answers to common questions about this code of conduct, see the FAQ at
|
||||||
|
[https://www.contributor-covenant.org/faq](https://www.contributor-covenant.org/faq). Translations are available at
|
||||||
|
[https://www.contributor-covenant.org/translations](https://www.contributor-covenant.org/translations).
|
||||||
184
CONTRIBUTING.md
Normal file
184
CONTRIBUTING.md
Normal file
@ -0,0 +1,184 @@
|
|||||||
|
## Contributing In General
|
||||||
|
Our project welcomes external contributions. If you have an itch, please feel
|
||||||
|
free to scratch it.
|
||||||
|
|
||||||
|
To contribute code or documentation, please submit a [pull request](https://github.com/DS4SD/docling/pulls).
|
||||||
|
|
||||||
|
A good way to familiarize yourself with the codebase and contribution process is
|
||||||
|
to look for and tackle low-hanging fruit in the [issue tracker](https://github.com/DS4SD/docling/issues).
|
||||||
|
Before embarking on a more ambitious contribution, please quickly [get in touch](#communication) with us.
|
||||||
|
|
||||||
|
For general questions or support requests, please refer to the [discussion section](https://github.com/DS4SD/docling/discussions).
|
||||||
|
|
||||||
|
**Note: We appreciate your effort, and want to avoid a situation where a contribution
|
||||||
|
requires extensive rework (by you or by us), sits in backlog for a long time, or
|
||||||
|
cannot be accepted at all!**
|
||||||
|
|
||||||
|
### Proposing new features
|
||||||
|
|
||||||
|
If you would like to implement a new feature, please [raise an issue](https://github.com/DS4SD/docling/issues)
|
||||||
|
before sending a pull request so the feature can be discussed. This is to avoid
|
||||||
|
you wasting your valuable time working on a feature that the project developers
|
||||||
|
are not interested in accepting into the code base.
|
||||||
|
|
||||||
|
### Fixing bugs
|
||||||
|
|
||||||
|
If you would like to fix a bug, please [raise an issue](https://github.com/DS4SD/docling/issues) before sending a
|
||||||
|
pull request so it can be tracked.
|
||||||
|
|
||||||
|
### Merge approval
|
||||||
|
|
||||||
|
The project maintainers use LGTM (Looks Good To Me) in comments on the code
|
||||||
|
review to indicate acceptance. A change requires LGTMs from two of the
|
||||||
|
maintainers of each component affected.
|
||||||
|
|
||||||
|
For a list of the maintainers, see the [MAINTAINERS.md](MAINTAINERS.md) page.
|
||||||
|
|
||||||
|
|
||||||
|
## Legal
|
||||||
|
|
||||||
|
Each source file must include a license header for the MIT
|
||||||
|
Software. Using the SPDX format is the simplest approach.
|
||||||
|
e.g.
|
||||||
|
|
||||||
|
```
|
||||||
|
/*
|
||||||
|
Copyright IBM Inc. All rights reserved.
|
||||||
|
|
||||||
|
SPDX-License-Identifier: MIT
|
||||||
|
*/
|
||||||
|
```
|
||||||
|
|
||||||
|
We have tried to make it as easy as possible to make contributions. This
|
||||||
|
applies to how we handle the legal aspects of contribution. We use the
|
||||||
|
same approach - the [Developer's Certificate of Origin 1.1 (DCO)](https://github.com/hyperledger/fabric/blob/master/docs/source/DCO1.1.txt) - that the Linux® Kernel [community](https://elinux.org/Developer_Certificate_Of_Origin)
|
||||||
|
uses to manage code contributions.
|
||||||
|
|
||||||
|
We simply ask that when submitting a patch for review, the developer
|
||||||
|
must include a sign-off statement in the commit message.
|
||||||
|
|
||||||
|
Here is an example Signed-off-by line, which indicates that the
|
||||||
|
submitter accepts the DCO:
|
||||||
|
|
||||||
|
```
|
||||||
|
Signed-off-by: John Doe <john.doe@example.com>
|
||||||
|
```
|
||||||
|
|
||||||
|
You can include this automatically when you commit a change to your
|
||||||
|
local git repository using the following command:
|
||||||
|
|
||||||
|
```
|
||||||
|
git commit -s
|
||||||
|
```
|
||||||
|
|
||||||
|
|
||||||
|
## Communication
|
||||||
|
|
||||||
|
Please feel free to connect with us using the [discussion section](https://github.com/DS4SD/docling/discussions).
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
## Developing
|
||||||
|
|
||||||
|
### Usage of Poetry
|
||||||
|
|
||||||
|
We use Poetry to manage dependencies.
|
||||||
|
|
||||||
|
|
||||||
|
#### Install
|
||||||
|
|
||||||
|
To install, see the documentation here: https://python-poetry.org/docs/master/#installing-with-the-official-installer
|
||||||
|
|
||||||
|
1. Install the Poetry globally in your machine
|
||||||
|
```bash
|
||||||
|
curl -sSL https://install.python-poetry.org | python3 -
|
||||||
|
```
|
||||||
|
The installation script will print the installation bin folder `POETRY_BIN` which you need in the next steps.
|
||||||
|
|
||||||
|
2. Make sure Poetry is in your `$PATH`
|
||||||
|
- for `zsh`
|
||||||
|
```sh
|
||||||
|
echo 'export PATH="POETRY_BIN:$PATH"' >> ~/.zshrc
|
||||||
|
```
|
||||||
|
- for `bash`
|
||||||
|
```sh
|
||||||
|
echo 'export PATH="POETRY_BIN:$PATH"' >> ~/.bashrc
|
||||||
|
```
|
||||||
|
|
||||||
|
3. The official guidelines linked above include useful details on the configuration of autocomplete for most shell environments, e.g. Bash and Zsh.
|
||||||
|
|
||||||
|
|
||||||
|
#### Create a Virtual Environment and Install Dependencies
|
||||||
|
|
||||||
|
To activate the Virtual Environment, run:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
poetry shell
|
||||||
|
```
|
||||||
|
|
||||||
|
To spawn a shell with the Virtual Environment activated. If the Virtual Environment doesn't exist, Poetry will create one for you. Then, to install dependencies, run:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
poetry install
|
||||||
|
```
|
||||||
|
|
||||||
|
**(Advanced) Use a Specific Python Version**
|
||||||
|
|
||||||
|
If for whatever reason you need to work in a specific (older) version of Python, run:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
poetry env use $(which python3.8)
|
||||||
|
```
|
||||||
|
|
||||||
|
This creates a Virtual Environment with Python 3.8. For other versions, replace `$(which python3.8)` by the path to the interpreter (e.g., `/usr/bin/python3.8`) or use `$(which pythonX.Y)`.
|
||||||
|
|
||||||
|
|
||||||
|
#### Add a new dependency
|
||||||
|
|
||||||
|
```bash
|
||||||
|
poetry add NAME
|
||||||
|
```
|
||||||
|
|
||||||
|
## Coding style guidelines
|
||||||
|
|
||||||
|
We use the following tools to enforce code style:
|
||||||
|
|
||||||
|
- iSort, to sort imports
|
||||||
|
- Black, to format code
|
||||||
|
|
||||||
|
|
||||||
|
We run a series of checks on the code base on every commit, using `pre-commit`. To install the hooks, run:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
pre-commit install
|
||||||
|
```
|
||||||
|
|
||||||
|
To run the checks on-demand, run:
|
||||||
|
|
||||||
|
```
|
||||||
|
pre-commit run --all-files
|
||||||
|
```
|
||||||
|
|
||||||
|
Note: Checks like `Black` and `isort` will "fail" if they modify files. This is because `pre-commit` doesn't like to see files modified by their Hooks. In these cases, `git add` the modified files and `git commit` again.
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
## Documentation
|
||||||
|
|
||||||
|
We use [MkDocs](https://www.mkdocs.org/) to write documentation.
|
||||||
|
|
||||||
|
To run the documentation server, do:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
mkdocs serve
|
||||||
|
```
|
||||||
|
|
||||||
|
The server will be available on [http://localhost:8000](http://localhost:8000).
|
||||||
|
|
||||||
|
### Pushing Documentation to GitHub pages
|
||||||
|
|
||||||
|
Run the following:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
mkdocs gh-deploy
|
||||||
|
```
|
||||||
23
Dockerfile
Normal file
23
Dockerfile
Normal file
@ -0,0 +1,23 @@
|
|||||||
|
FROM python:3.11-slim-bookworm
|
||||||
|
|
||||||
|
ENV GIT_SSH_COMMAND="ssh -o StrictHostKeyChecking=no"
|
||||||
|
|
||||||
|
RUN apt-get update \
|
||||||
|
&& apt-get install -y libgl1 libglib2.0-0 curl wget git \
|
||||||
|
&& apt-get clean
|
||||||
|
|
||||||
|
RUN --mount=type=ssh \
|
||||||
|
pip install --no-cache-dir https://github.com/DS4SD/docling.git
|
||||||
|
|
||||||
|
ENV HF_HOME=/tmp/
|
||||||
|
ENV TORCH_HOME=/tmp/
|
||||||
|
|
||||||
|
COPY examples/minimal.py /root/minimal.py
|
||||||
|
|
||||||
|
RUN python -c 'from deepsearch_glm.utils.load_pretrained_models import load_pretrained_nlp_models; load_pretrained_nlp_models(verbose=True);'
|
||||||
|
RUN python -c 'from docling.document_converter import DocumentConverter; artifacts_path = DocumentConverter.download_models_hf(force=True);'
|
||||||
|
RUN wget "https://www.ibm.com/docs/en/SSQRB8/com.ibm.spectrum.si.pdfs/IBM_Storage_Insights_Fact_Sheet.pdf" -O /root/factsheet.pdf
|
||||||
|
|
||||||
|
# On container shell:
|
||||||
|
# > cd /root/
|
||||||
|
# > python minimal.py
|
||||||
21
LICENSE
Normal file
21
LICENSE
Normal file
@ -0,0 +1,21 @@
|
|||||||
|
MIT License
|
||||||
|
|
||||||
|
Copyright (c) [year] [fullname]
|
||||||
|
|
||||||
|
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||||
|
of this software and associated documentation files (the "Software"), to deal
|
||||||
|
in the Software without restriction, including without limitation the rights
|
||||||
|
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||||
|
copies of the Software, and to permit persons to whom the Software is
|
||||||
|
furnished to do so, subject to the following conditions:
|
||||||
|
|
||||||
|
The above copyright notice and this permission notice shall be included in all
|
||||||
|
copies or substantial portions of the Software.
|
||||||
|
|
||||||
|
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||||
|
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||||
|
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||||
|
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||||
|
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||||
|
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||||
|
SOFTWARE.
|
||||||
10
MAINTAINERS.md
Normal file
10
MAINTAINERS.md
Normal file
@ -0,0 +1,10 @@
|
|||||||
|
# MAINTAINERS
|
||||||
|
|
||||||
|
- Christoph Auer - [@cau-git](https://github.com/cau-git)
|
||||||
|
- Michele Dolfi - [@dolfim-ibm](https://github.com/dolfim-ibm)
|
||||||
|
- Maxim Lysak - [@maxmnemonic](https://github.com/maxmnemonic)
|
||||||
|
- Nikos Livathinos - [@nikos-livathinos](https://github.com/nikos-livathinos)
|
||||||
|
- Ahmed Nassar [@nassarofficial](https://github.com/nassarofficial)
|
||||||
|
- Peter Staar - [@PeterStaar-IBM](https://github.com/PeterStaar-IBM)
|
||||||
|
|
||||||
|
Maintainers can be contacted at [deepsearch-core@zurich.ibm.com](mailto:deepsearch-core@zurich.ibm.com).
|
||||||
99
README.md
Normal file
99
README.md
Normal file
@ -0,0 +1,99 @@
|
|||||||
|
<p align="center">
|
||||||
|
<a href="https://github.com/ds4sd/docling"> <img loading="lazy" alt="Docling" src="logo.png" width="150" /> </a>
|
||||||
|
</p>
|
||||||
|
|
||||||
|
# Docling
|
||||||
|
|
||||||
|
Dockling bundles PDF document conversion to JSON and Markdown in an easy, self-contained package.
|
||||||
|
|
||||||
|
## Features
|
||||||
|
* ⚡ Converts any PDF document to JSON or Markdown format, stable and lightning fast
|
||||||
|
* 📑 Understands detailed page layout, reading order and recovers table structures
|
||||||
|
* 📝 Extracts metadata from the document, such as title, authors, references and language
|
||||||
|
* 🔍 Optionally applies OCR (use with scanned PDFs)
|
||||||
|
|
||||||
|
## Setup
|
||||||
|
|
||||||
|
You need Python 3.11 and poetry. Install poetry from [here](https://python-poetry.org/docs/#installing-with-the-official-installer).
|
||||||
|
|
||||||
|
Once you have `poetry` installed, create an environment and install the package:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
poetry env use $(which python3.11)
|
||||||
|
poetry shell
|
||||||
|
poetry install
|
||||||
|
```
|
||||||
|
|
||||||
|
**Notes**:
|
||||||
|
* Works on macOS and Linux environments. Windows platforms are currently not tested.
|
||||||
|
|
||||||
|
|
||||||
|
## Usage
|
||||||
|
|
||||||
|
For basic usage, see the [convert.py](examples/convert.py) example module. Run with:
|
||||||
|
|
||||||
|
```
|
||||||
|
python examples/convert.py
|
||||||
|
```
|
||||||
|
The output of the above command will be written to `./scratch`.
|
||||||
|
|
||||||
|
### Enable or disable pipeline features
|
||||||
|
|
||||||
|
You can control if table structure recognition or OCR should be performed by arguments passed to `DocumentConverter`
|
||||||
|
```python
|
||||||
|
doc_converter = DocumentConverter(
|
||||||
|
artifacts_path=artifacts_path,
|
||||||
|
pipeline_options=PipelineOptions(do_table_structure=False, # Controls if table structure is recovered.
|
||||||
|
do_ocr=True), # Controls if OCR is applied (ignores programmatic content)
|
||||||
|
)
|
||||||
|
```
|
||||||
|
|
||||||
|
### Impose limits on the document size
|
||||||
|
|
||||||
|
You can limit the file size and number of pages which should be allowed to process per document.
|
||||||
|
```python
|
||||||
|
paths = [Path("./test/data/2206.01062.pdf")]
|
||||||
|
|
||||||
|
input = DocumentConversionInput.from_paths(
|
||||||
|
paths, limits=DocumentLimits(max_num_pages=100, max_file_size=20971520)
|
||||||
|
)
|
||||||
|
```
|
||||||
|
|
||||||
|
### Convert from binary PDF streams
|
||||||
|
|
||||||
|
You can convert PDFs from a binary stream instead of from the filesystem as follows:
|
||||||
|
```python
|
||||||
|
buf = BytesIO(your_binary_stream)
|
||||||
|
docs = [DocumentStream(filename="my_doc.pdf", stream=buf)]
|
||||||
|
input = DocumentConversionInput.from_streams(docs)
|
||||||
|
converted_docs = doc_converter.convert(input)
|
||||||
|
```
|
||||||
|
### Limit resource usage
|
||||||
|
|
||||||
|
You can limit the CPU threads used by `docling` by setting the environment variable `OMP_NUM_THREADS` accordingly. The default setting is using 4 CPU threads.
|
||||||
|
|
||||||
|
|
||||||
|
## Contributing
|
||||||
|
|
||||||
|
Please read [Contributing to Docling](./CONTRIBUTING.md) for details.
|
||||||
|
|
||||||
|
|
||||||
|
## References
|
||||||
|
|
||||||
|
If you use `Docling` in your projects, please consider citing the following:
|
||||||
|
|
||||||
|
```bib
|
||||||
|
@software{Docling,
|
||||||
|
author = {Deep Search Team},
|
||||||
|
month = {7},
|
||||||
|
title = {{Docling}},
|
||||||
|
url = {https://github.com/DS4SD/docling},
|
||||||
|
version = {main},
|
||||||
|
year = {2024}
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
## License
|
||||||
|
|
||||||
|
The `Docling` codebase is under MIT license.
|
||||||
|
For individual model usage, please refer to the model licenses found in the original packages.
|
||||||
0
docling/__init__.py
Normal file
0
docling/__init__.py
Normal file
0
docling/backend/__init__.py
Normal file
0
docling/backend/__init__.py
Normal file
55
docling/backend/abstract_backend.py
Normal file
55
docling/backend/abstract_backend.py
Normal file
@ -0,0 +1,55 @@
|
|||||||
|
from abc import ABC, abstractmethod
|
||||||
|
from io import BytesIO
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import Any, Iterable, Optional, Union
|
||||||
|
|
||||||
|
from PIL import Image
|
||||||
|
|
||||||
|
|
||||||
|
class PdfPageBackend(ABC):
|
||||||
|
def __init__(self, page_obj: Any) -> object:
|
||||||
|
pass
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
|
def get_text_in_rect(self, bbox: "BoundingBox") -> str:
|
||||||
|
pass
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
|
def get_text_cells(self) -> Iterable["Cell"]:
|
||||||
|
pass
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
|
def get_page_image(
|
||||||
|
self, scale: int = 1, cropbox: Optional["BoundingBox"] = None
|
||||||
|
) -> Image.Image:
|
||||||
|
pass
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
|
def get_size(self) -> "PageSize":
|
||||||
|
pass
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
|
def unload(self):
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
class PdfDocumentBackend(ABC):
|
||||||
|
@abstractmethod
|
||||||
|
def __init__(self, path_or_stream: Iterable[Union[BytesIO, Path]]):
|
||||||
|
pass
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
|
def load_page(self, page_no: int) -> PdfPageBackend:
|
||||||
|
pass
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
|
def page_count(self) -> int:
|
||||||
|
pass
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
|
def is_valid(self) -> bool:
|
||||||
|
pass
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
|
def unload(self):
|
||||||
|
pass
|
||||||
223
docling/backend/pypdfium2_backend.py
Normal file
223
docling/backend/pypdfium2_backend.py
Normal file
@ -0,0 +1,223 @@
|
|||||||
|
import random
|
||||||
|
from io import BytesIO
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import Iterable, List, Optional, Union
|
||||||
|
|
||||||
|
import pypdfium2 as pdfium
|
||||||
|
from PIL import Image, ImageDraw
|
||||||
|
from pypdfium2 import PdfPage
|
||||||
|
|
||||||
|
from docling.backend.abstract_backend import PdfDocumentBackend, PdfPageBackend
|
||||||
|
from docling.datamodel.base_models import BoundingBox, Cell, CoordOrigin, PageSize
|
||||||
|
|
||||||
|
|
||||||
|
class PyPdfiumPageBackend(PdfPageBackend):
|
||||||
|
def __init__(self, page_obj: PdfPage):
|
||||||
|
super().__init__(page_obj)
|
||||||
|
self._ppage = page_obj
|
||||||
|
self.text_page = None
|
||||||
|
|
||||||
|
def get_text_in_rect(self, bbox: BoundingBox) -> str:
|
||||||
|
if not self.text_page:
|
||||||
|
self.text_page = self._ppage.get_textpage()
|
||||||
|
|
||||||
|
if bbox.coord_origin != CoordOrigin.BOTTOMLEFT:
|
||||||
|
bbox = bbox.to_bottom_left_origin(self.get_size().height)
|
||||||
|
|
||||||
|
text_piece = self.text_page.get_text_bounded(*bbox.as_tuple())
|
||||||
|
|
||||||
|
return text_piece
|
||||||
|
|
||||||
|
def get_text_cells(self) -> Iterable[Cell]:
|
||||||
|
if not self.text_page:
|
||||||
|
self.text_page = self._ppage.get_textpage()
|
||||||
|
|
||||||
|
cells = []
|
||||||
|
cell_counter = 0
|
||||||
|
|
||||||
|
page_size = self.get_size()
|
||||||
|
|
||||||
|
for i in range(self.text_page.count_rects()):
|
||||||
|
rect = self.text_page.get_rect(i)
|
||||||
|
text_piece = self.text_page.get_text_bounded(*rect)
|
||||||
|
x0, y0, x1, y1 = rect
|
||||||
|
cells.append(
|
||||||
|
Cell(
|
||||||
|
id=cell_counter,
|
||||||
|
text=text_piece,
|
||||||
|
bbox=BoundingBox(
|
||||||
|
l=x0, b=y0, r=x1, t=y1, coord_origin=CoordOrigin.BOTTOMLEFT
|
||||||
|
).to_top_left_origin(page_size.height),
|
||||||
|
)
|
||||||
|
)
|
||||||
|
cell_counter += 1
|
||||||
|
|
||||||
|
# PyPdfium2 produces very fragmented cells, with sub-word level boundaries, in many PDFs.
|
||||||
|
# The cell merging code below is to clean this up.
|
||||||
|
def merge_horizontal_cells(
|
||||||
|
cells: List[Cell],
|
||||||
|
horizontal_threshold_factor: float = 1.0,
|
||||||
|
vertical_threshold_factor: float = 0.5,
|
||||||
|
) -> List[Cell]:
|
||||||
|
if not cells:
|
||||||
|
return []
|
||||||
|
|
||||||
|
def group_rows(cells: List[Cell]) -> List[List[Cell]]:
|
||||||
|
rows = []
|
||||||
|
current_row = [cells[0]]
|
||||||
|
row_top = cells[0].bbox.t
|
||||||
|
row_bottom = cells[0].bbox.b
|
||||||
|
row_height = cells[0].bbox.height
|
||||||
|
|
||||||
|
for cell in cells[1:]:
|
||||||
|
vertical_threshold = row_height * vertical_threshold_factor
|
||||||
|
if (
|
||||||
|
abs(cell.bbox.t - row_top) <= vertical_threshold
|
||||||
|
and abs(cell.bbox.b - row_bottom) <= vertical_threshold
|
||||||
|
):
|
||||||
|
current_row.append(cell)
|
||||||
|
row_top = min(row_top, cell.bbox.t)
|
||||||
|
row_bottom = max(row_bottom, cell.bbox.b)
|
||||||
|
row_height = row_bottom - row_top
|
||||||
|
else:
|
||||||
|
rows.append(current_row)
|
||||||
|
current_row = [cell]
|
||||||
|
row_top = cell.bbox.t
|
||||||
|
row_bottom = cell.bbox.b
|
||||||
|
row_height = cell.bbox.height
|
||||||
|
|
||||||
|
if current_row:
|
||||||
|
rows.append(current_row)
|
||||||
|
|
||||||
|
return rows
|
||||||
|
|
||||||
|
def merge_row(row: List[Cell]) -> List[Cell]:
|
||||||
|
merged = []
|
||||||
|
current_group = [row[0]]
|
||||||
|
|
||||||
|
for cell in row[1:]:
|
||||||
|
prev_cell = current_group[-1]
|
||||||
|
avg_height = (prev_cell.bbox.height + cell.bbox.height) / 2
|
||||||
|
if (
|
||||||
|
cell.bbox.l - prev_cell.bbox.r
|
||||||
|
<= avg_height * horizontal_threshold_factor
|
||||||
|
):
|
||||||
|
current_group.append(cell)
|
||||||
|
else:
|
||||||
|
merged.append(merge_group(current_group))
|
||||||
|
current_group = [cell]
|
||||||
|
|
||||||
|
if current_group:
|
||||||
|
merged.append(merge_group(current_group))
|
||||||
|
|
||||||
|
return merged
|
||||||
|
|
||||||
|
def merge_group(group: List[Cell]) -> Cell:
|
||||||
|
if len(group) == 1:
|
||||||
|
return group[0]
|
||||||
|
|
||||||
|
merged_text = "".join(cell.text for cell in group)
|
||||||
|
merged_bbox = BoundingBox(
|
||||||
|
l=min(cell.bbox.l for cell in group),
|
||||||
|
t=min(cell.bbox.t for cell in group),
|
||||||
|
r=max(cell.bbox.r for cell in group),
|
||||||
|
b=max(cell.bbox.b for cell in group),
|
||||||
|
)
|
||||||
|
return Cell(id=group[0].id, text=merged_text, bbox=merged_bbox)
|
||||||
|
|
||||||
|
rows = group_rows(cells)
|
||||||
|
merged_cells = [cell for row in rows for cell in merge_row(row)]
|
||||||
|
|
||||||
|
for i, cell in enumerate(merged_cells, 1):
|
||||||
|
cell.id = i
|
||||||
|
|
||||||
|
return merged_cells
|
||||||
|
|
||||||
|
def draw_clusters_and_cells():
|
||||||
|
image = self.get_page_image()
|
||||||
|
draw = ImageDraw.Draw(image)
|
||||||
|
for c in cells:
|
||||||
|
x0, y0, x1, y1 = c.bbox.as_tuple()
|
||||||
|
cell_color = (
|
||||||
|
random.randint(30, 140),
|
||||||
|
random.randint(30, 140),
|
||||||
|
random.randint(30, 140),
|
||||||
|
)
|
||||||
|
draw.rectangle([(x0, y0), (x1, y1)], outline=cell_color)
|
||||||
|
image.show()
|
||||||
|
|
||||||
|
# before merge:
|
||||||
|
# draw_clusters_and_cells()
|
||||||
|
|
||||||
|
cells = merge_horizontal_cells(cells)
|
||||||
|
|
||||||
|
# after merge:
|
||||||
|
# draw_clusters_and_cells()
|
||||||
|
|
||||||
|
return cells
|
||||||
|
|
||||||
|
def get_page_image(
|
||||||
|
self, scale: int = 1, cropbox: Optional[BoundingBox] = None
|
||||||
|
) -> Image.Image:
|
||||||
|
|
||||||
|
page_size = self.get_size()
|
||||||
|
|
||||||
|
if not cropbox:
|
||||||
|
cropbox = BoundingBox(
|
||||||
|
l=0,
|
||||||
|
r=page_size.width,
|
||||||
|
t=0,
|
||||||
|
b=page_size.height,
|
||||||
|
coord_origin=CoordOrigin.TOPLEFT,
|
||||||
|
)
|
||||||
|
padbox = BoundingBox(
|
||||||
|
l=0, r=0, t=0, b=0, coord_origin=CoordOrigin.BOTTOMLEFT
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
padbox = cropbox.to_bottom_left_origin(page_size.height)
|
||||||
|
padbox.r = page_size.width - padbox.r
|
||||||
|
padbox.t = page_size.height - padbox.t
|
||||||
|
|
||||||
|
image = (
|
||||||
|
self._ppage.render(
|
||||||
|
scale=scale * 1.5,
|
||||||
|
rotation=0, # no additional rotation
|
||||||
|
crop=padbox.as_tuple(),
|
||||||
|
)
|
||||||
|
.to_pil()
|
||||||
|
.resize(size=(round(cropbox.width * scale), round(cropbox.height * scale)))
|
||||||
|
) # We resize the image from 1.5x the given scale to make it sharper.
|
||||||
|
|
||||||
|
return image
|
||||||
|
|
||||||
|
def get_size(self) -> PageSize:
|
||||||
|
return PageSize(width=self._ppage.get_width(), height=self._ppage.get_height())
|
||||||
|
|
||||||
|
def unload(self):
|
||||||
|
self._ppage = None
|
||||||
|
self.text_page = None
|
||||||
|
|
||||||
|
|
||||||
|
class PyPdfiumDocumentBackend(PdfDocumentBackend):
|
||||||
|
def __init__(self, path_or_stream: Iterable[Union[BytesIO, Path]]):
|
||||||
|
super().__init__(path_or_stream)
|
||||||
|
|
||||||
|
if isinstance(path_or_stream, Path):
|
||||||
|
self._pdoc = pdfium.PdfDocument(path_or_stream)
|
||||||
|
elif isinstance(path_or_stream, BytesIO):
|
||||||
|
self._pdoc = pdfium.PdfDocument(
|
||||||
|
path_or_stream
|
||||||
|
) # TODO Fix me, won't accept bytes.
|
||||||
|
|
||||||
|
def page_count(self) -> int:
|
||||||
|
return len(self._pdoc)
|
||||||
|
|
||||||
|
def load_page(self, page_no: int) -> PdfPage:
|
||||||
|
return PyPdfiumPageBackend(self._pdoc[page_no])
|
||||||
|
|
||||||
|
def is_valid(self) -> bool:
|
||||||
|
return self.page_count() > 0
|
||||||
|
|
||||||
|
def unload(self):
|
||||||
|
self._pdoc.close()
|
||||||
|
self._pdoc = None
|
||||||
0
docling/datamodel/__init__.py
Normal file
0
docling/datamodel/__init__.py
Normal file
247
docling/datamodel/base_models.py
Normal file
247
docling/datamodel/base_models.py
Normal file
@ -0,0 +1,247 @@
|
|||||||
|
from enum import Enum, auto
|
||||||
|
from io import BytesIO
|
||||||
|
from typing import Any, Dict, List, Optional, Tuple, Union
|
||||||
|
|
||||||
|
from PIL.Image import Image
|
||||||
|
from pydantic import BaseModel, ConfigDict, model_validator
|
||||||
|
|
||||||
|
from docling.backend.abstract_backend import PdfPageBackend
|
||||||
|
|
||||||
|
|
||||||
|
class ConversionStatus(str, Enum):
|
||||||
|
PENDING = auto()
|
||||||
|
STARTED = auto()
|
||||||
|
FAILURE = auto()
|
||||||
|
SUCCESS = auto()
|
||||||
|
SUCCESS_WITH_ERRORS = auto()
|
||||||
|
|
||||||
|
|
||||||
|
class DocInputType(str, Enum):
|
||||||
|
PATH = auto()
|
||||||
|
STREAM = auto()
|
||||||
|
|
||||||
|
|
||||||
|
class CoordOrigin(str, Enum):
|
||||||
|
TOPLEFT = auto()
|
||||||
|
BOTTOMLEFT = auto()
|
||||||
|
|
||||||
|
|
||||||
|
class PageSize(BaseModel):
|
||||||
|
width: float = 0.0
|
||||||
|
height: float = 0.0
|
||||||
|
|
||||||
|
|
||||||
|
class BoundingBox(BaseModel):
|
||||||
|
l: float # left
|
||||||
|
t: float # top
|
||||||
|
r: float # right
|
||||||
|
b: float # bottom
|
||||||
|
|
||||||
|
coord_origin: CoordOrigin = CoordOrigin.TOPLEFT
|
||||||
|
|
||||||
|
@property
|
||||||
|
def width(self):
|
||||||
|
return self.r - self.l
|
||||||
|
|
||||||
|
@property
|
||||||
|
def height(self):
|
||||||
|
return abs(self.t - self.b)
|
||||||
|
|
||||||
|
def as_tuple(self):
|
||||||
|
if self.coord_origin == CoordOrigin.TOPLEFT:
|
||||||
|
return (self.l, self.t, self.r, self.b)
|
||||||
|
elif self.coord_origin == CoordOrigin.BOTTOMLEFT:
|
||||||
|
return (self.l, self.b, self.r, self.t)
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def from_tuple(cls, coord: Tuple[float], origin: CoordOrigin):
|
||||||
|
if origin == CoordOrigin.TOPLEFT:
|
||||||
|
return BoundingBox(
|
||||||
|
l=coord[0], t=coord[1], r=coord[2], b=coord[3], coord_origin=origin
|
||||||
|
)
|
||||||
|
elif origin == CoordOrigin.BOTTOMLEFT:
|
||||||
|
return BoundingBox(
|
||||||
|
l=coord[0], b=coord[1], r=coord[2], t=coord[3], coord_origin=origin
|
||||||
|
)
|
||||||
|
|
||||||
|
def area(self) -> float:
|
||||||
|
return (self.r - self.l) * (self.b - self.t)
|
||||||
|
|
||||||
|
def intersection_area_with(self, other: "BoundingBox") -> float:
|
||||||
|
# Calculate intersection coordinates
|
||||||
|
left = max(self.l, other.l)
|
||||||
|
top = max(self.t, other.t)
|
||||||
|
right = min(self.r, other.r)
|
||||||
|
bottom = min(self.b, other.b)
|
||||||
|
|
||||||
|
# Calculate intersection dimensions
|
||||||
|
width = right - left
|
||||||
|
height = bottom - top
|
||||||
|
|
||||||
|
# If the bounding boxes do not overlap, width or height will be negative
|
||||||
|
if width <= 0 or height <= 0:
|
||||||
|
return 0.0
|
||||||
|
|
||||||
|
return width * height
|
||||||
|
|
||||||
|
def to_bottom_left_origin(self, page_height) -> "BoundingBox":
|
||||||
|
if self.coord_origin == CoordOrigin.BOTTOMLEFT:
|
||||||
|
return self
|
||||||
|
elif self.coord_origin == CoordOrigin.TOPLEFT:
|
||||||
|
return BoundingBox(
|
||||||
|
l=self.l,
|
||||||
|
r=self.r,
|
||||||
|
t=page_height - self.t,
|
||||||
|
b=page_height - self.b,
|
||||||
|
coord_origin=CoordOrigin.BOTTOMLEFT,
|
||||||
|
)
|
||||||
|
|
||||||
|
def to_top_left_origin(self, page_height):
|
||||||
|
if self.coord_origin == CoordOrigin.TOPLEFT:
|
||||||
|
return self
|
||||||
|
elif self.coord_origin == CoordOrigin.BOTTOMLEFT:
|
||||||
|
return BoundingBox(
|
||||||
|
l=self.l,
|
||||||
|
r=self.r,
|
||||||
|
t=page_height - self.t, # self.b
|
||||||
|
b=page_height - self.b, # self.t
|
||||||
|
coord_origin=CoordOrigin.TOPLEFT,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class Cell(BaseModel):
|
||||||
|
id: int
|
||||||
|
text: str
|
||||||
|
bbox: BoundingBox
|
||||||
|
|
||||||
|
|
||||||
|
class OcrCell(Cell):
|
||||||
|
confidence: float
|
||||||
|
|
||||||
|
|
||||||
|
class Cluster(BaseModel):
|
||||||
|
id: int
|
||||||
|
label: str
|
||||||
|
bbox: BoundingBox
|
||||||
|
confidence: float = 1.0
|
||||||
|
cells: List[Cell] = []
|
||||||
|
|
||||||
|
|
||||||
|
class BasePageElement(BaseModel):
|
||||||
|
label: str
|
||||||
|
id: int
|
||||||
|
page_no: int
|
||||||
|
cluster: Cluster
|
||||||
|
text: Optional[str] = None
|
||||||
|
|
||||||
|
|
||||||
|
class LayoutPrediction(BaseModel):
|
||||||
|
clusters: List[Cluster] = []
|
||||||
|
|
||||||
|
|
||||||
|
class TableCell(BaseModel):
|
||||||
|
bbox: BoundingBox
|
||||||
|
row_span: int
|
||||||
|
col_span: int
|
||||||
|
start_row_offset_idx: int
|
||||||
|
end_row_offset_idx: int
|
||||||
|
start_col_offset_idx: int
|
||||||
|
end_col_offset_idx: int
|
||||||
|
text: str
|
||||||
|
column_header: bool = False
|
||||||
|
row_header: bool = False
|
||||||
|
row_section: bool = False
|
||||||
|
|
||||||
|
@model_validator(mode="before")
|
||||||
|
@classmethod
|
||||||
|
def from_dict_format(cls, data: Any) -> Any:
|
||||||
|
if isinstance(data, Dict):
|
||||||
|
text = data["bbox"].get("token", "")
|
||||||
|
if not len(text):
|
||||||
|
text_cells = data.pop("text_cell_bboxes", None)
|
||||||
|
if text_cells:
|
||||||
|
for el in text_cells:
|
||||||
|
text += el["token"] + " "
|
||||||
|
|
||||||
|
text = text.strip()
|
||||||
|
data["text"] = text
|
||||||
|
|
||||||
|
return data
|
||||||
|
|
||||||
|
|
||||||
|
class TableElement(BasePageElement):
|
||||||
|
otsl_seq: List[str]
|
||||||
|
num_rows: int = 0
|
||||||
|
num_cols: int = 0
|
||||||
|
table_cells: List[TableCell]
|
||||||
|
|
||||||
|
|
||||||
|
class TableStructurePrediction(BaseModel):
|
||||||
|
table_map: Dict[int, TableElement] = {}
|
||||||
|
|
||||||
|
|
||||||
|
class TextElement(BasePageElement):
|
||||||
|
...
|
||||||
|
|
||||||
|
|
||||||
|
class FigureData(BaseModel):
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
class FigureElement(BasePageElement):
|
||||||
|
data: Optional[FigureData] = None
|
||||||
|
provenance: Optional[str] = None
|
||||||
|
predicted_class: Optional[str] = None
|
||||||
|
confidence: Optional[float] = None
|
||||||
|
|
||||||
|
|
||||||
|
class FigureClassificationPrediction(BaseModel):
|
||||||
|
figure_count: int = 0
|
||||||
|
figure_map: Dict[int, FigureElement] = {}
|
||||||
|
|
||||||
|
|
||||||
|
class EquationPrediction(BaseModel):
|
||||||
|
equation_count: int = 0
|
||||||
|
equation_map: Dict[int, TextElement] = {}
|
||||||
|
|
||||||
|
|
||||||
|
class PagePredictions(BaseModel):
|
||||||
|
layout: LayoutPrediction = None
|
||||||
|
tablestructure: TableStructurePrediction = None
|
||||||
|
figures_classification: FigureClassificationPrediction = None
|
||||||
|
equations_prediction: EquationPrediction = None
|
||||||
|
|
||||||
|
|
||||||
|
PageElement = Union[TextElement, TableElement, FigureElement]
|
||||||
|
|
||||||
|
|
||||||
|
class AssembledUnit(BaseModel):
|
||||||
|
elements: List[PageElement]
|
||||||
|
body: List[PageElement]
|
||||||
|
headers: List[PageElement]
|
||||||
|
|
||||||
|
|
||||||
|
class Page(BaseModel):
|
||||||
|
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||||
|
|
||||||
|
page_no: int
|
||||||
|
page_hash: str = None
|
||||||
|
size: PageSize = None
|
||||||
|
image: Image = None
|
||||||
|
cells: List[Cell] = None
|
||||||
|
predictions: PagePredictions = PagePredictions()
|
||||||
|
assembled: AssembledUnit = None
|
||||||
|
|
||||||
|
_backend: PdfPageBackend = None # Internal PDF backend
|
||||||
|
|
||||||
|
|
||||||
|
class DocumentStream(BaseModel):
|
||||||
|
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||||
|
|
||||||
|
filename: str
|
||||||
|
stream: BytesIO
|
||||||
|
|
||||||
|
|
||||||
|
class PipelineOptions(BaseModel):
|
||||||
|
do_table_structure: bool = True
|
||||||
|
do_ocr: bool = False
|
||||||
351
docling/datamodel/document.py
Normal file
351
docling/datamodel/document.py
Normal file
@ -0,0 +1,351 @@
|
|||||||
|
import logging
|
||||||
|
from io import BytesIO
|
||||||
|
from pathlib import Path, PurePath
|
||||||
|
from typing import ClassVar, Dict, Iterable, List, Optional, Type, Union
|
||||||
|
|
||||||
|
from deepsearch.documents.core.export import export_to_markdown
|
||||||
|
from docling_core.types import BaseCell, BaseText
|
||||||
|
from docling_core.types import BoundingBox as DsBoundingBox
|
||||||
|
from docling_core.types import Document as DsDocument
|
||||||
|
from docling_core.types import DocumentDescription as DsDocumentDescription
|
||||||
|
from docling_core.types import FileInfoObject as DsFileInfoObject
|
||||||
|
from docling_core.types import PageDimensions, PageReference, Prov, Ref
|
||||||
|
from docling_core.types import Table as DsSchemaTable
|
||||||
|
from docling_core.types import TableCell
|
||||||
|
from pydantic import BaseModel
|
||||||
|
|
||||||
|
from docling.backend.abstract_backend import PdfDocumentBackend
|
||||||
|
from docling.backend.pypdfium2_backend import PyPdfiumDocumentBackend
|
||||||
|
from docling.datamodel.base_models import (
|
||||||
|
AssembledUnit,
|
||||||
|
ConversionStatus,
|
||||||
|
DocumentStream,
|
||||||
|
FigureElement,
|
||||||
|
Page,
|
||||||
|
TableElement,
|
||||||
|
TextElement,
|
||||||
|
)
|
||||||
|
from docling.datamodel.settings import DocumentLimits
|
||||||
|
from docling.utils.utils import create_file_hash
|
||||||
|
|
||||||
|
_log = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
layout_label_to_ds_type = {
|
||||||
|
"Title": "title",
|
||||||
|
"Document Index": "table-of-path_or_stream",
|
||||||
|
"Section-header": "subtitle-level-1",
|
||||||
|
"Checkbox-Selected": "checkbox-selected",
|
||||||
|
"Checkbox-Unselected": "checkbox-unselected",
|
||||||
|
"Caption": "caption",
|
||||||
|
"Page-header": "page-header",
|
||||||
|
"Page-footer": "page-footer",
|
||||||
|
"Footnote": "footnote",
|
||||||
|
"Table": "table",
|
||||||
|
"Formula": "equation",
|
||||||
|
"List-item": "paragraph",
|
||||||
|
"Code": "paragraph",
|
||||||
|
"Picture": "figure",
|
||||||
|
"Text": "paragraph",
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
class InputDocument(BaseModel):
|
||||||
|
file: PurePath = None
|
||||||
|
document_hash: Optional[str] = None
|
||||||
|
valid: bool = False
|
||||||
|
limits: DocumentLimits = DocumentLimits()
|
||||||
|
|
||||||
|
filesize: Optional[int] = None
|
||||||
|
page_count: Optional[int] = None
|
||||||
|
|
||||||
|
_backend: PdfDocumentBackend = None # Internal PDF backend used
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
path_or_stream: Union[BytesIO, Path],
|
||||||
|
filename: Optional[str] = None,
|
||||||
|
limits: Optional[DocumentLimits] = None,
|
||||||
|
pdf_backend=PyPdfiumDocumentBackend,
|
||||||
|
):
|
||||||
|
super().__init__()
|
||||||
|
|
||||||
|
self.limits = limits or DocumentLimits()
|
||||||
|
|
||||||
|
try:
|
||||||
|
if isinstance(path_or_stream, Path):
|
||||||
|
self.file = path_or_stream
|
||||||
|
self.filesize = path_or_stream.stat().st_size
|
||||||
|
if self.filesize > self.limits.max_file_size:
|
||||||
|
self.valid = False
|
||||||
|
else:
|
||||||
|
self.document_hash = create_file_hash(path_or_stream)
|
||||||
|
self._backend = pdf_backend(path_or_stream=path_or_stream)
|
||||||
|
|
||||||
|
elif isinstance(path_or_stream, BytesIO):
|
||||||
|
self.file = PurePath(filename)
|
||||||
|
self.filesize = path_or_stream.getbuffer().nbytes
|
||||||
|
|
||||||
|
if self.filesize > self.limits.max_file_size:
|
||||||
|
self.valid = False
|
||||||
|
else:
|
||||||
|
self.document_hash = create_file_hash(path_or_stream)
|
||||||
|
self._backend = pdf_backend(path_or_stream=path_or_stream)
|
||||||
|
|
||||||
|
if self.document_hash and self._backend.page_count() > 0:
|
||||||
|
self.page_count = self._backend.page_count()
|
||||||
|
|
||||||
|
if self.page_count <= self.limits.max_num_pages:
|
||||||
|
self.valid = True
|
||||||
|
|
||||||
|
except (FileNotFoundError, OSError) as e:
|
||||||
|
_log.exception(
|
||||||
|
f"File {self.file.name} not found or cannot be opened.", exc_info=e
|
||||||
|
)
|
||||||
|
# raise
|
||||||
|
except RuntimeError as e:
|
||||||
|
_log.exception(
|
||||||
|
f"An unexpected error occurred while opening the document {self.file.name}",
|
||||||
|
exc_info=e,
|
||||||
|
)
|
||||||
|
# raise
|
||||||
|
|
||||||
|
|
||||||
|
class ConvertedDocument(BaseModel):
|
||||||
|
input: InputDocument
|
||||||
|
|
||||||
|
status: ConversionStatus = ConversionStatus.PENDING # failure, success
|
||||||
|
errors: List[Dict] = [] # structure to keep errors
|
||||||
|
|
||||||
|
pages: List[Page] = []
|
||||||
|
assembled: AssembledUnit = None
|
||||||
|
|
||||||
|
output: DsDocument = None
|
||||||
|
|
||||||
|
def to_ds_document(self) -> DsDocument:
|
||||||
|
title = ""
|
||||||
|
desc = DsDocumentDescription(logs=[])
|
||||||
|
|
||||||
|
page_hashes = [
|
||||||
|
PageReference(hash=p.page_hash, page=p.page_no, model="default")
|
||||||
|
for p in self.pages
|
||||||
|
]
|
||||||
|
|
||||||
|
file_info = DsFileInfoObject(
|
||||||
|
filename=self.input.file.name,
|
||||||
|
document_hash=self.input.document_hash,
|
||||||
|
num_pages=self.input.page_count,
|
||||||
|
page_hashes=page_hashes,
|
||||||
|
)
|
||||||
|
|
||||||
|
main_text = []
|
||||||
|
tables = []
|
||||||
|
figures = []
|
||||||
|
|
||||||
|
page_no_to_page = {p.page_no: p for p in self.pages}
|
||||||
|
|
||||||
|
for element in self.assembled.elements:
|
||||||
|
# Convert bboxes to lower-left origin.
|
||||||
|
target_bbox = DsBoundingBox(
|
||||||
|
element.cluster.bbox.to_bottom_left_origin(
|
||||||
|
page_no_to_page[element.page_no].size.height
|
||||||
|
).as_tuple()
|
||||||
|
)
|
||||||
|
|
||||||
|
if isinstance(element, TextElement):
|
||||||
|
main_text.append(
|
||||||
|
BaseText(
|
||||||
|
text=element.text,
|
||||||
|
obj_type=layout_label_to_ds_type.get(element.label),
|
||||||
|
name=element.label,
|
||||||
|
prov=[
|
||||||
|
Prov(
|
||||||
|
bbox=target_bbox,
|
||||||
|
page=element.page_no,
|
||||||
|
span=[0, len(element.text)],
|
||||||
|
)
|
||||||
|
],
|
||||||
|
)
|
||||||
|
)
|
||||||
|
elif isinstance(element, TableElement):
|
||||||
|
index = len(tables)
|
||||||
|
ref_str = f"#/tables/{index}"
|
||||||
|
main_text.append(
|
||||||
|
Ref(
|
||||||
|
name=element.label,
|
||||||
|
obj_type=layout_label_to_ds_type.get(element.label),
|
||||||
|
ref=ref_str,
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
# Initialise empty table data grid (only empty cells)
|
||||||
|
table_data = [
|
||||||
|
[
|
||||||
|
TableCell(
|
||||||
|
text="",
|
||||||
|
# bbox=[0,0,0,0],
|
||||||
|
spans=[[i, j]],
|
||||||
|
obj_type="body",
|
||||||
|
)
|
||||||
|
for j in range(element.num_cols)
|
||||||
|
]
|
||||||
|
for i in range(element.num_rows)
|
||||||
|
]
|
||||||
|
|
||||||
|
# Overwrite cells in table data for which there is actual cell content.
|
||||||
|
for cell in element.table_cells:
|
||||||
|
for i in range(
|
||||||
|
min(cell.start_row_offset_idx, element.num_rows),
|
||||||
|
min(cell.end_row_offset_idx, element.num_rows),
|
||||||
|
):
|
||||||
|
for j in range(
|
||||||
|
min(cell.start_col_offset_idx, element.num_cols),
|
||||||
|
min(cell.end_col_offset_idx, element.num_cols),
|
||||||
|
):
|
||||||
|
celltype = "body"
|
||||||
|
if cell.column_header:
|
||||||
|
celltype = "col_header"
|
||||||
|
elif cell.row_header:
|
||||||
|
celltype = "row_header"
|
||||||
|
|
||||||
|
def make_spans(cell):
|
||||||
|
for rspan in range(
|
||||||
|
min(cell.start_row_offset_idx, element.num_rows),
|
||||||
|
min(cell.end_row_offset_idx, element.num_rows),
|
||||||
|
):
|
||||||
|
for cspan in range(
|
||||||
|
min(
|
||||||
|
cell.start_col_offset_idx, element.num_cols
|
||||||
|
),
|
||||||
|
min(cell.end_col_offset_idx, element.num_cols),
|
||||||
|
):
|
||||||
|
yield [rspan, cspan]
|
||||||
|
|
||||||
|
spans = list(make_spans(cell))
|
||||||
|
table_data[i][j] = TableCell(
|
||||||
|
text=cell.text,
|
||||||
|
bbox=cell.bbox.to_bottom_left_origin(
|
||||||
|
page_no_to_page[element.page_no].size.height
|
||||||
|
).as_tuple(),
|
||||||
|
# col=j,
|
||||||
|
# row=i,
|
||||||
|
spans=spans,
|
||||||
|
obj_type=celltype,
|
||||||
|
# col_span=[cell.start_col_offset_idx, cell.end_col_offset_idx],
|
||||||
|
# row_span=[cell.start_row_offset_idx, cell.end_row_offset_idx]
|
||||||
|
)
|
||||||
|
|
||||||
|
tables.append(
|
||||||
|
DsSchemaTable(
|
||||||
|
num_cols=element.num_cols,
|
||||||
|
num_rows=element.num_rows,
|
||||||
|
obj_type=layout_label_to_ds_type.get(element.label),
|
||||||
|
data=table_data,
|
||||||
|
prov=[
|
||||||
|
Prov(
|
||||||
|
bbox=target_bbox,
|
||||||
|
page=element.page_no,
|
||||||
|
span=[0, 0],
|
||||||
|
)
|
||||||
|
],
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
elif isinstance(element, FigureElement):
|
||||||
|
index = len(figures)
|
||||||
|
ref_str = f"#/figures/{index}"
|
||||||
|
main_text.append(
|
||||||
|
Ref(
|
||||||
|
name=element.label,
|
||||||
|
obj_type=layout_label_to_ds_type.get(element.label),
|
||||||
|
ref=ref_str,
|
||||||
|
),
|
||||||
|
)
|
||||||
|
figures.append(
|
||||||
|
BaseCell(
|
||||||
|
prov=[
|
||||||
|
Prov(
|
||||||
|
bbox=target_bbox,
|
||||||
|
page=element.page_no,
|
||||||
|
span=[0, 0],
|
||||||
|
)
|
||||||
|
],
|
||||||
|
obj_type=layout_label_to_ds_type.get(element.label),
|
||||||
|
# data=[[]],
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
page_dimensions = [
|
||||||
|
PageDimensions(page=p.page_no, height=p.size.height, width=p.size.width)
|
||||||
|
for p in self.pages
|
||||||
|
]
|
||||||
|
|
||||||
|
ds_doc = DsDocument(
|
||||||
|
name=title,
|
||||||
|
description=desc,
|
||||||
|
file_info=file_info,
|
||||||
|
main_text=main_text,
|
||||||
|
tables=tables,
|
||||||
|
figures=figures,
|
||||||
|
page_dimensions=page_dimensions,
|
||||||
|
)
|
||||||
|
|
||||||
|
return ds_doc
|
||||||
|
|
||||||
|
def render_as_dict(self):
|
||||||
|
if self.output:
|
||||||
|
return self.output.model_dump(by_alias=True, exclude_none=True)
|
||||||
|
else:
|
||||||
|
return {}
|
||||||
|
|
||||||
|
def render_as_markdown(self):
|
||||||
|
if self.output:
|
||||||
|
return export_to_markdown(
|
||||||
|
self.output.model_dump(by_alias=True, exclude_none=True)
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
return ""
|
||||||
|
|
||||||
|
|
||||||
|
class DocumentConversionInput(BaseModel):
|
||||||
|
|
||||||
|
_path_or_stream_iterator: Iterable[Union[Path, DocumentStream]] = None
|
||||||
|
limits: Optional[DocumentLimits] = DocumentLimits()
|
||||||
|
|
||||||
|
DEFAULT_BACKEND: ClassVar = PyPdfiumDocumentBackend
|
||||||
|
|
||||||
|
def docs(
|
||||||
|
self, pdf_backend: Optional[Type[PdfDocumentBackend]] = None
|
||||||
|
) -> Iterable[InputDocument]:
|
||||||
|
|
||||||
|
pdf_backend = pdf_backend or DocumentConversionInput.DEFAULT_BACKEND
|
||||||
|
|
||||||
|
for obj in self._path_or_stream_iterator:
|
||||||
|
if isinstance(obj, Path):
|
||||||
|
yield InputDocument(
|
||||||
|
path_or_stream=obj, limits=self.limits, pdf_backend=pdf_backend
|
||||||
|
)
|
||||||
|
elif isinstance(obj, DocumentStream):
|
||||||
|
yield InputDocument(
|
||||||
|
path_or_stream=obj.stream,
|
||||||
|
filename=obj.filename,
|
||||||
|
limits=self.limits,
|
||||||
|
pdf_backend=pdf_backend,
|
||||||
|
)
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def from_paths(cls, paths: Iterable[Path], limits: Optional[DocumentLimits] = None):
|
||||||
|
paths = [Path(p) for p in paths]
|
||||||
|
|
||||||
|
doc_input = cls(limits=limits)
|
||||||
|
doc_input._path_or_stream_iterator = paths
|
||||||
|
|
||||||
|
return doc_input
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def from_streams(
|
||||||
|
cls, streams: Iterable[DocumentStream], limits: Optional[DocumentLimits] = None
|
||||||
|
):
|
||||||
|
doc_input = cls(limits=limits)
|
||||||
|
doc_input._path_or_stream_iterator = streams
|
||||||
|
|
||||||
|
return doc_input
|
||||||
32
docling/datamodel/settings.py
Normal file
32
docling/datamodel/settings.py
Normal file
@ -0,0 +1,32 @@
|
|||||||
|
import sys
|
||||||
|
|
||||||
|
from pydantic import BaseModel
|
||||||
|
from pydantic_settings import BaseSettings
|
||||||
|
|
||||||
|
|
||||||
|
class DocumentLimits(BaseModel):
|
||||||
|
max_num_pages: int = sys.maxsize
|
||||||
|
max_file_size: int = sys.maxsize
|
||||||
|
|
||||||
|
|
||||||
|
class BatchConcurrencySettings(BaseModel):
|
||||||
|
doc_batch_size: int = 2
|
||||||
|
doc_batch_concurrency: int = 2
|
||||||
|
page_batch_size: int = 4
|
||||||
|
page_batch_concurrency: int = 2
|
||||||
|
|
||||||
|
# doc_batch_size: int = 1
|
||||||
|
# doc_batch_concurrency: int = 1
|
||||||
|
# page_batch_size: int = 1
|
||||||
|
# page_batch_concurrency: int = 1
|
||||||
|
|
||||||
|
# model_concurrency: int = 2
|
||||||
|
|
||||||
|
# To force models into single core: export OMP_NUM_THREADS=1
|
||||||
|
|
||||||
|
|
||||||
|
class AppSettings(BaseSettings):
|
||||||
|
perf: BatchConcurrencySettings
|
||||||
|
|
||||||
|
|
||||||
|
settings = AppSettings(perf=BatchConcurrencySettings())
|
||||||
207
docling/document_converter.py
Normal file
207
docling/document_converter.py
Normal file
@ -0,0 +1,207 @@
|
|||||||
|
import functools
|
||||||
|
import logging
|
||||||
|
import time
|
||||||
|
import traceback
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import Iterable, Optional, Type, Union
|
||||||
|
|
||||||
|
from PIL import ImageDraw
|
||||||
|
|
||||||
|
from docling.backend.abstract_backend import PdfDocumentBackend
|
||||||
|
from docling.datamodel.base_models import (
|
||||||
|
AssembledUnit,
|
||||||
|
ConversionStatus,
|
||||||
|
Page,
|
||||||
|
PipelineOptions,
|
||||||
|
)
|
||||||
|
from docling.datamodel.document import (
|
||||||
|
ConvertedDocument,
|
||||||
|
DocumentConversionInput,
|
||||||
|
InputDocument,
|
||||||
|
)
|
||||||
|
from docling.datamodel.settings import settings
|
||||||
|
from docling.models.ds_glm_model import GlmModel
|
||||||
|
from docling.models.page_assemble_model import PageAssembleModel
|
||||||
|
from docling.pipeline.base_model_pipeline import BaseModelPipeline
|
||||||
|
from docling.pipeline.standard_model_pipeline import StandardModelPipeline
|
||||||
|
from docling.utils.utils import chunkify, create_hash
|
||||||
|
|
||||||
|
_log = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
class DocumentConverter:
|
||||||
|
_layout_model_path = "model_artifacts/layout/beehive_v0.0.5"
|
||||||
|
_table_model_path = "model_artifacts/tableformer"
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
artifacts_path: Optional[Union[Path, str]] = None,
|
||||||
|
pipeline_options: PipelineOptions = PipelineOptions(),
|
||||||
|
pdf_backend: Type[PdfDocumentBackend] = DocumentConversionInput.DEFAULT_BACKEND,
|
||||||
|
pipeline_cls: Type[BaseModelPipeline] = StandardModelPipeline,
|
||||||
|
):
|
||||||
|
if not artifacts_path:
|
||||||
|
artifacts_path = self.download_models_hf()
|
||||||
|
|
||||||
|
artifacts_path = Path(artifacts_path)
|
||||||
|
|
||||||
|
self.model_pipeline = pipeline_cls(
|
||||||
|
artifacts_path=artifacts_path, pipeline_options=pipeline_options
|
||||||
|
)
|
||||||
|
|
||||||
|
self.page_assemble_model = PageAssembleModel(config={})
|
||||||
|
self.glm_model = GlmModel(config={})
|
||||||
|
self.pdf_backend = pdf_backend
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def download_models_hf(
|
||||||
|
local_dir: Optional[Path] = None, force: bool = False
|
||||||
|
) -> Path:
|
||||||
|
from huggingface_hub import snapshot_download
|
||||||
|
|
||||||
|
download_path = snapshot_download(
|
||||||
|
repo_id="ds4sd/docling-models", force_download=force, local_dir=local_dir
|
||||||
|
)
|
||||||
|
|
||||||
|
return Path(download_path)
|
||||||
|
|
||||||
|
def convert(self, input: DocumentConversionInput) -> Iterable[ConvertedDocument]:
|
||||||
|
|
||||||
|
for input_batch in chunkify(
|
||||||
|
input.docs(pdf_backend=self.pdf_backend), settings.perf.doc_batch_size
|
||||||
|
):
|
||||||
|
_log.info(f"Going to convert document batch...")
|
||||||
|
# parallel processing only within input_batch
|
||||||
|
# with ThreadPoolExecutor(
|
||||||
|
# max_workers=settings.perf.doc_batch_concurrency
|
||||||
|
# ) as pool:
|
||||||
|
# yield from pool.map(self.process_document, input_batch)
|
||||||
|
|
||||||
|
# Note: Pdfium backend is not thread-safe, thread pool usage was disabled.
|
||||||
|
yield from map(self.process_document, input_batch)
|
||||||
|
|
||||||
|
def process_document(self, in_doc: InputDocument) -> ConvertedDocument:
|
||||||
|
start_doc_time = time.time()
|
||||||
|
converted_doc = ConvertedDocument(input=in_doc)
|
||||||
|
|
||||||
|
if not in_doc.valid:
|
||||||
|
converted_doc.status = ConversionStatus.FAILURE
|
||||||
|
return converted_doc
|
||||||
|
|
||||||
|
for i in range(0, in_doc.page_count):
|
||||||
|
converted_doc.pages.append(Page(page_no=i))
|
||||||
|
|
||||||
|
all_assembled_pages = []
|
||||||
|
|
||||||
|
try:
|
||||||
|
# Iterate batches of pages (page_batch_size) in the doc
|
||||||
|
for page_batch in chunkify(
|
||||||
|
converted_doc.pages, settings.perf.page_batch_size
|
||||||
|
):
|
||||||
|
|
||||||
|
start_pb_time = time.time()
|
||||||
|
# Pipeline
|
||||||
|
|
||||||
|
# 1. Initialise the page resources
|
||||||
|
init_pages = map(
|
||||||
|
functools.partial(self.initialize_page, in_doc), page_batch
|
||||||
|
)
|
||||||
|
|
||||||
|
# 2. Populate page image
|
||||||
|
pages_with_images = map(
|
||||||
|
functools.partial(self.populate_page_images, in_doc), init_pages
|
||||||
|
)
|
||||||
|
|
||||||
|
# 3. Populate programmatic page cells
|
||||||
|
pages_with_cells = map(
|
||||||
|
functools.partial(self.parse_page_cells, in_doc),
|
||||||
|
pages_with_images,
|
||||||
|
)
|
||||||
|
|
||||||
|
pipeline_pages = self.model_pipeline.apply(pages_with_cells)
|
||||||
|
|
||||||
|
# 7. Assemble page elements (per page)
|
||||||
|
assembled_pages = self.page_assemble_model(pipeline_pages)
|
||||||
|
|
||||||
|
# exhaust assembled_pages
|
||||||
|
for assembled_page in assembled_pages:
|
||||||
|
# Free up mem resources before moving on with next batch
|
||||||
|
assembled_page.image = (
|
||||||
|
None # Comment this if you want to visualize page images
|
||||||
|
)
|
||||||
|
assembled_page._backend.unload()
|
||||||
|
|
||||||
|
all_assembled_pages.append(assembled_page)
|
||||||
|
|
||||||
|
end_pb_time = time.time() - start_pb_time
|
||||||
|
_log.info(f"Finished converting page batch time={end_pb_time:.3f}")
|
||||||
|
|
||||||
|
# Free up mem resources of PDF backend
|
||||||
|
in_doc._backend.unload()
|
||||||
|
|
||||||
|
converted_doc.pages = all_assembled_pages
|
||||||
|
self.assemble_doc(converted_doc)
|
||||||
|
|
||||||
|
converted_doc.status = ConversionStatus.SUCCESS
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
converted_doc.status = ConversionStatus.FAILURE
|
||||||
|
trace = "\n".join(traceback.format_exception(e))
|
||||||
|
_log.info(f"Encountered an error during conversion: {trace}")
|
||||||
|
|
||||||
|
end_doc_time = time.time() - start_doc_time
|
||||||
|
_log.info(
|
||||||
|
f"Finished converting document time-pages={end_doc_time:.2f}/{in_doc.page_count}"
|
||||||
|
)
|
||||||
|
|
||||||
|
return converted_doc
|
||||||
|
|
||||||
|
# Initialise and load resources for a page, before downstream steps (populate images, cells, ...)
|
||||||
|
def initialize_page(self, doc: InputDocument, page: Page) -> Page:
|
||||||
|
page._backend = doc._backend.load_page(page.page_no)
|
||||||
|
page.size = page._backend.get_size()
|
||||||
|
page.page_hash = create_hash(doc.document_hash + ":" + str(page.page_no))
|
||||||
|
|
||||||
|
return page
|
||||||
|
|
||||||
|
# Generate the page image and store it in the page object
|
||||||
|
def populate_page_images(self, doc: InputDocument, page: Page) -> Page:
|
||||||
|
page.image = page._backend.get_page_image()
|
||||||
|
|
||||||
|
return page
|
||||||
|
|
||||||
|
# Extract and populate the page cells and store it in the page object
|
||||||
|
def parse_page_cells(self, doc: InputDocument, page: Page) -> Page:
|
||||||
|
page.cells = page._backend.get_text_cells()
|
||||||
|
|
||||||
|
# DEBUG code:
|
||||||
|
def draw_text_boxes(image, cells):
|
||||||
|
draw = ImageDraw.Draw(image)
|
||||||
|
for c in cells:
|
||||||
|
x0, y0, x1, y1 = c.bbox.as_tuple()
|
||||||
|
draw.rectangle([(x0, y0), (x1, y1)], outline="red")
|
||||||
|
image.show()
|
||||||
|
|
||||||
|
# draw_text_boxes(page.image, cells)
|
||||||
|
|
||||||
|
return page
|
||||||
|
|
||||||
|
def assemble_doc(self, converted_doc: ConvertedDocument):
|
||||||
|
all_elements = []
|
||||||
|
all_headers = []
|
||||||
|
all_body = []
|
||||||
|
|
||||||
|
for p in converted_doc.pages:
|
||||||
|
|
||||||
|
for el in p.assembled.body:
|
||||||
|
all_body.append(el)
|
||||||
|
for el in p.assembled.headers:
|
||||||
|
all_headers.append(el)
|
||||||
|
for el in p.assembled.elements:
|
||||||
|
all_elements.append(el)
|
||||||
|
|
||||||
|
converted_doc.assembled = AssembledUnit(
|
||||||
|
elements=all_elements, headers=all_headers, body=all_body
|
||||||
|
)
|
||||||
|
|
||||||
|
converted_doc.output = self.glm_model(converted_doc)
|
||||||
0
docling/models/__init__.py
Normal file
0
docling/models/__init__.py
Normal file
82
docling/models/ds_glm_model.py
Normal file
82
docling/models/ds_glm_model.py
Normal file
@ -0,0 +1,82 @@
|
|||||||
|
import copy
|
||||||
|
import random
|
||||||
|
|
||||||
|
from deepsearch_glm.nlp_utils import init_nlp_model
|
||||||
|
from deepsearch_glm.utils.ds_utils import to_legacy_document_format
|
||||||
|
from deepsearch_glm.utils.load_pretrained_models import load_pretrained_nlp_models
|
||||||
|
from docling_core.types import BaseText
|
||||||
|
from docling_core.types import Document as DsDocument
|
||||||
|
from docling_core.types import Ref
|
||||||
|
from PIL import ImageDraw
|
||||||
|
|
||||||
|
from docling.datamodel.base_models import BoundingBox, Cluster, CoordOrigin
|
||||||
|
from docling.datamodel.document import ConvertedDocument
|
||||||
|
|
||||||
|
|
||||||
|
class GlmModel:
|
||||||
|
def __init__(self, config):
|
||||||
|
self.config = config
|
||||||
|
load_pretrained_nlp_models()
|
||||||
|
model = init_nlp_model(model_names="language;term;reference")
|
||||||
|
self.model = model
|
||||||
|
|
||||||
|
def __call__(self, document: ConvertedDocument) -> DsDocument:
|
||||||
|
ds_doc = document.to_ds_document()
|
||||||
|
ds_doc_dict = ds_doc.model_dump(by_alias=True)
|
||||||
|
|
||||||
|
glm_doc = self.model.apply_on_doc(ds_doc_dict)
|
||||||
|
ds_doc_dict = to_legacy_document_format(
|
||||||
|
glm_doc, ds_doc_dict, update_name_label=True
|
||||||
|
)
|
||||||
|
|
||||||
|
exported_doc = DsDocument.model_validate(ds_doc_dict)
|
||||||
|
|
||||||
|
# DEBUG code:
|
||||||
|
def draw_clusters_and_cells(ds_document, page_no):
|
||||||
|
clusters_to_draw = []
|
||||||
|
image = copy.deepcopy(document.pages[page_no].image)
|
||||||
|
for ix, elem in enumerate(ds_document.main_text):
|
||||||
|
if isinstance(elem, BaseText):
|
||||||
|
prov = elem.prov[0]
|
||||||
|
elif isinstance(elem, Ref):
|
||||||
|
_, arr, index = elem.ref.split("/")
|
||||||
|
index = int(index)
|
||||||
|
if arr == "tables":
|
||||||
|
prov = ds_document.tables[index].prov[0]
|
||||||
|
elif arr == "figures":
|
||||||
|
prov = ds_document.figures[index].prov[0]
|
||||||
|
else:
|
||||||
|
prov = None
|
||||||
|
|
||||||
|
if prov and prov.page == page_no:
|
||||||
|
clusters_to_draw.append(
|
||||||
|
Cluster(
|
||||||
|
id=ix,
|
||||||
|
label=elem.name,
|
||||||
|
bbox=BoundingBox.from_tuple(
|
||||||
|
coord=prov.bbox,
|
||||||
|
origin=CoordOrigin.BOTTOMLEFT,
|
||||||
|
).to_top_left_origin(document.pages[page_no].size.height),
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
draw = ImageDraw.Draw(image)
|
||||||
|
for c in clusters_to_draw:
|
||||||
|
x0, y0, x1, y1 = c.bbox.as_tuple()
|
||||||
|
draw.rectangle([(x0, y0), (x1, y1)], outline="red")
|
||||||
|
draw.text((x0 + 2, y0 + 2), f"{c.id}:{c.label}", fill=(255, 0, 0, 255))
|
||||||
|
|
||||||
|
cell_color = (
|
||||||
|
random.randint(30, 140),
|
||||||
|
random.randint(30, 140),
|
||||||
|
random.randint(30, 140),
|
||||||
|
)
|
||||||
|
for tc in c.cells: # [:1]:
|
||||||
|
x0, y0, x1, y1 = tc.bbox.as_tuple()
|
||||||
|
draw.rectangle([(x0, y0), (x1, y1)], outline=cell_color)
|
||||||
|
image.show()
|
||||||
|
|
||||||
|
# draw_clusters_and_cells(ds_doc, 0)
|
||||||
|
# draw_clusters_and_cells(exported_doc, 0)
|
||||||
|
|
||||||
|
return exported_doc
|
||||||
77
docling/models/easyocr_model.py
Normal file
77
docling/models/easyocr_model.py
Normal file
@ -0,0 +1,77 @@
|
|||||||
|
import copy
|
||||||
|
import logging
|
||||||
|
import random
|
||||||
|
from typing import Iterable
|
||||||
|
|
||||||
|
import numpy
|
||||||
|
from PIL import ImageDraw
|
||||||
|
|
||||||
|
from docling.datamodel.base_models import BoundingBox, CoordOrigin, OcrCell, Page
|
||||||
|
|
||||||
|
_log = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
class EasyOcrModel:
|
||||||
|
def __init__(self, config):
|
||||||
|
self.config = config
|
||||||
|
self.enabled = config["enabled"]
|
||||||
|
self.scale = 3 # multiplier for 72 dpi == 216 dpi.
|
||||||
|
|
||||||
|
if self.enabled:
|
||||||
|
import easyocr
|
||||||
|
|
||||||
|
self.reader = easyocr.Reader(config["lang"])
|
||||||
|
|
||||||
|
def __call__(self, page_batch: Iterable[Page]) -> Iterable[Page]:
|
||||||
|
|
||||||
|
if not self.enabled:
|
||||||
|
yield from page_batch
|
||||||
|
return
|
||||||
|
|
||||||
|
for page in page_batch:
|
||||||
|
# rects = page._fpage.
|
||||||
|
high_res_image = page._backend.get_page_image(scale=self.scale)
|
||||||
|
im = numpy.array(high_res_image)
|
||||||
|
result = self.reader.readtext(im)
|
||||||
|
|
||||||
|
del high_res_image
|
||||||
|
del im
|
||||||
|
|
||||||
|
cells = [
|
||||||
|
OcrCell(
|
||||||
|
id=ix,
|
||||||
|
text=line[1],
|
||||||
|
confidence=line[2],
|
||||||
|
bbox=BoundingBox.from_tuple(
|
||||||
|
coord=(
|
||||||
|
line[0][0][0] / self.scale,
|
||||||
|
line[0][0][1] / self.scale,
|
||||||
|
line[0][2][0] / self.scale,
|
||||||
|
line[0][2][1] / self.scale,
|
||||||
|
),
|
||||||
|
origin=CoordOrigin.TOPLEFT,
|
||||||
|
),
|
||||||
|
)
|
||||||
|
for ix, line in enumerate(result)
|
||||||
|
]
|
||||||
|
|
||||||
|
page.cells = cells # For now, just overwrites all digital cells.
|
||||||
|
|
||||||
|
# DEBUG code:
|
||||||
|
def draw_clusters_and_cells():
|
||||||
|
image = copy.deepcopy(page.image)
|
||||||
|
draw = ImageDraw.Draw(image)
|
||||||
|
|
||||||
|
cell_color = (
|
||||||
|
random.randint(30, 140),
|
||||||
|
random.randint(30, 140),
|
||||||
|
random.randint(30, 140),
|
||||||
|
)
|
||||||
|
for tc in cells:
|
||||||
|
x0, y0, x1, y1 = tc.bbox.as_tuple()
|
||||||
|
draw.rectangle([(x0, y0), (x1, y1)], outline=cell_color)
|
||||||
|
image.show()
|
||||||
|
|
||||||
|
# draw_clusters_and_cells()
|
||||||
|
|
||||||
|
yield page
|
||||||
318
docling/models/layout_model.py
Normal file
318
docling/models/layout_model.py
Normal file
@ -0,0 +1,318 @@
|
|||||||
|
import copy
|
||||||
|
import logging
|
||||||
|
import random
|
||||||
|
import time
|
||||||
|
from typing import Iterable, List
|
||||||
|
|
||||||
|
from docling_ibm_models.layoutmodel.layout_predictor import LayoutPredictor
|
||||||
|
from PIL import ImageDraw
|
||||||
|
|
||||||
|
from docling.datamodel.base_models import (
|
||||||
|
BoundingBox,
|
||||||
|
Cell,
|
||||||
|
Cluster,
|
||||||
|
CoordOrigin,
|
||||||
|
LayoutPrediction,
|
||||||
|
Page,
|
||||||
|
)
|
||||||
|
from docling.utils import layout_utils as lu
|
||||||
|
|
||||||
|
_log = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
class LayoutModel:
|
||||||
|
|
||||||
|
TEXT_ELEM_LABELS = [
|
||||||
|
"Text",
|
||||||
|
"Footnote",
|
||||||
|
"Caption",
|
||||||
|
"Checkbox-Unselected",
|
||||||
|
"Checkbox-Selected",
|
||||||
|
"Section-header",
|
||||||
|
"Page-header",
|
||||||
|
"Page-footer",
|
||||||
|
"Code",
|
||||||
|
"List-item",
|
||||||
|
# "Formula",
|
||||||
|
]
|
||||||
|
PAGE_HEADER_LABELS = ["Page-header", "Page-footer"]
|
||||||
|
|
||||||
|
TABLE_LABEL = "Table"
|
||||||
|
FIGURE_LABEL = "Picture"
|
||||||
|
FORMULA_LABEL = "Formula"
|
||||||
|
|
||||||
|
def __init__(self, config):
|
||||||
|
self.config = config
|
||||||
|
self.layout_predictor = LayoutPredictor(
|
||||||
|
config["artifacts_path"]
|
||||||
|
) # TODO temporary
|
||||||
|
|
||||||
|
def postprocess(self, clusters: List[Cluster], cells: List[Cell], page_height):
|
||||||
|
MIN_INTERSECTION = 0.2
|
||||||
|
CLASS_THRESHOLDS = {
|
||||||
|
"Caption": 0.35,
|
||||||
|
"Footnote": 0.35,
|
||||||
|
"Formula": 0.35,
|
||||||
|
"List-item": 0.35,
|
||||||
|
"Page-footer": 0.35,
|
||||||
|
"Page-header": 0.35,
|
||||||
|
"Picture": 0.2, # low threshold adjust to capture chemical structures for examples.
|
||||||
|
"Section-header": 0.45,
|
||||||
|
"Table": 0.35,
|
||||||
|
"Text": 0.45,
|
||||||
|
"Title": 0.45,
|
||||||
|
"Document Index": 0.45,
|
||||||
|
"Code": 0.45,
|
||||||
|
"Checkbox-Selected": 0.45,
|
||||||
|
"Checkbox-Unselected": 0.45,
|
||||||
|
"Form": 0.45,
|
||||||
|
"Key-Value Region": 0.45,
|
||||||
|
}
|
||||||
|
|
||||||
|
_log.debug("================= Start postprocess function ====================")
|
||||||
|
start_time = time.time()
|
||||||
|
# Apply Confidence Threshold to cluster predictions
|
||||||
|
# confidence = self.conf_threshold
|
||||||
|
clusters_out = []
|
||||||
|
|
||||||
|
for cluster in clusters:
|
||||||
|
confidence = CLASS_THRESHOLDS[cluster.label]
|
||||||
|
if cluster.confidence >= confidence:
|
||||||
|
# annotation["created_by"] = "high_conf_pred"
|
||||||
|
clusters_out.append(cluster)
|
||||||
|
|
||||||
|
# map to dictionary clusters and cells, with bottom left origin
|
||||||
|
clusters = [
|
||||||
|
{
|
||||||
|
"id": c.id,
|
||||||
|
"bbox": list(
|
||||||
|
c.bbox.to_bottom_left_origin(page_height).as_tuple()
|
||||||
|
), # TODO
|
||||||
|
"confidence": c.confidence,
|
||||||
|
"cell_ids": [],
|
||||||
|
"type": c.label,
|
||||||
|
}
|
||||||
|
for c in clusters
|
||||||
|
]
|
||||||
|
|
||||||
|
clusters_out = [
|
||||||
|
{
|
||||||
|
"id": c.id,
|
||||||
|
"bbox": list(
|
||||||
|
c.bbox.to_bottom_left_origin(page_height).as_tuple()
|
||||||
|
), # TODO
|
||||||
|
"confidence": c.confidence,
|
||||||
|
"created_by": "high_conf_pred",
|
||||||
|
"cell_ids": [],
|
||||||
|
"type": c.label,
|
||||||
|
}
|
||||||
|
for c in clusters_out
|
||||||
|
]
|
||||||
|
|
||||||
|
raw_cells = [
|
||||||
|
{
|
||||||
|
"id": c.id,
|
||||||
|
"bbox": list(
|
||||||
|
c.bbox.to_bottom_left_origin(page_height).as_tuple()
|
||||||
|
), # TODO
|
||||||
|
"text": c.text,
|
||||||
|
}
|
||||||
|
for c in cells
|
||||||
|
]
|
||||||
|
cell_count = len(raw_cells)
|
||||||
|
|
||||||
|
_log.debug("---- 0. Treat cluster overlaps ------")
|
||||||
|
clusters_out = lu.remove_cluster_duplicates_by_conf(clusters_out, 0.8)
|
||||||
|
|
||||||
|
_log.debug(
|
||||||
|
"---- 1. Initially assign cells to clusters based on minimum intersection ------"
|
||||||
|
)
|
||||||
|
## Check for cells included in or touched by clusters:
|
||||||
|
clusters_out = lu.assigning_cell_ids_to_clusters(
|
||||||
|
clusters_out, raw_cells, MIN_INTERSECTION
|
||||||
|
)
|
||||||
|
|
||||||
|
_log.debug("---- 2. Assign Orphans with Low Confidence Detections")
|
||||||
|
# Creates a map of cell_id->cluster_id
|
||||||
|
(
|
||||||
|
clusters_around_cells,
|
||||||
|
orphan_cell_indices,
|
||||||
|
ambiguous_cell_indices,
|
||||||
|
) = lu.cell_id_state_map(clusters_out, cell_count)
|
||||||
|
|
||||||
|
# Assign orphan cells with lower confidence predictions
|
||||||
|
clusters_out, orphan_cell_indices = lu.assign_orphans_with_low_conf_pred(
|
||||||
|
clusters_out, clusters, raw_cells, orphan_cell_indices
|
||||||
|
)
|
||||||
|
|
||||||
|
# Refresh the cell_ids assignment, after creating new clusters using low conf predictions
|
||||||
|
clusters_out = lu.assigning_cell_ids_to_clusters(
|
||||||
|
clusters_out, raw_cells, MIN_INTERSECTION
|
||||||
|
)
|
||||||
|
|
||||||
|
_log.debug("---- 3. Settle Ambigous Cells")
|
||||||
|
# Creates an update map after assignment of cell_id->cluster_id
|
||||||
|
(
|
||||||
|
clusters_around_cells,
|
||||||
|
orphan_cell_indices,
|
||||||
|
ambiguous_cell_indices,
|
||||||
|
) = lu.cell_id_state_map(clusters_out, cell_count)
|
||||||
|
|
||||||
|
# Settle pdf cells that belong to multiple clusters
|
||||||
|
clusters_out, ambiguous_cell_indices = lu.remove_ambigous_pdf_cell_by_conf(
|
||||||
|
clusters_out, raw_cells, ambiguous_cell_indices
|
||||||
|
)
|
||||||
|
|
||||||
|
_log.debug("---- 4. Set Orphans as Text")
|
||||||
|
(
|
||||||
|
clusters_around_cells,
|
||||||
|
orphan_cell_indices,
|
||||||
|
ambiguous_cell_indices,
|
||||||
|
) = lu.cell_id_state_map(clusters_out, cell_count)
|
||||||
|
|
||||||
|
clusters_out, orphan_cell_indices = lu.set_orphan_as_text(
|
||||||
|
clusters_out, clusters, raw_cells, orphan_cell_indices
|
||||||
|
)
|
||||||
|
|
||||||
|
_log.debug("---- 5. Merge Cells & and adapt the bounding boxes")
|
||||||
|
# Merge cells orphan cells
|
||||||
|
clusters_out = lu.merge_cells(clusters_out)
|
||||||
|
|
||||||
|
# Clean up clusters that remain from merged and unreasonable clusters
|
||||||
|
clusters_out = lu.clean_up_clusters(
|
||||||
|
clusters_out,
|
||||||
|
raw_cells,
|
||||||
|
merge_cells=True,
|
||||||
|
img_table=True,
|
||||||
|
one_cell_table=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
new_clusters = lu.adapt_bboxes(raw_cells, clusters_out, orphan_cell_indices)
|
||||||
|
clusters_out = new_clusters
|
||||||
|
|
||||||
|
## We first rebuild where every cell is now:
|
||||||
|
## Now we write into a prediction cells list, not into the raw cells list.
|
||||||
|
## As we don't need previous labels, we best overwrite any old list, because that might
|
||||||
|
## have been sorted differently.
|
||||||
|
(
|
||||||
|
clusters_around_cells,
|
||||||
|
orphan_cell_indices,
|
||||||
|
ambiguous_cell_indices,
|
||||||
|
) = lu.cell_id_state_map(clusters_out, cell_count)
|
||||||
|
|
||||||
|
target_cells = []
|
||||||
|
for ix, cell in enumerate(raw_cells):
|
||||||
|
new_cell = {
|
||||||
|
"id": ix,
|
||||||
|
"rawcell_id": ix,
|
||||||
|
"label": "None",
|
||||||
|
"bbox": cell["bbox"],
|
||||||
|
"text": cell["text"],
|
||||||
|
}
|
||||||
|
for cluster_index in clusters_around_cells[
|
||||||
|
ix
|
||||||
|
]: # By previous analysis, this is always 1 cluster.
|
||||||
|
new_cell["label"] = clusters_out[cluster_index]["type"]
|
||||||
|
target_cells.append(new_cell)
|
||||||
|
# _log.debug("New label of cell " + str(ix) + " is " + str(new_cell["label"]))
|
||||||
|
cells_out = target_cells
|
||||||
|
|
||||||
|
## -------------------------------
|
||||||
|
## Sort clusters into reasonable reading order, and sort the cells inside each cluster
|
||||||
|
_log.debug("---- 5. Sort clusters in reading order ------")
|
||||||
|
sorted_clusters = lu.produce_reading_order(
|
||||||
|
clusters_out, "raw_cell_ids", "raw_cell_ids", True
|
||||||
|
)
|
||||||
|
clusters_out = sorted_clusters
|
||||||
|
|
||||||
|
# end_time = timer()
|
||||||
|
_log.debug("---- End of postprocessing function ------")
|
||||||
|
end_time = time.time() - start_time
|
||||||
|
_log.debug(f"Finished post processing in seconds={end_time:.3f}")
|
||||||
|
|
||||||
|
cells_out = [
|
||||||
|
Cell(
|
||||||
|
id=c["id"],
|
||||||
|
bbox=BoundingBox.from_tuple(
|
||||||
|
coord=c["bbox"], origin=CoordOrigin.BOTTOMLEFT
|
||||||
|
).to_top_left_origin(page_height),
|
||||||
|
text=c["text"],
|
||||||
|
)
|
||||||
|
for c in cells_out
|
||||||
|
]
|
||||||
|
clusters_out_new = []
|
||||||
|
for c in clusters_out:
|
||||||
|
cluster_cells = [ccell for ccell in cells_out if ccell.id in c["cell_ids"]]
|
||||||
|
c_new = Cluster(
|
||||||
|
id=c["id"],
|
||||||
|
bbox=BoundingBox.from_tuple(
|
||||||
|
coord=c["bbox"], origin=CoordOrigin.BOTTOMLEFT
|
||||||
|
).to_top_left_origin(page_height),
|
||||||
|
confidence=c["confidence"],
|
||||||
|
label=c["type"],
|
||||||
|
cells=cluster_cells,
|
||||||
|
)
|
||||||
|
clusters_out_new.append(c_new)
|
||||||
|
|
||||||
|
return clusters_out_new, cells_out
|
||||||
|
|
||||||
|
def __call__(self, page_batch: Iterable[Page]) -> Iterable[Page]:
|
||||||
|
for page in page_batch:
|
||||||
|
clusters = []
|
||||||
|
for ix, pred_item in enumerate(self.layout_predictor.predict(page.image)):
|
||||||
|
cluster = Cluster(
|
||||||
|
id=ix,
|
||||||
|
label=pred_item["label"],
|
||||||
|
confidence=pred_item["confidence"],
|
||||||
|
bbox=BoundingBox.model_validate(pred_item),
|
||||||
|
cells=[],
|
||||||
|
)
|
||||||
|
clusters.append(cluster)
|
||||||
|
|
||||||
|
# Map cells to clusters
|
||||||
|
# TODO: Remove, postprocess should take care of it anyway.
|
||||||
|
for cell in page.cells:
|
||||||
|
for cluster in clusters:
|
||||||
|
if not cell.bbox.area() > 0:
|
||||||
|
overlap_frac = 0.0
|
||||||
|
else:
|
||||||
|
overlap_frac = (
|
||||||
|
cell.bbox.intersection_area_with(cluster.bbox)
|
||||||
|
/ cell.bbox.area()
|
||||||
|
)
|
||||||
|
|
||||||
|
if overlap_frac > 0.5:
|
||||||
|
cluster.cells.append(cell)
|
||||||
|
|
||||||
|
# Pre-sort clusters
|
||||||
|
# clusters = self.sort_clusters_by_cell_order(clusters)
|
||||||
|
|
||||||
|
# DEBUG code:
|
||||||
|
def draw_clusters_and_cells():
|
||||||
|
image = copy.deepcopy(page.image)
|
||||||
|
draw = ImageDraw.Draw(image)
|
||||||
|
for c in clusters:
|
||||||
|
x0, y0, x1, y1 = c.bbox.as_tuple()
|
||||||
|
draw.rectangle([(x0, y0), (x1, y1)], outline="green")
|
||||||
|
|
||||||
|
cell_color = (
|
||||||
|
random.randint(30, 140),
|
||||||
|
random.randint(30, 140),
|
||||||
|
random.randint(30, 140),
|
||||||
|
)
|
||||||
|
for tc in c.cells: # [:1]:
|
||||||
|
x0, y0, x1, y1 = tc.bbox.as_tuple()
|
||||||
|
draw.rectangle([(x0, y0), (x1, y1)], outline=cell_color)
|
||||||
|
image.show()
|
||||||
|
|
||||||
|
# draw_clusters_and_cells()
|
||||||
|
|
||||||
|
clusters, page.cells = self.postprocess(
|
||||||
|
clusters, page.cells, page.size.height
|
||||||
|
)
|
||||||
|
|
||||||
|
# draw_clusters_and_cells()
|
||||||
|
|
||||||
|
page.predictions.layout = LayoutPrediction(clusters=clusters)
|
||||||
|
|
||||||
|
yield page
|
||||||
160
docling/models/page_assemble_model.py
Normal file
160
docling/models/page_assemble_model.py
Normal file
@ -0,0 +1,160 @@
|
|||||||
|
import logging
|
||||||
|
import re
|
||||||
|
from typing import Iterable, List
|
||||||
|
|
||||||
|
from docling.datamodel.base_models import (
|
||||||
|
AssembledUnit,
|
||||||
|
FigureElement,
|
||||||
|
Page,
|
||||||
|
PageElement,
|
||||||
|
TableElement,
|
||||||
|
TextElement,
|
||||||
|
)
|
||||||
|
from docling.models.layout_model import LayoutModel
|
||||||
|
|
||||||
|
_log = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
class PageAssembleModel:
|
||||||
|
def __init__(self, config):
|
||||||
|
self.config = config
|
||||||
|
|
||||||
|
# self.line_wrap_pattern = re.compile(r'(?<=[^\W_])- \n(?=\w)')
|
||||||
|
|
||||||
|
# def sanitize_text_poor(self, lines):
|
||||||
|
# text = '\n'.join(lines)
|
||||||
|
#
|
||||||
|
# # treat line wraps.
|
||||||
|
# sanitized_text = self.line_wrap_pattern.sub('', text)
|
||||||
|
#
|
||||||
|
# sanitized_text = sanitized_text.replace('\n', ' ')
|
||||||
|
#
|
||||||
|
# return sanitized_text
|
||||||
|
|
||||||
|
def sanitize_text(self, lines):
|
||||||
|
if len(lines) <= 1:
|
||||||
|
return " ".join(lines)
|
||||||
|
|
||||||
|
for ix, line in enumerate(lines[1:]):
|
||||||
|
prev_line = lines[ix]
|
||||||
|
|
||||||
|
if prev_line.endswith("-"):
|
||||||
|
prev_words = re.findall(r"\b[\w]+\b", prev_line)
|
||||||
|
line_words = re.findall(r"\b[\w]+\b", line)
|
||||||
|
|
||||||
|
if (
|
||||||
|
len(prev_words)
|
||||||
|
and len(line_words)
|
||||||
|
and prev_words[-1].isalnum()
|
||||||
|
and line_words[0].isalnum()
|
||||||
|
):
|
||||||
|
lines[ix] = prev_line[:-1]
|
||||||
|
else:
|
||||||
|
lines[ix] += " "
|
||||||
|
|
||||||
|
sanitized_text = "".join(lines)
|
||||||
|
|
||||||
|
return sanitized_text.strip() # Strip any leading or trailing whitespace
|
||||||
|
|
||||||
|
def __call__(self, page_batch: Iterable[Page]) -> Iterable[Page]:
|
||||||
|
for page in page_batch:
|
||||||
|
# assembles some JSON output page by page.
|
||||||
|
|
||||||
|
elements: List[PageElement] = []
|
||||||
|
headers: List[PageElement] = []
|
||||||
|
body: List[PageElement] = []
|
||||||
|
|
||||||
|
for cluster in page.predictions.layout.clusters:
|
||||||
|
# _log.info("Cluster label seen:", cluster.label)
|
||||||
|
if cluster.label in LayoutModel.TEXT_ELEM_LABELS:
|
||||||
|
|
||||||
|
textlines = [
|
||||||
|
cell.text.replace("\x02", "-").strip()
|
||||||
|
for cell in cluster.cells
|
||||||
|
if len(cell.text.strip()) > 0
|
||||||
|
]
|
||||||
|
text = self.sanitize_text(textlines)
|
||||||
|
text_el = TextElement(
|
||||||
|
label=cluster.label,
|
||||||
|
id=cluster.id,
|
||||||
|
text=text,
|
||||||
|
page_no=page.page_no,
|
||||||
|
cluster=cluster,
|
||||||
|
)
|
||||||
|
elements.append(text_el)
|
||||||
|
|
||||||
|
if cluster.label in LayoutModel.PAGE_HEADER_LABELS:
|
||||||
|
headers.append(text_el)
|
||||||
|
else:
|
||||||
|
body.append(text_el)
|
||||||
|
elif cluster.label == LayoutModel.TABLE_LABEL:
|
||||||
|
tbl = None
|
||||||
|
if page.predictions.tablestructure:
|
||||||
|
tbl = page.predictions.tablestructure.table_map.get(
|
||||||
|
cluster.id, None
|
||||||
|
)
|
||||||
|
if (
|
||||||
|
not tbl
|
||||||
|
): # fallback: add table without structure, if it isn't present
|
||||||
|
tbl = TableElement(
|
||||||
|
label=cluster.label,
|
||||||
|
id=cluster.id,
|
||||||
|
text="",
|
||||||
|
otsl_seq=[],
|
||||||
|
table_cells=[],
|
||||||
|
cluster=cluster,
|
||||||
|
page_no=page.page_no,
|
||||||
|
)
|
||||||
|
|
||||||
|
elements.append(tbl)
|
||||||
|
body.append(tbl)
|
||||||
|
elif cluster.label == LayoutModel.FIGURE_LABEL:
|
||||||
|
fig = None
|
||||||
|
if page.predictions.figures_classification:
|
||||||
|
fig = page.predictions.figures_classification.figure_map.get(
|
||||||
|
cluster.id, None
|
||||||
|
)
|
||||||
|
if (
|
||||||
|
not fig
|
||||||
|
): # fallback: add figure without classification, if it isn't present
|
||||||
|
fig = FigureElement(
|
||||||
|
label=cluster.label,
|
||||||
|
id=cluster.id,
|
||||||
|
text="",
|
||||||
|
data=None,
|
||||||
|
cluster=cluster,
|
||||||
|
page_no=page.page_no,
|
||||||
|
)
|
||||||
|
elements.append(fig)
|
||||||
|
body.append(fig)
|
||||||
|
elif cluster.label == LayoutModel.FORMULA_LABEL:
|
||||||
|
equation = None
|
||||||
|
if page.predictions.equations_prediction:
|
||||||
|
equation = (
|
||||||
|
page.predictions.equations_prediction.equation_map.get(
|
||||||
|
cluster.id, None
|
||||||
|
)
|
||||||
|
)
|
||||||
|
if not equation: # fallback: add empty formula, if it isn't present
|
||||||
|
text = self.sanitize_text(
|
||||||
|
[
|
||||||
|
cell.text.replace("\x02", "-").strip()
|
||||||
|
for cell in cluster.cells
|
||||||
|
if len(cell.text.strip()) > 0
|
||||||
|
]
|
||||||
|
)
|
||||||
|
equation = TextElement(
|
||||||
|
label=cluster.label,
|
||||||
|
id=cluster.id,
|
||||||
|
cluster=cluster,
|
||||||
|
page_no=page.page_no,
|
||||||
|
text=text,
|
||||||
|
)
|
||||||
|
elements.append(equation)
|
||||||
|
body.append(equation)
|
||||||
|
|
||||||
|
page.assembled = AssembledUnit(
|
||||||
|
elements=elements, headers=headers, body=body
|
||||||
|
)
|
||||||
|
|
||||||
|
yield page
|
||||||
114
docling/models/table_structure_model.py
Normal file
114
docling/models/table_structure_model.py
Normal file
@ -0,0 +1,114 @@
|
|||||||
|
from typing import Iterable
|
||||||
|
|
||||||
|
import numpy
|
||||||
|
from docling_ibm_models.tableformer.data_management.tf_predictor import TFPredictor
|
||||||
|
|
||||||
|
from docling.datamodel.base_models import (
|
||||||
|
BoundingBox,
|
||||||
|
Page,
|
||||||
|
TableCell,
|
||||||
|
TableElement,
|
||||||
|
TableStructurePrediction,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class TableStructureModel:
|
||||||
|
def __init__(self, config):
|
||||||
|
self.config = config
|
||||||
|
self.do_cell_matching = config["do_cell_matching"]
|
||||||
|
|
||||||
|
self.enabled = config["enabled"]
|
||||||
|
if self.enabled:
|
||||||
|
artifacts_path = config["artifacts_path"]
|
||||||
|
# Third Party
|
||||||
|
import docling_ibm_models.tableformer.common as c
|
||||||
|
|
||||||
|
self.tm_config = c.read_config(f"{artifacts_path}/tm_config.json")
|
||||||
|
self.tm_config["model"]["save_dir"] = artifacts_path
|
||||||
|
self.tm_model_type = self.tm_config["model"]["type"]
|
||||||
|
|
||||||
|
self.tf_predictor = TFPredictor(self.tm_config)
|
||||||
|
|
||||||
|
def __call__(self, page_batch: Iterable[Page]) -> Iterable[Page]:
|
||||||
|
|
||||||
|
if not self.enabled:
|
||||||
|
yield from page_batch
|
||||||
|
return
|
||||||
|
|
||||||
|
for page in page_batch:
|
||||||
|
page.predictions.tablestructure = TableStructurePrediction() # dummy
|
||||||
|
|
||||||
|
in_tables = [
|
||||||
|
(
|
||||||
|
cluster,
|
||||||
|
[
|
||||||
|
round(cluster.bbox.l),
|
||||||
|
round(cluster.bbox.t),
|
||||||
|
round(cluster.bbox.r),
|
||||||
|
round(cluster.bbox.b),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
for cluster in page.predictions.layout.clusters
|
||||||
|
if cluster.label == "Table"
|
||||||
|
]
|
||||||
|
if not len(in_tables):
|
||||||
|
yield page
|
||||||
|
continue
|
||||||
|
|
||||||
|
tokens = []
|
||||||
|
for c in page.cells:
|
||||||
|
for cluster, _ in in_tables:
|
||||||
|
if c.bbox.area() > 0:
|
||||||
|
if (
|
||||||
|
c.bbox.intersection_area_with(cluster.bbox) / c.bbox.area()
|
||||||
|
> 0.2
|
||||||
|
):
|
||||||
|
# Only allow non empty stings (spaces) into the cells of a table
|
||||||
|
if len(c.text.strip()) > 0:
|
||||||
|
tokens.append(c.model_dump())
|
||||||
|
|
||||||
|
iocr_page = {
|
||||||
|
"image": numpy.asarray(page.image),
|
||||||
|
"tokens": tokens,
|
||||||
|
"width": page.size.width,
|
||||||
|
"height": page.size.height,
|
||||||
|
}
|
||||||
|
|
||||||
|
table_clusters, table_bboxes = zip(*in_tables)
|
||||||
|
|
||||||
|
if len(table_bboxes):
|
||||||
|
tf_output = self.tf_predictor.multi_table_predict(
|
||||||
|
iocr_page, table_bboxes, do_matching=self.do_cell_matching
|
||||||
|
)
|
||||||
|
|
||||||
|
for table_cluster, table_out in zip(table_clusters, tf_output):
|
||||||
|
table_cells = []
|
||||||
|
for element in table_out["tf_responses"]:
|
||||||
|
|
||||||
|
if not self.do_cell_matching:
|
||||||
|
the_bbox = BoundingBox.model_validate(element["bbox"])
|
||||||
|
text_piece = page._backend.get_text_in_rect(the_bbox)
|
||||||
|
element["bbox"]["token"] = text_piece
|
||||||
|
|
||||||
|
tc = TableCell.model_validate(element)
|
||||||
|
table_cells.append(tc)
|
||||||
|
|
||||||
|
# Retrieving cols/rows, after post processing:
|
||||||
|
num_rows = table_out["predict_details"]["num_rows"]
|
||||||
|
num_cols = table_out["predict_details"]["num_cols"]
|
||||||
|
otsl_seq = table_out["predict_details"]["prediction"]["rs_seq"]
|
||||||
|
|
||||||
|
tbl = TableElement(
|
||||||
|
otsl_seq=otsl_seq,
|
||||||
|
table_cells=table_cells,
|
||||||
|
num_rows=num_rows,
|
||||||
|
num_cols=num_cols,
|
||||||
|
id=table_cluster.id,
|
||||||
|
page_no=page.page_no,
|
||||||
|
cluster=table_cluster,
|
||||||
|
label="Table",
|
||||||
|
)
|
||||||
|
|
||||||
|
page.predictions.tablestructure.table_map[table_cluster.id] = tbl
|
||||||
|
|
||||||
|
yield page
|
||||||
0
docling/pipeline/__init__.py
Normal file
0
docling/pipeline/__init__.py
Normal file
18
docling/pipeline/base_model_pipeline.py
Normal file
18
docling/pipeline/base_model_pipeline.py
Normal file
@ -0,0 +1,18 @@
|
|||||||
|
from abc import abstractmethod
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import Iterable
|
||||||
|
|
||||||
|
from docling.datamodel.base_models import Page, PipelineOptions
|
||||||
|
|
||||||
|
|
||||||
|
class BaseModelPipeline:
|
||||||
|
def __init__(self, artifacts_path: Path, pipeline_options: PipelineOptions):
|
||||||
|
self.model_pipe = []
|
||||||
|
self.artifacts_path = artifacts_path
|
||||||
|
self.pipeline_options = pipeline_options
|
||||||
|
|
||||||
|
def apply(self, page_batch: Iterable[Page]) -> Iterable[Page]:
|
||||||
|
for model in self.model_pipe:
|
||||||
|
page_batch = model(page_batch)
|
||||||
|
|
||||||
|
yield from page_batch
|
||||||
40
docling/pipeline/standard_model_pipeline.py
Normal file
40
docling/pipeline/standard_model_pipeline.py
Normal file
@ -0,0 +1,40 @@
|
|||||||
|
from pathlib import Path
|
||||||
|
from typing import Iterable
|
||||||
|
|
||||||
|
from docling.datamodel.base_models import Page, PipelineOptions
|
||||||
|
from docling.models.easyocr_model import EasyOcrModel
|
||||||
|
from docling.models.layout_model import LayoutModel
|
||||||
|
from docling.models.page_assemble_model import PageAssembleModel
|
||||||
|
from docling.models.table_structure_model import TableStructureModel
|
||||||
|
from docling.pipeline.base_model_pipeline import BaseModelPipeline
|
||||||
|
|
||||||
|
|
||||||
|
class StandardModelPipeline(BaseModelPipeline):
|
||||||
|
_layout_model_path = "model_artifacts/layout/beehive_v0.0.5"
|
||||||
|
_table_model_path = "model_artifacts/tableformer"
|
||||||
|
|
||||||
|
def __init__(self, artifacts_path: Path, pipeline_options: PipelineOptions):
|
||||||
|
super().__init__(artifacts_path, pipeline_options)
|
||||||
|
|
||||||
|
self.model_pipe = [
|
||||||
|
EasyOcrModel(
|
||||||
|
config={
|
||||||
|
"lang": ["fr", "de", "es", "en"],
|
||||||
|
"enabled": pipeline_options.do_ocr,
|
||||||
|
}
|
||||||
|
),
|
||||||
|
LayoutModel(
|
||||||
|
config={
|
||||||
|
"artifacts_path": artifacts_path
|
||||||
|
/ StandardModelPipeline._layout_model_path
|
||||||
|
}
|
||||||
|
),
|
||||||
|
TableStructureModel(
|
||||||
|
config={
|
||||||
|
"artifacts_path": artifacts_path
|
||||||
|
/ StandardModelPipeline._table_model_path,
|
||||||
|
"enabled": pipeline_options.do_table_structure,
|
||||||
|
"do_cell_matching": False,
|
||||||
|
}
|
||||||
|
),
|
||||||
|
]
|
||||||
0
docling/utils/__init__.py
Normal file
0
docling/utils/__init__.py
Normal file
806
docling/utils/layout_utils.py
Normal file
806
docling/utils/layout_utils.py
Normal file
@ -0,0 +1,806 @@
|
|||||||
|
import copy
|
||||||
|
import logging
|
||||||
|
|
||||||
|
import networkx as nx
|
||||||
|
|
||||||
|
logger = logging.getLogger("layout_utils")
|
||||||
|
|
||||||
|
|
||||||
|
## -------------------------------
|
||||||
|
## Geometric helper functions
|
||||||
|
## The coordinates grow left to right, and bottom to top.
|
||||||
|
## The bounding box list elements 0 to 3 are x_left, y_bottom, x_right, y_top.
|
||||||
|
|
||||||
|
|
||||||
|
def area(bbox):
|
||||||
|
return (bbox[2] - bbox[0]) * (bbox[3] - bbox[1])
|
||||||
|
|
||||||
|
|
||||||
|
def contains(bbox_i, bbox_j):
|
||||||
|
## Returns True if bbox_i contains bbox_j, else False
|
||||||
|
return (
|
||||||
|
bbox_i[0] <= bbox_j[0]
|
||||||
|
and bbox_i[1] <= bbox_j[1]
|
||||||
|
and bbox_i[2] >= bbox_j[2]
|
||||||
|
and bbox_i[3] >= bbox_j[3]
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def is_intersecting(bbox_i, bbox_j):
|
||||||
|
return not (
|
||||||
|
bbox_i[2] < bbox_j[0]
|
||||||
|
or bbox_i[0] > bbox_j[2]
|
||||||
|
or bbox_i[3] < bbox_j[1]
|
||||||
|
or bbox_i[1] > bbox_j[3]
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def bb_iou(boxA, boxB):
|
||||||
|
# determine the (x, y)-coordinates of the intersection rectangle
|
||||||
|
xA = max(boxA[0], boxB[0])
|
||||||
|
yA = max(boxA[1], boxB[1])
|
||||||
|
xB = min(boxA[2], boxB[2])
|
||||||
|
yB = min(boxA[3], boxB[3])
|
||||||
|
# compute the area of intersection rectangle
|
||||||
|
interArea = max(0, xB - xA + 1) * max(0, yB - yA + 1)
|
||||||
|
# compute the area of both the prediction and ground-truth
|
||||||
|
# rectangles
|
||||||
|
boxAArea = (boxA[2] - boxA[0] + 1) * (boxA[3] - boxA[1] + 1)
|
||||||
|
boxBArea = (boxB[2] - boxB[0] + 1) * (boxB[3] - boxB[1] + 1)
|
||||||
|
# compute the intersection over union by taking the intersection
|
||||||
|
# area and dividing it by the sum of prediction + ground-truth
|
||||||
|
# areas - the interesection area
|
||||||
|
iou = interArea / float(boxAArea + boxBArea - interArea)
|
||||||
|
# return the intersection over union value
|
||||||
|
return iou
|
||||||
|
|
||||||
|
|
||||||
|
def compute_intersection(bbox_i, bbox_j):
|
||||||
|
## Returns the size of the intersection area of the two boxes
|
||||||
|
if not is_intersecting(bbox_i, bbox_j):
|
||||||
|
return 0
|
||||||
|
## Determine the (x, y)-coordinates of the intersection rectangle:
|
||||||
|
xA = max(bbox_i[0], bbox_j[0])
|
||||||
|
yA = max(bbox_i[1], bbox_j[1])
|
||||||
|
xB = min(bbox_i[2], bbox_j[2])
|
||||||
|
yB = min(bbox_i[3], bbox_j[3])
|
||||||
|
## Compute the area of intersection rectangle:
|
||||||
|
interArea = (xB - xA) * (yB - yA)
|
||||||
|
if interArea < 0:
|
||||||
|
logger.debug("Warning: Negative intersection detected!")
|
||||||
|
return 0
|
||||||
|
return interArea
|
||||||
|
|
||||||
|
|
||||||
|
def surrounding(bbox_i, bbox_j):
|
||||||
|
## Computes minimal box that contains both input boxes
|
||||||
|
sbox = []
|
||||||
|
sbox.append(min(bbox_i[0], bbox_j[0]))
|
||||||
|
sbox.append(min(bbox_i[1], bbox_j[1]))
|
||||||
|
sbox.append(max(bbox_i[2], bbox_j[2]))
|
||||||
|
sbox.append(max(bbox_i[3], bbox_j[3]))
|
||||||
|
return sbox
|
||||||
|
|
||||||
|
|
||||||
|
def surrounding_list(bbox_list):
|
||||||
|
## Computes minimal box that contains all boxes in the input list
|
||||||
|
## The list should be non-empty, but just in case it's not:
|
||||||
|
if len(bbox_list) == 0:
|
||||||
|
sbox = [0, 0, 0, 0]
|
||||||
|
else:
|
||||||
|
sbox = []
|
||||||
|
sbox.append(min([bbox[0] for bbox in bbox_list]))
|
||||||
|
sbox.append(min([bbox[1] for bbox in bbox_list]))
|
||||||
|
sbox.append(max([bbox[2] for bbox in bbox_list]))
|
||||||
|
sbox.append(max([bbox[3] for bbox in bbox_list]))
|
||||||
|
return sbox
|
||||||
|
|
||||||
|
|
||||||
|
def vertical_overlap(bboxA, bboxB):
|
||||||
|
## bbox[1] is the lower bound, bbox[3] the upper bound (larger number)
|
||||||
|
if bboxB[3] < bboxA[1]: ## B below A
|
||||||
|
return False
|
||||||
|
elif bboxA[3] < bboxB[1]: ## A below B
|
||||||
|
return False
|
||||||
|
else:
|
||||||
|
return True
|
||||||
|
|
||||||
|
|
||||||
|
def vertical_overlap_fraction(bboxA, bboxB):
|
||||||
|
## Returns the vertical overlap as fraction of the lower bbox height.
|
||||||
|
## bbox[1] is the lower bound, bbox[3] the upper bound (larger number)
|
||||||
|
## Height 0 is permitted in the input.
|
||||||
|
heightA = bboxA[3] - bboxA[1]
|
||||||
|
heightB = bboxB[3] - bboxB[1]
|
||||||
|
min_height = min(heightA, heightB)
|
||||||
|
if bboxA[3] >= bboxB[3]: ## A starts higher or equal
|
||||||
|
if (
|
||||||
|
bboxA[1] <= bboxB[1]
|
||||||
|
): ## B is completely in A; this can include height of B = 0:
|
||||||
|
fraction = 1
|
||||||
|
else:
|
||||||
|
overlap = max(bboxB[3] - bboxA[1], 0)
|
||||||
|
fraction = overlap / max(min_height, 0.001)
|
||||||
|
else:
|
||||||
|
if (
|
||||||
|
bboxB[1] <= bboxA[1]
|
||||||
|
): ## A is completely in B; this can include height of A = 0:
|
||||||
|
fraction = 1
|
||||||
|
else:
|
||||||
|
overlap = max(bboxA[3] - bboxB[1], 0)
|
||||||
|
fraction = overlap / max(min_height, 0.001)
|
||||||
|
return fraction
|
||||||
|
|
||||||
|
|
||||||
|
## -------------------------------
|
||||||
|
## Cluster-and-cell relations
|
||||||
|
|
||||||
|
|
||||||
|
def compute_enclosed_cells(
|
||||||
|
cluster_bbox, raw_cells, min_cell_intersection_with_cluster=0.2
|
||||||
|
):
|
||||||
|
cells_in_cluster = []
|
||||||
|
cells_in_cluster_int = []
|
||||||
|
for ix, cell in enumerate(raw_cells):
|
||||||
|
cell_bbox = cell["bbox"]
|
||||||
|
intersection = compute_intersection(cell_bbox, cluster_bbox)
|
||||||
|
frac_area = area(cell_bbox) * min_cell_intersection_with_cluster
|
||||||
|
|
||||||
|
if (
|
||||||
|
intersection > frac_area and frac_area > 0
|
||||||
|
): # intersect > certain fraction of cell
|
||||||
|
cells_in_cluster.append(ix)
|
||||||
|
cells_in_cluster_int.append(intersection)
|
||||||
|
elif contains(
|
||||||
|
cluster_bbox,
|
||||||
|
[cell_bbox[0] + 3, cell_bbox[1] + 3, cell_bbox[2] - 3, cell_bbox[3] - 3],
|
||||||
|
):
|
||||||
|
cells_in_cluster.append(ix)
|
||||||
|
return cells_in_cluster, cells_in_cluster_int
|
||||||
|
|
||||||
|
|
||||||
|
def find_clusters_around_cells(cell_count, clusters):
|
||||||
|
## Per raw cell, find to which clusters it belongs.
|
||||||
|
## Return list of these indices in the raw-cell order.
|
||||||
|
clusters_around_cells = [[] for _ in range(cell_count)]
|
||||||
|
for cl_ix, cluster in enumerate(clusters):
|
||||||
|
for ix in cluster["cell_ids"]:
|
||||||
|
clusters_around_cells[ix].append(cl_ix)
|
||||||
|
return clusters_around_cells
|
||||||
|
|
||||||
|
|
||||||
|
def find_cell_index(raw_ix, cell_array):
|
||||||
|
## "raw_ix" is a rawcell_id.
|
||||||
|
## "cell_array" has the structure of an (annotation) cells array.
|
||||||
|
## Returns index of cell in cell_array that has this rawcell_id.
|
||||||
|
for ix, cell in enumerate(cell_array):
|
||||||
|
if cell["rawcell_id"] == raw_ix:
|
||||||
|
return ix
|
||||||
|
|
||||||
|
|
||||||
|
def find_cell_indices(cluster, cell_array):
|
||||||
|
## "cluster" must have the structure as in a clusters array in a prediction,
|
||||||
|
## "cell_array" that of a cells array.
|
||||||
|
## Returns list of indices of cells in cell_array that have the rawcell_ids as in the cluster,
|
||||||
|
## in the order of the rawcell_ids.
|
||||||
|
result = []
|
||||||
|
for raw_ix in sorted(cluster["cell_ids"]):
|
||||||
|
## Find the cell with this rawcell_id (if any)
|
||||||
|
for ix, cell in enumerate(cell_array):
|
||||||
|
if cell["rawcell_id"] == raw_ix:
|
||||||
|
result.append(ix)
|
||||||
|
return result
|
||||||
|
|
||||||
|
|
||||||
|
def find_first_cell_index(cluster, cell_array):
|
||||||
|
## "cluster" must be a dict with key "cell_ids"; it can also be a line.
|
||||||
|
## "cell_array" has the structure of a cells array in an annotation.
|
||||||
|
## Returns index of cell in cell_array that has the lowest rawcell_id from the cluster.
|
||||||
|
result = [] ## We keep it a list as it can be empty (picture without text cells)
|
||||||
|
if len(cluster["cell_ids"]) == 0:
|
||||||
|
return result
|
||||||
|
raw_ix = min(cluster["cell_ids"])
|
||||||
|
## Find the cell with this rawcell_id (if any)
|
||||||
|
for ix, cell in enumerate(cell_array):
|
||||||
|
if cell["rawcell_id"] == raw_ix:
|
||||||
|
result.append(ix)
|
||||||
|
break ## One is enough; should be only one anyway.
|
||||||
|
if result == []:
|
||||||
|
logger.debug(
|
||||||
|
" Warning: Raw cell " + str(raw_ix) + " not found in annotation cells"
|
||||||
|
)
|
||||||
|
return result
|
||||||
|
|
||||||
|
|
||||||
|
## -------------------------------
|
||||||
|
## Cluster labels and text
|
||||||
|
|
||||||
|
|
||||||
|
def relabel_cluster(cluster, cl_ix, new_label, target_pred):
|
||||||
|
## "cluster" must have the structure as in a clusters array in a prediction,
|
||||||
|
## "cl_ix" is its index in target_pred,
|
||||||
|
## "new_label" is the intended new label,
|
||||||
|
## "target_pred" is the entire current target prediction.
|
||||||
|
## Sets label on the cluster itself, and on the cells in the target_pred.
|
||||||
|
## Returns new_label so that also the cl_label variable in the main code is easily set.
|
||||||
|
target_pred["clusters"][cl_ix]["type"] = new_label
|
||||||
|
cluster_target_cells = find_cell_indices(cluster, target_pred["cells"])
|
||||||
|
for ix in cluster_target_cells:
|
||||||
|
target_pred["cells"][ix]["label"] = new_label
|
||||||
|
return new_label
|
||||||
|
|
||||||
|
|
||||||
|
def find_cluster_text(cluster, raw_cells):
|
||||||
|
## "cluster" must be a dict with "cell_ids"; it can also be a line.
|
||||||
|
## "raw_cells" must have the format of item["raw"]["cells"]
|
||||||
|
## Returns the text of the cluster, with blanks between the cell contents
|
||||||
|
## (which seem to be words or phrases without starting or trailing blanks).
|
||||||
|
## Note that in formulas, this may give a lot more blanks than originally
|
||||||
|
cluster_text = ""
|
||||||
|
for raw_ix in sorted(cluster["cell_ids"]):
|
||||||
|
cluster_text = cluster_text + raw_cells[raw_ix]["text"] + " "
|
||||||
|
return cluster_text.rstrip()
|
||||||
|
|
||||||
|
|
||||||
|
def find_cluster_text_without_blanks(cluster, raw_cells):
|
||||||
|
## "cluster" must be a dict with "cell_ids"; it can also be a line.
|
||||||
|
## "raw_cells" must have the format of item["raw"]["cells"]
|
||||||
|
## Returns the text of the cluster, without blanks between the cell contents
|
||||||
|
## Interesting in formula analysis.
|
||||||
|
cluster_text = ""
|
||||||
|
for raw_ix in sorted(cluster["cell_ids"]):
|
||||||
|
cluster_text = cluster_text + raw_cells[raw_ix]["text"]
|
||||||
|
return cluster_text.rstrip()
|
||||||
|
|
||||||
|
|
||||||
|
## -------------------------------
|
||||||
|
## Clusters and lines
|
||||||
|
## (Most line-oriented functions are only needed in TextAnalysisGivenClusters,
|
||||||
|
## but this one also in FormulaAnalysis)
|
||||||
|
|
||||||
|
|
||||||
|
def build_cluster_from_lines(lines, label, id):
|
||||||
|
## Lines must be a non-empty list of dicts (lines) with elements "cell_ids" and "bbox"
|
||||||
|
## (There is no condition that they are really geometrically lines)
|
||||||
|
## A cluster in standard format is returned with given label and id
|
||||||
|
local_lines = copy.deepcopy(
|
||||||
|
lines
|
||||||
|
) ## without this, it changes "lines" also outside this function
|
||||||
|
first_line = local_lines.pop(0)
|
||||||
|
cluster = {
|
||||||
|
"id": id,
|
||||||
|
"type": label,
|
||||||
|
"cell_ids": first_line["cell_ids"],
|
||||||
|
"bbox": first_line["bbox"],
|
||||||
|
"confidence": 0,
|
||||||
|
"created_by": "merged_cells",
|
||||||
|
}
|
||||||
|
confidence = 0
|
||||||
|
counter = 0
|
||||||
|
for line in local_lines:
|
||||||
|
new_cell_ids = cluster["cell_ids"] + line["cell_ids"]
|
||||||
|
cluster["cell_ids"] = new_cell_ids
|
||||||
|
cluster["bbox"] = surrounding(cluster["bbox"], line["bbox"])
|
||||||
|
counter += 1
|
||||||
|
confidence += line["confidence"]
|
||||||
|
confidence = confidence / counter
|
||||||
|
cluster["confidence"] = confidence
|
||||||
|
return cluster
|
||||||
|
|
||||||
|
|
||||||
|
## -------------------------------
|
||||||
|
## Reading order
|
||||||
|
|
||||||
|
|
||||||
|
def produce_reading_order(clusters, cluster_sort_type, cell_sort_type, sort_ids):
|
||||||
|
## In:
|
||||||
|
## Clusters: list as in predictions.
|
||||||
|
## cluster_sort_type: string, currently only "raw_cells".
|
||||||
|
## cell_sort_type: string, currently only "raw_cells".
|
||||||
|
## sort_ids: Boolean, whether the cluster ids should be adapted to their new position
|
||||||
|
## Out: Another clusters list, sorted according to the type.
|
||||||
|
|
||||||
|
logger.debug("---- Start cluster sorting ------")
|
||||||
|
|
||||||
|
if cell_sort_type == "raw_cell_ids":
|
||||||
|
for cl in clusters:
|
||||||
|
sorted_cell_ids = sorted(cl["cell_ids"])
|
||||||
|
cl["cell_ids"] = sorted_cell_ids
|
||||||
|
else:
|
||||||
|
logger.debug(
|
||||||
|
"Unknown cell_sort_type `"
|
||||||
|
+ cell_sort_type
|
||||||
|
+ "`, no cell sorting will happen."
|
||||||
|
)
|
||||||
|
|
||||||
|
if cluster_sort_type == "raw_cell_ids":
|
||||||
|
clusters_with_cells = [cl for cl in clusters if cl["cell_ids"] != []]
|
||||||
|
clusters_without_cells = [cl for cl in clusters if cl["cell_ids"] == []]
|
||||||
|
logger.debug(
|
||||||
|
"Clusters with cells: " + str([cl["id"] for cl in clusters_with_cells])
|
||||||
|
)
|
||||||
|
logger.debug(
|
||||||
|
" Their first cell ids: "
|
||||||
|
+ str([cl["cell_ids"][0] for cl in clusters_with_cells])
|
||||||
|
)
|
||||||
|
logger.debug(
|
||||||
|
"Clusters without cells: "
|
||||||
|
+ str([cl["id"] for cl in clusters_without_cells])
|
||||||
|
)
|
||||||
|
clusters_with_cells_sorted = sorted(
|
||||||
|
clusters_with_cells, key=lambda cluster: cluster["cell_ids"][0]
|
||||||
|
)
|
||||||
|
logger.debug(
|
||||||
|
" First cell ids after sorting: "
|
||||||
|
+ str([cl["cell_ids"][0] for cl in clusters_with_cells_sorted])
|
||||||
|
)
|
||||||
|
sorted_clusters = clusters_with_cells_sorted + clusters_without_cells
|
||||||
|
else:
|
||||||
|
logger.debug(
|
||||||
|
"Unknown cluster_sort_type: `"
|
||||||
|
+ cluster_sort_type
|
||||||
|
+ "`, no cluster sorting will happen."
|
||||||
|
)
|
||||||
|
|
||||||
|
if sort_ids:
|
||||||
|
for i, cl in enumerate(sorted_clusters):
|
||||||
|
cl["id"] = i
|
||||||
|
return sorted_clusters
|
||||||
|
|
||||||
|
|
||||||
|
## -------------------------------
|
||||||
|
## Line Splitting
|
||||||
|
|
||||||
|
|
||||||
|
def sort_cells_horizontal(line_cell_ids, raw_cells):
|
||||||
|
## "line_cells" should be a non-empty list of (raw) cell_ids
|
||||||
|
## "raw_cells" has the structure of item["raw"]["cells"].
|
||||||
|
## Sorts the cells in the line by x0 (left start).
|
||||||
|
new_line_cell_ids = sorted(
|
||||||
|
line_cell_ids, key=lambda cell_id: raw_cells[cell_id]["bbox"][0]
|
||||||
|
)
|
||||||
|
return new_line_cell_ids
|
||||||
|
|
||||||
|
|
||||||
|
def adapt_bboxes(raw_cells, clusters, orphan_cell_indices):
|
||||||
|
new_clusters = []
|
||||||
|
for ix, cluster in enumerate(clusters):
|
||||||
|
new_cluster = copy.deepcopy(cluster)
|
||||||
|
logger.debug(
|
||||||
|
"Treating cluster " + str(ix) + ", type " + str(new_cluster["type"])
|
||||||
|
)
|
||||||
|
logger.debug(" with cells: " + str(new_cluster["cell_ids"]))
|
||||||
|
if len(cluster["cell_ids"]) == 0 and cluster["type"] != "Picture":
|
||||||
|
logger.debug(" Empty non-picture, removed")
|
||||||
|
continue ## Skip this former cluster, now without cells.
|
||||||
|
new_bbox = adapt_bbox(raw_cells, new_cluster, orphan_cell_indices)
|
||||||
|
new_cluster["bbox"] = new_bbox
|
||||||
|
new_clusters.append(new_cluster)
|
||||||
|
return new_clusters
|
||||||
|
|
||||||
|
|
||||||
|
def adapt_bbox(raw_cells, cluster, orphan_cell_indices):
|
||||||
|
if not (cluster["type"] in ["Table", "Picture"]):
|
||||||
|
## A text-like cluster. The bbox only needs to be around the text cells:
|
||||||
|
logger.debug(" Initial bbox: " + str(cluster["bbox"]))
|
||||||
|
new_bbox = surrounding_list(
|
||||||
|
[raw_cells[cid]["bbox"] for cid in cluster["cell_ids"]]
|
||||||
|
)
|
||||||
|
logger.debug(" New bounding box:" + str(new_bbox))
|
||||||
|
if cluster["type"] == "Picture":
|
||||||
|
## We only make the bbox completely comprise included text cells:
|
||||||
|
logger.debug(" Picture")
|
||||||
|
if len(cluster["cell_ids"]) != 0:
|
||||||
|
min_bbox = surrounding_list(
|
||||||
|
[raw_cells[cid]["bbox"] for cid in cluster["cell_ids"]]
|
||||||
|
)
|
||||||
|
logger.debug(" Minimum bbox: " + str(min_bbox))
|
||||||
|
logger.debug(" Initial bbox: " + str(cluster["bbox"]))
|
||||||
|
new_bbox = surrounding(min_bbox, cluster["bbox"])
|
||||||
|
logger.debug(" New bbox (initial and text cells): " + str(new_bbox))
|
||||||
|
else:
|
||||||
|
logger.debug(" without text cells, no change.")
|
||||||
|
new_bbox = cluster["bbox"]
|
||||||
|
else: ## A table
|
||||||
|
## At least we have to keep the included text cells, and we make the bbox completely comprise them
|
||||||
|
min_bbox = surrounding_list(
|
||||||
|
[raw_cells[cid]["bbox"] for cid in cluster["cell_ids"]]
|
||||||
|
)
|
||||||
|
logger.debug(" Minimum bbox: " + str(min_bbox))
|
||||||
|
logger.debug(" Initial bbox: " + str(cluster["bbox"]))
|
||||||
|
new_bbox = surrounding(min_bbox, cluster["bbox"])
|
||||||
|
logger.debug(" Possibly increased bbox: " + str(new_bbox))
|
||||||
|
|
||||||
|
## Now we look which non-belonging cells are covered.
|
||||||
|
## (To decrease dependencies, we don't make use of which cells we actually removed.)
|
||||||
|
## We don't worry about orphan cells, those could still be added to the table.
|
||||||
|
enclosed_cells = compute_enclosed_cells(
|
||||||
|
new_bbox, raw_cells, min_cell_intersection_with_cluster=0.3
|
||||||
|
)[0]
|
||||||
|
additional_cells = set(enclosed_cells) - set(cluster["cell_ids"])
|
||||||
|
logger.debug(
|
||||||
|
" Additional cells enclosed by Table bbox: " + str(additional_cells)
|
||||||
|
)
|
||||||
|
spurious_cells = additional_cells - set(orphan_cell_indices)
|
||||||
|
logger.debug(
|
||||||
|
" Spurious cells enclosed by Table bbox (additional minus orphans): "
|
||||||
|
+ str(spurious_cells)
|
||||||
|
)
|
||||||
|
if len(spurious_cells) == 0:
|
||||||
|
return new_bbox
|
||||||
|
|
||||||
|
## Else we want to keep as much as possible, e.g., grid lines, but not the spurious cells if we can.
|
||||||
|
## We initialize possible cuts with the current bbox.
|
||||||
|
left_cut = new_bbox[0]
|
||||||
|
right_cut = new_bbox[2]
|
||||||
|
upper_cut = new_bbox[3]
|
||||||
|
lower_cut = new_bbox[1]
|
||||||
|
|
||||||
|
for cell_ix in spurious_cells:
|
||||||
|
cell = raw_cells[cell_ix]
|
||||||
|
# logger.debug(" Spurious cell bbox: " + str(cell["bbox"]))
|
||||||
|
is_left = cell["bbox"][2] < min_bbox[0]
|
||||||
|
is_right = cell["bbox"][0] > min_bbox[2]
|
||||||
|
is_above = cell["bbox"][1] > min_bbox[3]
|
||||||
|
is_below = cell["bbox"][3] < min_bbox[1]
|
||||||
|
# logger.debug(" Left, right, above, below? " + str([is_left, is_right, is_above, is_below]))
|
||||||
|
|
||||||
|
if is_left:
|
||||||
|
if cell["bbox"][2] > left_cut:
|
||||||
|
## We move the left cut to exclude this cell:
|
||||||
|
left_cut = cell["bbox"][2]
|
||||||
|
if is_right:
|
||||||
|
if cell["bbox"][0] < right_cut:
|
||||||
|
## We move the right cut to exclude this cell:
|
||||||
|
right_cut = cell["bbox"][0]
|
||||||
|
if is_above:
|
||||||
|
if cell["bbox"][1] < upper_cut:
|
||||||
|
## We move the upper cut to exclude this cell:
|
||||||
|
upper_cut = cell["bbox"][1]
|
||||||
|
if is_below:
|
||||||
|
if cell["bbox"][3] > lower_cut:
|
||||||
|
## We move the left cut to exclude this cell:
|
||||||
|
lower_cut = cell["bbox"][3]
|
||||||
|
# logger.debug(" Current bbox: " + str([left_cut, lower_cut, right_cut, upper_cut]))
|
||||||
|
|
||||||
|
new_bbox = [left_cut, lower_cut, right_cut, upper_cut]
|
||||||
|
|
||||||
|
logger.debug(" Final bbox: " + str(new_bbox))
|
||||||
|
return new_bbox
|
||||||
|
|
||||||
|
|
||||||
|
def remove_cluster_duplicates_by_conf(cluster_predictions, threshold=0.5):
|
||||||
|
DuplicateDeletedClusterIDs = []
|
||||||
|
for cluster_1 in cluster_predictions:
|
||||||
|
for cluster_2 in cluster_predictions:
|
||||||
|
if cluster_1["id"] != cluster_2["id"]:
|
||||||
|
if_conf = False
|
||||||
|
if cluster_1["confidence"] > cluster_2["confidence"]:
|
||||||
|
if_conf = True
|
||||||
|
if if_conf == True:
|
||||||
|
if bb_iou(cluster_1["bbox"], cluster_2["bbox"]) > threshold:
|
||||||
|
DuplicateDeletedClusterIDs.append(cluster_2["id"])
|
||||||
|
elif contains(
|
||||||
|
cluster_1["bbox"],
|
||||||
|
[
|
||||||
|
cluster_2["bbox"][0] + 3,
|
||||||
|
cluster_2["bbox"][1] + 3,
|
||||||
|
cluster_2["bbox"][2] - 3,
|
||||||
|
cluster_2["bbox"][3] - 3,
|
||||||
|
],
|
||||||
|
):
|
||||||
|
DuplicateDeletedClusterIDs.append(cluster_2["id"])
|
||||||
|
|
||||||
|
DuplicateDeletedClusterIDs = list(set(DuplicateDeletedClusterIDs))
|
||||||
|
|
||||||
|
for cl_id in DuplicateDeletedClusterIDs:
|
||||||
|
for cluster in cluster_predictions:
|
||||||
|
if cl_id == cluster["id"]:
|
||||||
|
cluster_predictions.remove(cluster)
|
||||||
|
return cluster_predictions
|
||||||
|
|
||||||
|
|
||||||
|
# Assign orphan cells by a low confidence prediction that is below the assigned confidence
|
||||||
|
def assign_orphans_with_low_conf_pred(
|
||||||
|
cluster_predictions, cluster_predictions_low, raw_cells, orphan_cell_indices
|
||||||
|
):
|
||||||
|
for orph_id in orphan_cell_indices:
|
||||||
|
cluster_chosen = {}
|
||||||
|
iou_thresh = 0.05
|
||||||
|
confidence = 0.05
|
||||||
|
|
||||||
|
# Loop over all predictions, and find the one with the highest IOU, and confidence
|
||||||
|
for cluster in cluster_predictions_low:
|
||||||
|
calc_iou = bb_iou(cluster["bbox"], raw_cells[orph_id]["bbox"])
|
||||||
|
cluster_area = (cluster["bbox"][3] - cluster["bbox"][1]) * (
|
||||||
|
cluster["bbox"][2] - cluster["bbox"][0]
|
||||||
|
)
|
||||||
|
cell_area = (
|
||||||
|
raw_cells[orph_id]["bbox"][3] - raw_cells[orph_id]["bbox"][1]
|
||||||
|
) * (raw_cells[orph_id]["bbox"][2] - raw_cells[orph_id]["bbox"][0])
|
||||||
|
|
||||||
|
if (
|
||||||
|
(iou_thresh < calc_iou)
|
||||||
|
and (cluster["confidence"] > confidence)
|
||||||
|
and (cell_area * 3 > cluster_area)
|
||||||
|
):
|
||||||
|
cluster_chosen = cluster
|
||||||
|
iou_thresh = calc_iou
|
||||||
|
confidence = cluster["confidence"]
|
||||||
|
# If a candidate is found, assign to it the PDF cell ids, and tag that it was created by this function for tracking
|
||||||
|
if iou_thresh != 0.05 and confidence != 0.05:
|
||||||
|
cluster_chosen["cell_ids"].append(orph_id)
|
||||||
|
cluster_chosen["created_by"] = "orph_low_conf"
|
||||||
|
cluster_predictions.append(cluster_chosen)
|
||||||
|
orphan_cell_indices.remove(orph_id)
|
||||||
|
return cluster_predictions, orphan_cell_indices
|
||||||
|
|
||||||
|
|
||||||
|
def remove_ambigous_pdf_cell_by_conf(cluster_predictions, raw_cells, amb_cell_idxs):
|
||||||
|
for amb_cell_id in amb_cell_idxs:
|
||||||
|
highest_conf = 0
|
||||||
|
highest_bbox_iou = 0
|
||||||
|
cluster_chosen = None
|
||||||
|
problamatic_clusters = []
|
||||||
|
|
||||||
|
# Find clusters in question
|
||||||
|
for cluster in cluster_predictions:
|
||||||
|
|
||||||
|
if amb_cell_id in cluster["cell_ids"]:
|
||||||
|
problamatic_clusters.append(amb_cell_id)
|
||||||
|
|
||||||
|
# If the cell_id is in a cluster of high conf, and highest iou score, and smaller in area
|
||||||
|
bbox_iou_val = bb_iou(cluster["bbox"], raw_cells[amb_cell_id]["bbox"])
|
||||||
|
|
||||||
|
if (
|
||||||
|
cluster["confidence"] > highest_conf
|
||||||
|
and bbox_iou_val > highest_bbox_iou
|
||||||
|
):
|
||||||
|
cluster_chosen = cluster
|
||||||
|
highest_conf = cluster["confidence"]
|
||||||
|
highest_bbox_iou = bbox_iou_val
|
||||||
|
if cluster["id"] in problamatic_clusters:
|
||||||
|
problamatic_clusters.remove(cluster["id"])
|
||||||
|
|
||||||
|
# now remove the assigning of cell id from lower confidence, and threshold
|
||||||
|
for cluster in cluster_predictions:
|
||||||
|
for prob_amb_id in problamatic_clusters:
|
||||||
|
if prob_amb_id in cluster["cell_ids"]:
|
||||||
|
cluster["cell_ids"].remove(prob_amb_id)
|
||||||
|
amb_cell_idxs.remove(amb_cell_id)
|
||||||
|
|
||||||
|
return cluster_predictions, amb_cell_idxs
|
||||||
|
|
||||||
|
|
||||||
|
def ranges(nums):
|
||||||
|
# Find if consecutive numbers exist within pdf cells
|
||||||
|
# Used to remove line numbers for review manuscripts
|
||||||
|
nums = sorted(set(nums))
|
||||||
|
gaps = [[s, e] for s, e in zip(nums, nums[1:]) if s + 1 < e]
|
||||||
|
edges = iter(nums[:1] + sum(gaps, []) + nums[-1:])
|
||||||
|
return list(zip(edges, edges))
|
||||||
|
|
||||||
|
|
||||||
|
def set_orphan_as_text(
|
||||||
|
cluster_predictions, cluster_predictions_low, raw_cells, orphan_cell_indices
|
||||||
|
):
|
||||||
|
max_id = -1
|
||||||
|
figures = []
|
||||||
|
for cluster in cluster_predictions:
|
||||||
|
if cluster["type"] == "Picture":
|
||||||
|
figures.append(cluster)
|
||||||
|
|
||||||
|
if cluster["id"] > max_id:
|
||||||
|
max_id = cluster["id"]
|
||||||
|
max_id += 1
|
||||||
|
|
||||||
|
lines_detector = False
|
||||||
|
content_of_orphans = []
|
||||||
|
for orph_id in orphan_cell_indices:
|
||||||
|
orph_cell = raw_cells[orph_id]
|
||||||
|
content_of_orphans.append(raw_cells[orph_id]["text"])
|
||||||
|
|
||||||
|
fil_content_of_orphans = []
|
||||||
|
for cell_content in content_of_orphans:
|
||||||
|
if cell_content.isnumeric():
|
||||||
|
try:
|
||||||
|
num = int(cell_content)
|
||||||
|
fil_content_of_orphans.append(num)
|
||||||
|
except ValueError: # ignore the cell
|
||||||
|
pass
|
||||||
|
|
||||||
|
# line_orphans = []
|
||||||
|
# Check if there are more than 2 pdf orphan cells, if there are more than 2,
|
||||||
|
# then check between the orphan cells if they are numeric
|
||||||
|
# and if they are a consecutive series of numbers (using ranges function) to decide
|
||||||
|
|
||||||
|
if len(fil_content_of_orphans) > 2:
|
||||||
|
out_ranges = ranges(fil_content_of_orphans)
|
||||||
|
if len(out_ranges) > 1:
|
||||||
|
cnt_range = 0
|
||||||
|
for ranges_ in out_ranges:
|
||||||
|
if ranges_[0] != ranges_[1]:
|
||||||
|
# If there are more than 75 (half the total line number of a review manuscript page)
|
||||||
|
# decide that there are line numbers on page to be ignored.
|
||||||
|
if len(list(range(ranges_[0], ranges_[1]))) > 75:
|
||||||
|
lines_detector = True
|
||||||
|
# line_orphans = line_orphans + list(range(ranges_[0], ranges_[1]))
|
||||||
|
|
||||||
|
for orph_id in orphan_cell_indices:
|
||||||
|
orph_cell = raw_cells[orph_id]
|
||||||
|
if bool(orph_cell["text"] and not orph_cell["text"].isspace()):
|
||||||
|
fig_flag = False
|
||||||
|
# Do not assign orphan cells if they are inside a figure
|
||||||
|
for fig in figures:
|
||||||
|
if contains(fig["bbox"], orph_cell["bbox"]):
|
||||||
|
fig_flag = True
|
||||||
|
|
||||||
|
# if fig_flag == False and raw_cells[orph_id]["text"] not in line_orphans:
|
||||||
|
if fig_flag == False and lines_detector == False:
|
||||||
|
# get class from low confidence detections if not set as text:
|
||||||
|
class_type = "Text"
|
||||||
|
|
||||||
|
for cluster in cluster_predictions_low:
|
||||||
|
intersection = compute_intersection(
|
||||||
|
orph_cell["bbox"], cluster["bbox"]
|
||||||
|
)
|
||||||
|
class_type = "Text"
|
||||||
|
if (
|
||||||
|
cluster["confidence"] > 0.1
|
||||||
|
and bb_iou(cluster["bbox"], orph_cell["bbox"]) > 0.4
|
||||||
|
):
|
||||||
|
class_type = cluster["type"]
|
||||||
|
elif contains(
|
||||||
|
cluster["bbox"],
|
||||||
|
[
|
||||||
|
orph_cell["bbox"][0] + 3,
|
||||||
|
orph_cell["bbox"][1] + 3,
|
||||||
|
orph_cell["bbox"][2] - 3,
|
||||||
|
orph_cell["bbox"][3] - 3,
|
||||||
|
],
|
||||||
|
):
|
||||||
|
class_type = cluster["type"]
|
||||||
|
elif intersection > area(orph_cell["bbox"]) * 0.2:
|
||||||
|
class_type = cluster["type"]
|
||||||
|
|
||||||
|
new_cluster = {
|
||||||
|
"id": max_id,
|
||||||
|
"bbox": orph_cell["bbox"],
|
||||||
|
"type": class_type,
|
||||||
|
"cell_ids": [orph_id],
|
||||||
|
"confidence": -1,
|
||||||
|
"created_by": "orphan_default",
|
||||||
|
}
|
||||||
|
max_id += 1
|
||||||
|
cluster_predictions.append(new_cluster)
|
||||||
|
return cluster_predictions, orphan_cell_indices
|
||||||
|
|
||||||
|
|
||||||
|
def merge_cells(cluster_predictions):
|
||||||
|
# Using graph component creates clusters if orphan cells are touching or too close.
|
||||||
|
G = nx.Graph()
|
||||||
|
for cluster in cluster_predictions:
|
||||||
|
if cluster["created_by"] == "orphan_default":
|
||||||
|
G.add_node(cluster["id"])
|
||||||
|
|
||||||
|
for cluster_1 in cluster_predictions:
|
||||||
|
for cluster_2 in cluster_predictions:
|
||||||
|
if (
|
||||||
|
cluster_1["id"] != cluster_2["id"]
|
||||||
|
and cluster_2["created_by"] == "orphan_default"
|
||||||
|
and cluster_1["created_by"] == "orphan_default"
|
||||||
|
):
|
||||||
|
cl1 = copy.deepcopy(cluster_1["bbox"])
|
||||||
|
cl2 = copy.deepcopy(cluster_2["bbox"])
|
||||||
|
cl1[0] = cl1[0] - 2
|
||||||
|
cl1[1] = cl1[1] - 2
|
||||||
|
cl1[2] = cl1[2] + 2
|
||||||
|
cl1[3] = cl1[3] + 2
|
||||||
|
cl2[0] = cl2[0] - 2
|
||||||
|
cl2[1] = cl2[1] - 2
|
||||||
|
cl2[2] = cl2[2] + 2
|
||||||
|
cl2[3] = cl2[3] + 2
|
||||||
|
if is_intersecting(cl1, cl2):
|
||||||
|
G.add_edge(cluster_1["id"], cluster_2["id"])
|
||||||
|
|
||||||
|
component = sorted(map(sorted, nx.k_edge_components(G, k=1)))
|
||||||
|
max_id = -1
|
||||||
|
for cluster_1 in cluster_predictions:
|
||||||
|
if cluster_1["id"] > max_id:
|
||||||
|
max_id = cluster_1["id"]
|
||||||
|
|
||||||
|
for nodes in component:
|
||||||
|
if len(nodes) > 1:
|
||||||
|
max_id += 1
|
||||||
|
lines = []
|
||||||
|
for node in nodes:
|
||||||
|
for cluster in cluster_predictions:
|
||||||
|
if cluster["id"] == node:
|
||||||
|
lines.append(cluster)
|
||||||
|
cluster_predictions.remove(cluster)
|
||||||
|
new_merged_cluster = build_cluster_from_lines(lines, "Text", max_id)
|
||||||
|
cluster_predictions.append(new_merged_cluster)
|
||||||
|
return cluster_predictions
|
||||||
|
|
||||||
|
|
||||||
|
def clean_up_clusters(
|
||||||
|
cluster_predictions,
|
||||||
|
raw_cells,
|
||||||
|
merge_cells=False,
|
||||||
|
img_table=False,
|
||||||
|
one_cell_table=False,
|
||||||
|
):
|
||||||
|
DuplicateDeletedClusterIDs = []
|
||||||
|
|
||||||
|
for cluster_1 in cluster_predictions:
|
||||||
|
for cluster_2 in cluster_predictions:
|
||||||
|
if cluster_1["id"] != cluster_2["id"]:
|
||||||
|
# remove any artifcats created by merging clusters
|
||||||
|
if merge_cells == True:
|
||||||
|
if contains(
|
||||||
|
cluster_1["bbox"],
|
||||||
|
[
|
||||||
|
cluster_2["bbox"][0] + 3,
|
||||||
|
cluster_2["bbox"][1] + 3,
|
||||||
|
cluster_2["bbox"][2] - 3,
|
||||||
|
cluster_2["bbox"][3] - 3,
|
||||||
|
],
|
||||||
|
):
|
||||||
|
cluster_1["cell_ids"] = (
|
||||||
|
cluster_1["cell_ids"] + cluster_2["cell_ids"]
|
||||||
|
)
|
||||||
|
DuplicateDeletedClusterIDs.append(cluster_2["id"])
|
||||||
|
# remove clusters that might appear inside tables, or images (such as pdf cells in graphs)
|
||||||
|
elif img_table == True:
|
||||||
|
if (
|
||||||
|
cluster_1["type"] == "Text"
|
||||||
|
and cluster_2["type"] == "Picture"
|
||||||
|
or cluster_2["type"] == "Table"
|
||||||
|
):
|
||||||
|
if bb_iou(cluster_1["bbox"], cluster_2["bbox"]) > 0.5:
|
||||||
|
DuplicateDeletedClusterIDs.append(cluster_1["id"])
|
||||||
|
elif contains(
|
||||||
|
[
|
||||||
|
cluster_2["bbox"][0] - 3,
|
||||||
|
cluster_2["bbox"][1] - 3,
|
||||||
|
cluster_2["bbox"][2] + 3,
|
||||||
|
cluster_2["bbox"][3] + 3,
|
||||||
|
],
|
||||||
|
cluster_1["bbox"],
|
||||||
|
):
|
||||||
|
DuplicateDeletedClusterIDs.append(cluster_1["id"])
|
||||||
|
# remove tables that have one pdf cell
|
||||||
|
if one_cell_table == True:
|
||||||
|
if cluster_1["type"] == "Table" and len(cluster_1["cell_ids"]) < 2:
|
||||||
|
DuplicateDeletedClusterIDs.append(cluster_1["id"])
|
||||||
|
|
||||||
|
DuplicateDeletedClusterIDs = list(set(DuplicateDeletedClusterIDs))
|
||||||
|
|
||||||
|
for cl_id in DuplicateDeletedClusterIDs:
|
||||||
|
for cluster in cluster_predictions:
|
||||||
|
if cl_id == cluster["id"]:
|
||||||
|
cluster_predictions.remove(cluster)
|
||||||
|
return cluster_predictions
|
||||||
|
|
||||||
|
|
||||||
|
def assigning_cell_ids_to_clusters(clusters, raw_cells, threshold):
|
||||||
|
for cluster in clusters:
|
||||||
|
cells_in_cluster, _ = compute_enclosed_cells(
|
||||||
|
cluster["bbox"], raw_cells, min_cell_intersection_with_cluster=threshold
|
||||||
|
)
|
||||||
|
cluster["cell_ids"] = cells_in_cluster
|
||||||
|
## These cell_ids are ids of the raw cells.
|
||||||
|
## They are often, but not always, the same as the "id" or the index of the "cells" list in a prediction.
|
||||||
|
return clusters
|
||||||
|
|
||||||
|
|
||||||
|
# Creates a map of cell_id->cluster_id
|
||||||
|
def cell_id_state_map(clusters, cell_count):
|
||||||
|
clusters_around_cells = find_clusters_around_cells(cell_count, clusters)
|
||||||
|
orphan_cell_indices = [
|
||||||
|
ix for ix in range(cell_count) if len(clusters_around_cells[ix]) == 0
|
||||||
|
] # which cells are assigned no cluster?
|
||||||
|
ambiguous_cell_indices = [
|
||||||
|
ix for ix in range(cell_count) if len(clusters_around_cells[ix]) > 1
|
||||||
|
] # which cells are assigned > 1 clusters?
|
||||||
|
return clusters_around_cells, orphan_cell_indices, ambiguous_cell_indices
|
||||||
41
docling/utils/utils.py
Normal file
41
docling/utils/utils.py
Normal file
@ -0,0 +1,41 @@
|
|||||||
|
import hashlib
|
||||||
|
from io import BytesIO
|
||||||
|
from itertools import islice
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import List, Union
|
||||||
|
|
||||||
|
|
||||||
|
def chunkify(iterator, chunk_size):
|
||||||
|
"""Yield successive chunks of chunk_size from the iterable."""
|
||||||
|
if isinstance(iterator, List):
|
||||||
|
iterator = iter(iterator)
|
||||||
|
for first in iterator: # Take the first element from the iterator
|
||||||
|
yield [first] + list(islice(iterator, chunk_size - 1))
|
||||||
|
|
||||||
|
|
||||||
|
def create_file_hash(path_or_stream: Union[BytesIO, Path]) -> str:
|
||||||
|
"""Create a stable page_hash of the path_or_stream of a file"""
|
||||||
|
|
||||||
|
block_size = 65536
|
||||||
|
hasher = hashlib.sha256()
|
||||||
|
|
||||||
|
def _hash_buf(binary_stream):
|
||||||
|
buf = binary_stream.read(block_size) # read and page_hash in chunks
|
||||||
|
while len(buf) > 0:
|
||||||
|
hasher.update(buf)
|
||||||
|
buf = binary_stream.read(block_size)
|
||||||
|
|
||||||
|
if isinstance(path_or_stream, Path):
|
||||||
|
with path_or_stream.open("rb") as afile:
|
||||||
|
_hash_buf(afile)
|
||||||
|
elif isinstance(path_or_stream, BytesIO):
|
||||||
|
_hash_buf(path_or_stream)
|
||||||
|
|
||||||
|
return hasher.hexdigest()
|
||||||
|
|
||||||
|
|
||||||
|
def create_hash(string: str):
|
||||||
|
hasher = hashlib.sha256()
|
||||||
|
hasher.update(string.encode("utf-8"))
|
||||||
|
|
||||||
|
return hasher.hexdigest()
|
||||||
73
examples/convert.py
Normal file
73
examples/convert.py
Normal file
@ -0,0 +1,73 @@
|
|||||||
|
import json
|
||||||
|
import logging
|
||||||
|
import time
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import Iterable
|
||||||
|
|
||||||
|
from docling.backend.pypdfium2_backend import PyPdfiumDocumentBackend
|
||||||
|
from docling.datamodel.base_models import ConversionStatus, PipelineOptions
|
||||||
|
from docling.datamodel.document import ConvertedDocument, DocumentConversionInput
|
||||||
|
from docling.document_converter import DocumentConverter
|
||||||
|
|
||||||
|
_log = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
def export_documents(
|
||||||
|
converted_docs: Iterable[ConvertedDocument],
|
||||||
|
output_dir: Path,
|
||||||
|
):
|
||||||
|
output_dir.mkdir(parents=True, exist_ok=True)
|
||||||
|
|
||||||
|
success_count = 0
|
||||||
|
failure_count = 0
|
||||||
|
|
||||||
|
for doc in converted_docs:
|
||||||
|
if doc.status == ConversionStatus.SUCCESS:
|
||||||
|
success_count += 1
|
||||||
|
doc_filename = doc.input.file.stem
|
||||||
|
|
||||||
|
# Export Deep Search document JSON format:
|
||||||
|
with (output_dir / f"{doc_filename}.json").open("w") as fp:
|
||||||
|
fp.write(json.dumps(doc.render_as_dict()))
|
||||||
|
|
||||||
|
# Export Markdown format:
|
||||||
|
with (output_dir / f"{doc_filename}.md").open("w") as fp:
|
||||||
|
fp.write(doc.render_as_markdown())
|
||||||
|
else:
|
||||||
|
_log.info(f"Document {doc.input.file} failed to convert.")
|
||||||
|
failure_count += 1
|
||||||
|
|
||||||
|
_log.info(
|
||||||
|
f"Processed {success_count + failure_count} docs, of which {failure_count} failed"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
logging.basicConfig(level=logging.INFO)
|
||||||
|
|
||||||
|
input_doc_paths = [
|
||||||
|
# Path("/Users/cau/Downloads/Issue-36122.pdf"),
|
||||||
|
# Path("/Users/cau/Downloads/IBM_Storage_Insights_Fact_Sheet.pdf"),
|
||||||
|
Path("./test/data/2206.01062.pdf"),
|
||||||
|
Path("./test/data/2203.01017v2.pdf"),
|
||||||
|
Path("./test/data/2305.03393v1.pdf"),
|
||||||
|
]
|
||||||
|
|
||||||
|
artifacts_path = DocumentConverter.download_models_hf()
|
||||||
|
|
||||||
|
doc_converter = DocumentConverter(artifacts_path=artifacts_path)
|
||||||
|
|
||||||
|
input = DocumentConversionInput.from_paths(input_doc_paths)
|
||||||
|
|
||||||
|
start_time = time.time()
|
||||||
|
|
||||||
|
converted_docs = doc_converter.convert(input)
|
||||||
|
export_documents(converted_docs, output_dir=Path("./scratch"))
|
||||||
|
|
||||||
|
end_time = time.time() - start_time
|
||||||
|
|
||||||
|
_log.info(f"All documents were converted in {end_time:.2f} seconds.")
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
11
examples/minimal.py
Normal file
11
examples/minimal.py
Normal file
@ -0,0 +1,11 @@
|
|||||||
|
from docling.datamodel.document import DocumentConversionInput
|
||||||
|
from docling.document_converter import DocumentConverter
|
||||||
|
|
||||||
|
artifacts_path = DocumentConverter.download_models_hf()
|
||||||
|
doc_converter = DocumentConverter(artifacts_path=artifacts_path)
|
||||||
|
|
||||||
|
input = DocumentConversionInput.from_paths(["factsheet.pdf"])
|
||||||
|
converted_docs = doc_converter.convert(input)
|
||||||
|
|
||||||
|
for d in converted_docs:
|
||||||
|
print(d.render_as_dict())
|
||||||
4865
poetry.lock
generated
Normal file
4865
poetry.lock
generated
Normal file
File diff suppressed because it is too large
Load Diff
72
pyproject.toml
Normal file
72
pyproject.toml
Normal file
@ -0,0 +1,72 @@
|
|||||||
|
[tool.poetry]
|
||||||
|
name = "docling"
|
||||||
|
version = "0.1.0"
|
||||||
|
description = "Docling PDF conversion package"
|
||||||
|
authors = ["Christoph Auer <cau@zurich.ibm.com>", "Michele Dolfi <dol@zurich.ibm.com>", "Maxim Lysak <mly@zurich.ibm.com>", "Nikos Livathinos <nli@zurich.ibm.com>", "Ahmed Nassar <ahn@zurich.ibm.com>", "Peter Staar <taa@zurich.ibm.com>"]
|
||||||
|
license = "MIT"
|
||||||
|
readme = "README.md"
|
||||||
|
keywords= ["docling", "convert", "document", "pdf", "layout model", "segmentation", "table structure", "table former"]
|
||||||
|
classifiers = [
|
||||||
|
"License :: OSI Approved :: MIT License",
|
||||||
|
"Operating System :: MacOS :: MacOS X",
|
||||||
|
"Operating System :: POSIX :: Linux",
|
||||||
|
"Development Status :: 5 - Production/Stable",
|
||||||
|
"Intended Audience :: Developers",
|
||||||
|
"Intended Audience :: Science/Research",
|
||||||
|
"Topic :: Scientific/Engineering :: Artificial Intelligence",
|
||||||
|
"Programming Language :: Python :: 3"
|
||||||
|
]
|
||||||
|
packages = [{include = "docling"}]
|
||||||
|
|
||||||
|
[tool.poetry.dependencies]
|
||||||
|
python = "^3.11"
|
||||||
|
pydantic = "^2.0.0"
|
||||||
|
docling-core = "^0.2.0"
|
||||||
|
docling-ibm-models = "^0.2.0"
|
||||||
|
deepsearch-glm = ">=0.18.4,<1"
|
||||||
|
deepsearch-toolkit = ">=0.47.0,<1"
|
||||||
|
filetype = "^1.2.0"
|
||||||
|
pypdfium2 = "^4.30.0"
|
||||||
|
pydantic-settings = "^2.3.0"
|
||||||
|
huggingface_hub = ">=0.23,<1"
|
||||||
|
|
||||||
|
[tool.poetry.group.ocr.dependencies]
|
||||||
|
easyocr = "^1.7"
|
||||||
|
|
||||||
|
[tool.poetry.group.dev.dependencies]
|
||||||
|
black = {extras = ["jupyter"], version = "^24.4.2"}
|
||||||
|
pytest = "^7.2.2"
|
||||||
|
pre-commit = "^3.7.1"
|
||||||
|
mypy = "^1.10.1"
|
||||||
|
isort = "^5.10.1"
|
||||||
|
python-semantic-release = "^7.32.2"
|
||||||
|
flake8 = "^6.0.0"
|
||||||
|
pyproject-flake8 = "^6.0.0"
|
||||||
|
pytest-xdist = "^3.3.1"
|
||||||
|
types-requests = "^2.31.0.2"
|
||||||
|
flake8-pyproject = "^1.2.3"
|
||||||
|
pylint = "^2.17.5"
|
||||||
|
|
||||||
|
[build-system]
|
||||||
|
requires = ["poetry-core"]
|
||||||
|
build-backend = "poetry.core.masonry.api"
|
||||||
|
|
||||||
|
[tool.black]
|
||||||
|
line-length = 88
|
||||||
|
target-version = ["py311"]
|
||||||
|
include = '\.pyi?$'
|
||||||
|
|
||||||
|
[tool.isort]
|
||||||
|
profile = "black"
|
||||||
|
line_length = 88
|
||||||
|
py_version=311
|
||||||
|
|
||||||
|
[tool.mypy]
|
||||||
|
pretty = true
|
||||||
|
# strict = true
|
||||||
|
no_implicit_optional = true
|
||||||
|
python_version = "3.11"
|
||||||
|
|
||||||
|
[tool.flake8]
|
||||||
|
max-line-length = 88
|
||||||
|
extend-ignore = ["E203", "E501"]
|
||||||
BIN
test/data/2203.01017v2.pdf
Normal file
BIN
test/data/2203.01017v2.pdf
Normal file
Binary file not shown.
BIN
test/data/2206.01062.pdf
Normal file
BIN
test/data/2206.01062.pdf
Normal file
Binary file not shown.
BIN
test/data/2305.03393v1.pdf
Normal file
BIN
test/data/2305.03393v1.pdf
Normal file
Binary file not shown.
33
test/test_backend_pdfium.py
Normal file
33
test/test_backend_pdfium.py
Normal file
@ -0,0 +1,33 @@
|
|||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from docling.backend.pypdfium2_backend import PyPdfiumDocumentBackend, PyPdfiumPageBackend
|
||||||
|
from docling.datamodel.base_models import BoundingBox
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def test_doc_path():
|
||||||
|
return Path("./data/2206.01062.pdf")
|
||||||
|
|
||||||
|
def test_get_text_from_rect(test_doc_path):
|
||||||
|
doc_backend = PyPdfiumDocumentBackend(test_doc_path)
|
||||||
|
page_backend: PyPdfiumPageBackend = doc_backend.load_page(0)
|
||||||
|
|
||||||
|
# Get the title text of the DocLayNet paper
|
||||||
|
textpiece = page_backend.get_text_in_rect(bbox=BoundingBox(l=102,t=77,r=511,b=124))
|
||||||
|
ref = "DocLayNet: A Large Human-Annotated Dataset for\r\nDocument-Layout Analysis"
|
||||||
|
|
||||||
|
assert textpiece.strip() == ref
|
||||||
|
|
||||||
|
def test_crop_page_image(test_doc_path):
|
||||||
|
doc_backend = PyPdfiumDocumentBackend(test_doc_path)
|
||||||
|
page_backend: PyPdfiumPageBackend = doc_backend.load_page(0)
|
||||||
|
|
||||||
|
# Crop out "Figure 1" from the DocLayNet paper
|
||||||
|
im = page_backend.get_page_image(scale=2, cropbox=BoundingBox(l=317,t=246,r=574,b=527))
|
||||||
|
# im.show()
|
||||||
|
|
||||||
|
def test_num_pages(test_doc_path):
|
||||||
|
doc_backend = PyPdfiumDocumentBackend(test_doc_path)
|
||||||
|
doc_backend.page_count() == 9
|
||||||
Loading…
Reference in New Issue
Block a user