[
{
"path": ".conda/meta.yaml",
"content": "{% set pyproject = load_file_data('../pyproject.toml', from_recipe_dir=True) %}\n{% set project = pyproject.get('project') %}\n{% set urls = pyproject.get('project', {}).get('urls') %}\n{% set version = environ.get('BUILD_VERSION', '0.8.2a0') %}\n\npackage:\n name: onnxtr\n version: {{ version }}\n\nsource:\n fn: onnxtr-{{ version }}.tar.gz\n url: ../dist/onnxtr-{{ version }}.tar.gz\n\nbuild:\n script: python setup.py install --single-version-externally-managed --record=record.txt\n\nrequirements:\n host:\n - python>=3.10, <3.12\n - setuptools\n\n run:\n - numpy >=1.16.0, <3.0.0\n - scipy >=1.4.0, <2.0.0\n - pillow >=9.2.0\n - opencv >=4.5.0, <5.0.0\n - pypdfium2-team::pypdfium2_helpers >=4.11.0, <5.0.0\n - pyclipper >=1.2.0, <2.0.0\n - langdetect >=1.0.9, <2.0.0\n - rapidfuzz >=3.0.0, <4.0.0\n - huggingface_hub >=0.20.0, <1.0.0\n - defusedxml >=0.7.0\n - anyascii >=0.3.2\n - tqdm >=4.30.0\n\ntest:\n requires:\n - pip\n - onnxruntime\n\n imports:\n - onnxtr\n\nabout:\n home: {{ urls.get('repository') }}\n license: Apache-2.0\n license_file: {{ project.get('license', {}).get('file') }}\n summary: {{ project.get('description') | replace(\":\", \" -\")}}\n dev_url: {{ urls.get('repository') }}\n"
},
{
"path": ".github/CODEOWNERS",
"content": "* @felixdittrich92"
},
{
"path": ".github/FUNDING.yml",
"content": "# These are supported funding model platforms\n\ngithub: felixdittrich92\npatreon: # Replace with a single Patreon username\nopen_collective: # Replace with a single Open Collective username\nko_fi: # Replace with a single Ko-fi username\ntidelift: # Replace with a single Tidelift platform-name/package-name e.g., npm/babel\ncommunity_bridge: # Replace with a single Community Bridge project-name e.g., cloud-foundry\nliberapay: # Replace with a single Liberapay username\nissuehunt: # Replace with a single IssueHunt username\nlfx_crowdfunding: # Replace with a single LFX Crowdfunding project-name e.g., cloud-foundry\npolar: # Replace with a single Polar username\nbuy_me_a_coffee: # Replace with a single Buy Me a Coffee username\nthanks_dev: # Replace with a single thanks.dev username\ncustom: # Replace with up to 4 custom sponsorship URLs e.g., ['link1', 'link2']\n"
},
{
"path": ".github/ISSUE_TEMPLATE/bug_report.yml",
"content": "name: 🐛 Bug report\ndescription: Create a report to help us improve the library\nlabels: 'type: bug'\n\nbody:\n- type: markdown\n attributes:\n value: >\n #### Before reporting a bug, please check that the issue hasn't already been addressed in [the existing and past issues](https://github.com/felixdittrich92/onnxtr/issues).\n- type: textarea\n attributes:\n label: Bug description\n description: |\n A clear and concise description of what the bug is.\n\n Please explain the result you observed and the behavior you were expecting.\n placeholder: |\n A clear and concise description of what the bug is.\n validations:\n required: true\n\n- type: textarea\n attributes:\n label: Code snippet to reproduce the bug\n description: |\n Sample code to reproduce the problem.\n\n Please wrap your code snippet with ```` ```triple quotes blocks``` ```` for readability.\n placeholder: |\n ```python\n Sample code to reproduce the problem\n ```\n validations:\n required: true\n- type: textarea\n attributes:\n label: Error traceback\n description: |\n The error message you received running the code snippet, with the full traceback.\n\n Please wrap your error message with ```` ```triple quotes blocks``` ```` for readability.\n placeholder: |\n ```\n The error message you got, with the full traceback.\n ```\n validations:\n required: true\n- type: textarea\n attributes:\n label: Environment\n description: |\n Please describe your environment:\n OS:\n Python version:\n Library version:\n Onnxruntime version:\n validations:\n required: true\n- type: markdown\n attributes:\n value: >\n Thanks for helping us improve the library!\n"
},
{
"path": ".github/ISSUE_TEMPLATE/config.yml",
"content": "blank_issues_enabled: true\ncontact_links:\n - name: Usage questions\n url: https://github.com/felixdittrich92/OnnxTR/discussions\n about: Ask questions and discuss with other OnnxTR community members"
},
{
"path": ".github/ISSUE_TEMPLATE/feature_request.yml",
"content": "name: 🚀 Feature request\ndescription: >\n Submit a proposal/request for a new feature for OnnxTR. Please search for existing issues before creating a new one.\n For non-onnx related features please use the [main repository](https://github.com/mindee/doctr/issues).\nlabels: 'type: enhancement'\n\nbody:\n- type: textarea\n attributes:\n label: 🚀 The feature\n description: >\n A clear and concise description of the feature proposal\n validations:\n required: true\n- type: textarea\n attributes:\n label: Additional context\n description: >\n Add any other context or screenshots about the feature request.\n- type: markdown\n attributes:\n value: >\n Thanks for contributing 🎉\n"
},
{
"path": ".github/dependabot.yml",
"content": "version: 2\nupdates:\n - package-ecosystem: \"pip\"\n directory: \"/\"\n open-pull-requests-limit: 10\n target-branch: \"main\"\n labels: [\"topic: build\"]\n schedule:\n interval: weekly\n day: sunday\n - package-ecosystem: \"github-actions\"\n directory: \"/\"\n open-pull-requests-limit: 10\n target-branch: \"main\"\n labels: [\"topic: CI/CD\"]\n schedule:\n interval: weekly\n day: sunday\n groups:\n github-actions:\n patterns:\n - \"*\"\n"
},
{
"path": ".github/release.yml",
"content": "changelog:\n exclude:\n labels:\n - ignore-for-release\n categories:\n - title: Breaking Changes 🛠\n labels:\n - \"type: breaking change\"\n # NEW FEATURES\n - title: New Features\n labels:\n - \"type: new feature\"\n # BUG FIXES\n - title: Bug Fixes\n labels:\n - \"type: bug\"\n # IMPROVEMENTS\n - title: Improvements\n labels:\n - \"type: enhancement\"\n # MISC\n - title: Miscellaneous\n labels:\n - \"type: misc\"\n"
},
{
"path": ".github/workflows/builds.yml",
"content": "name: builds\n\non:\n push:\n branches: main\n pull_request:\n branches: main\n schedule:\n # Runs every 2 weeks on Monday at 03:00 UTC\n - cron: '0 3 * * 1'\n\njobs:\n build:\n runs-on: ${{ matrix.os }}\n strategy:\n fail-fast: false\n matrix:\n os: [ubuntu-latest, macos-latest, windows-latest]\n python: [\"3.10\", \"3.11\", \"3.12\", \"3.13\"]\n steps:\n - uses: actions/checkout@v6\n - name: Set up Python\n uses: actions/setup-python@v6\n with:\n # MacOS issue ref.: https://github.com/actions/setup-python/issues/855 & https://github.com/actions/setup-python/issues/865\n python-version: ${{ matrix.os == 'macos-latest' && matrix.python == '3.10' && '3.11' || matrix.python }}\n architecture: x64\n - name: Cache python modules\n uses: actions/cache@v5\n with:\n path: ~/.cache/pip\n key: ${{ runner.os }}-pkg-deps-${{ matrix.python }}-${{ hashFiles('pyproject.toml') }}\n - name: Install package\n run: |\n python -m pip install --upgrade pip\n pip install -e .[cpu-headless,viz] --upgrade\n - name: Import package\n run: python -c \"import onnxtr; print(onnxtr.__version__)\"\n\n conda:\n runs-on: ubuntu-latest\n steps:\n - uses: actions/checkout@v6\n - uses: conda-incubator/setup-miniconda@v4\n with:\n auto-update-conda: true\n python-version: \"3.10\"\n channels: pypdfium2-team,bblanchon,defaults,conda-forge\n channel-priority: strict\n - name: Install dependencies\n shell: bash -el {0}\n run: conda install -y conda-build conda-verify anaconda-client\n - name: Install libEGL\n run: sudo apt-get update && sudo apt-get install -y libegl1\n - name: Build and verify\n shell: bash -el {0}\n run: |\n python setup.py sdist\n mkdir conda-dist\n conda build .conda/ --output-folder conda-dist\n conda-verify conda-dist/linux-64/*conda --ignore=C1115\n"
},
{
"path": ".github/workflows/clear_caches.yml",
"content": "name: Clear GitHub runner caches\n\non:\n workflow_dispatch:\n schedule:\n - cron: '0 0 * * *' # Runs once a day\n\njobs:\n clear:\n name: Clear caches\n runs-on: ubuntu-latest\n steps:\n - uses: MyAlbum/purge-cache@v2\n with:\n max-age: 172800 # Caches older than 2 days are deleted\n"
},
{
"path": ".github/workflows/demo.yml",
"content": "name: Sync Hugging Face demo\n\non:\n # Run 'test-demo' on every pull request to the main branch\n pull_request:\n branches: [main]\n\n # Run 'sync-to-hub' on push when tagging (e.g., 'v*') and on a scheduled cron job\n push:\n tags:\n - 'v*'\n\n schedule:\n - cron: '0 2 10 * *' # At 02:00 on day-of-month 10 (every month)\n\n # Allow manual triggering of the workflow\n workflow_dispatch:\n\njobs:\n # This job runs on every pull request to main\n test-demo:\n runs-on: ${{ matrix.os }}\n strategy:\n fail-fast: false\n matrix:\n os: [ubuntu-latest]\n python: [\"3.10\"]\n steps:\n - uses: actions/checkout@v6\n - name: Set up Python\n uses: actions/setup-python@v6\n with:\n python-version: ${{ matrix.python }}\n architecture: x64\n - name: Cache python modules\n uses: actions/cache@v5\n with:\n path: ~/.cache/pip\n key: ${{ runner.os }}-pkg-deps-${{ matrix.python }}-${{ hashFiles('requirements.txt') }}-${{ hashFiles('demo/requirements.txt') }}\n - name: Install dependencies\n run: |\n python -m pip install --upgrade pip\n pip install -r demo/requirements.txt --upgrade\n - name: Start Gradio demo\n run: |\n nohup python demo/app.py &\n sleep 10 # Allow some time for the Gradio server to start\n - name: Check demo build\n run: |\n curl --fail http://127.0.0.1:7860/ || exit 1\n\n # This job only runs when a new version tag is pushed or during the cron job\n sync-to-hub:\n if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags/v') || github.event_name == 'schedule' || github.event_name == 'workflow_dispatch'\n needs: test-demo\n runs-on: ubuntu-latest\n steps:\n - uses: actions/checkout@v6\n with:\n fetch-depth: 0\n - name: Set up Python\n uses: actions/setup-python@v6\n with:\n python-version: \"3.10\"\n - name: Install huggingface_hub\n run: pip install huggingface-hub\n - name: Upload folder to Hugging Face\n env:\n HF_TOKEN: ${{ secrets.HF_TOKEN }}\n run: |\n python -c \"\n from huggingface_hub import HfApi\n api = HfApi(token='${{ secrets.HF_TOKEN }}')\n repo_id = 'Felix92/OnnxTR-OCR'\n api.upload_folder(repo_id=repo_id, repo_type='space', folder_path='demo/')\n api.restart_space(repo_id=repo_id, factory_reboot=True)\n \"\n"
},
{
"path": ".github/workflows/docker.yml",
"content": "# https://docs.github.com/en/actions/publishing-packages/publishing-docker-images#publishing-images-to-github-packages\n#\nname: Docker image on ghcr.io\n\non:\n push:\n tags:\n - 'v*'\n pull_request:\n branches: main\n schedule:\n - cron: '0 2 1 6 *' # At 02:00 on day-of-month 1 in June (once a year actually)\n\nenv:\n REGISTRY: ghcr.io\n\njobs:\n build-and-push-image:\n runs-on: ubuntu-latest\n\n strategy:\n fail-fast: false\n matrix:\n image:\n - \"ubuntu:24.04\" # Base image for CPU variants\n - \"nvidia/cuda:12.6.2-base-ubuntu24.04\" # Base image for GPU\n variant:\n - \"cpu-headless\" # CPU variant 1\n - \"openvino-headless\" # CPU variant 2\n - \"gpu-headless\" # GPU variant\n python: [3.10.13]\n\n # Exclude invalid combinations\n exclude:\n - image: \"nvidia/cuda:12.6.2-base-ubuntu24.04\"\n variant: \"cpu-headless\"\n - image: \"nvidia/cuda:12.6.2-base-ubuntu24.04\"\n variant: \"openvino-headless\"\n - image: \"ubuntu:24.04\"\n variant: \"gpu-headless\"\n\n permissions:\n contents: read\n packages: write\n\n steps:\n - name: Checkout repository\n uses: actions/checkout@v6\n\n - name: Log in to the Container registry\n uses: docker/login-action@v4\n with:\n registry: ${{ env.REGISTRY }}\n username: ${{ github.actor }}\n password: ${{ secrets.GITHUB_TOKEN }}\n\n - name: Sanitize docker tag\n run: |\n # Start with the base prefix\n PREFIX_DOCKER_TAG=\"OnnxTR-${{ matrix.variant }}-py${{ matrix.python }}\"\n\n # Replace any commas with hyphens (if needed)\n PREFIX_DOCKER_TAG=$(echo \"$PREFIX_DOCKER_TAG\" | sed 's/,/-/g')\n\n # Determine suffix based on image\n IMAGE=\"${{ matrix.image }}\"\n case \"$IMAGE\" in\n \"nvidia/cuda:\"*)\n SUFFIX=$(echo \"$IMAGE\" | sed -E 's|.*/cuda:([0-9]+\\.[0-9]+\\.[0-9]+)-base-(ubuntu[0-9]+\\.[0-9]+)|-\\2-cuda\\1|')\n ;;\n \"ubuntu:\"*)\n SUFFIX=$(echo \"$IMAGE\" | sed -E 's|ubuntu:([0-9]+\\.[0-9]+)|-ubuntu\\1|')\n ;;\n *)\n SUFFIX=\"\"\n ;;\n esac\n\n # Combine the prefix, suffix, and ensure ending hyphen\n PREFIX_DOCKER_TAG=\"${PREFIX_DOCKER_TAG}${SUFFIX}-\"\n\n # Export to environment\n echo \"PREFIX_DOCKER_TAG=${PREFIX_DOCKER_TAG}\" >> $GITHUB_ENV\n\n # Debugging output\n echo \"Final Docker Tag: $PREFIX_DOCKER_TAG\"\n\n - name: Extract metadata (tags, labels) for Docker\n id: meta\n uses: docker/metadata-action@v6\n with:\n images: ${{ env.REGISTRY }}/${{ github.repository }}\n tags: |\n # used only on schedule event\n type=schedule,pattern={{date 'YYYY-MM'}},prefix=${{ env.PREFIX_DOCKER_TAG }}\n # used only if a tag following semver is published\n type=semver,pattern={{raw}},prefix=${{ env.PREFIX_DOCKER_TAG }}\n\n - name: Build Docker image\n id: build\n uses: docker/build-push-action@v7\n with:\n context: .\n build-args: |\n BASE_IMAGE=${{ matrix.image }}\n SYSTEM=${{ matrix.variant }}\n PYTHON_VERSION=${{ matrix.python }}\n ONNXTR_REPO=${{ github.repository }}\n ONNXTR_VERSION=${{ github.sha }}\n push: false # push only if `import onnxtr` works\n tags: ${{ steps.meta.outputs.tags }}\n\n - name: Check if `import onnxtr` works\n run: docker run ${{ steps.build.outputs.imageid }} python3 -c 'import onnxtr; print(onnxtr.__version__)'\n\n - name: Push Docker image\n if: ${{ (github.ref == 'refs/heads/main' && github.event_name != 'pull_request') || (startsWith(github.ref, 'refs/tags') && github.event_name == 'push') }}\n uses: docker/build-push-action@v7\n with:\n context: .\n build-args: |\n BASE_IMAGE=${{ matrix.image }}\n SYSTEM=${{ matrix.variant }}\n PYTHON_VERSION=${{ matrix.python }}\n ONNXTR_REPO=${{ github.repository }}\n ONNXTR_VERSION=${{ github.sha }}\n push: true\n tags: ${{ steps.meta.outputs.tags }}"
},
{
"path": ".github/workflows/main.yml",
"content": "name: tests\n\non:\n push:\n branches: main\n pull_request:\n branches: main\n schedule:\n # Runs every 2 weeks on Monday at 03:00 UTC\n - cron: '0 3 * * 1'\n\njobs:\n pytest-common:\n runs-on: ${{ matrix.os }}\n strategy:\n matrix:\n os: [ubuntu-latest]\n python: [\"3.10\", \"3.11\", \"3.12\"]\n backend: [\"cpu-headless\", \"openvino-headless\"]\n steps:\n - uses: actions/checkout@v6\n - name: Set up Python\n uses: actions/setup-python@v6\n with:\n python-version: ${{ matrix.python }}\n architecture: x64\n - name: Cache python modules\n uses: actions/cache@v5\n with:\n path: ~/.cache/pip\n key: ${{ runner.os }}-pkg-deps-${{ matrix.python }}-${{ hashFiles('pyproject.toml') }}-tests\n - name: Install dependencies\n run: |\n python -m pip install --upgrade pip\n pip install -e .[${{ matrix.backend }},viz,html,testing] --upgrade\n - name: Run unittests\n run: |\n coverage run -m pytest tests/common/ -rs --memray\n coverage xml -o coverage-common-${{ matrix.backend }}-${{ matrix.python }}.xml\n - uses: actions/upload-artifact@v7\n with:\n name: coverage-common-${{ matrix.backend }}-${{ matrix.python }}\n path: ./coverage-common-${{ matrix.backend }}-${{ matrix.python }}.xml\n if-no-files-found: error\n\n codecov-upload:\n runs-on: ubuntu-latest\n needs: [ pytest-common ]\n steps:\n - uses: actions/checkout@v6\n - uses: actions/download-artifact@v8\n - name: Upload coverage to Codecov\n uses: codecov/codecov-action@v6\n with:\n flags: unittests\n fail_ci_if_error: true\n token: ${{ secrets.CODECOV_TOKEN }}\n"
},
{
"path": ".github/workflows/publish.yml",
"content": "name: publish\n\non:\n release:\n types: [published]\n\njobs:\n pypi:\n if: \"!github.event.release.prerelease\"\n strategy:\n fail-fast: false\n matrix:\n os: [ubuntu-latest]\n python: [\"3.10\"]\n runs-on: ${{ matrix.os }}\n steps:\n - uses: actions/checkout@v6\n - name: Set up Python\n uses: actions/setup-python@v6\n with:\n python-version: ${{ matrix.python }}\n architecture: x64\n - name: Cache python modules\n uses: actions/cache@v5\n with:\n path: ~/.cache/pip\n key: ${{ runner.os }}-pkg-deps-${{ matrix.python }}-${{ hashFiles('pyproject.toml') }}\n - name: Install dependencies\n run: |\n python -m pip install --upgrade pip\n pip install setuptools wheel twine --upgrade\n - name: Get release tag\n id: release_tag\n run: echo \"VERSION=${GITHUB_REF/refs\\/tags\\//}\" >> $GITHUB_ENV\n - name: Build and publish\n env:\n TWINE_USERNAME: ${{ secrets.PYPI_USERNAME }}\n TWINE_PASSWORD: ${{ secrets.PYPI_PASSWORD }}\n VERSION: ${{ env.VERSION }}\n run: |\n BUILD_VERSION=$VERSION python setup.py sdist bdist_wheel\n twine check dist/*\n twine upload dist/*\n\n pypi-check:\n needs: pypi\n if: \"!github.event.release.prerelease\"\n strategy:\n fail-fast: false\n matrix:\n os: [ubuntu-latest]\n python: [\"3.10\"]\n runs-on: ${{ matrix.os }}\n steps:\n - uses: actions/checkout@v6\n - name: Set up Python\n uses: actions/setup-python@v6\n with:\n python-version: ${{ matrix.python }}\n architecture: x64\n - name: Install package\n run: |\n python -m pip install --upgrade pip\n pip install onnxtr[cpu] --upgrade\n python -c \"from importlib.metadata import version; print(version('onnxtr'))\"\n\n conda:\n if: \"!github.event.release.prerelease\"\n runs-on: ubuntu-latest\n steps:\n - uses: actions/checkout@v6\n - uses: conda-incubator/setup-miniconda@v4\n with:\n auto-update-conda: true\n python-version: \"3.10\"\n channels: pypdfium2-team,bblanchon,defaults,conda-forge\n channel-priority: strict\n - name: Install dependencies\n shell: bash -el {0}\n run: conda install -y conda-build conda-verify anaconda-client\n - name: Install libEGL\n run: sudo apt-get update && sudo apt-get install -y libegl1\n - name: Get release tag\n id: release_tag\n run: echo \"VERSION=${GITHUB_REF/refs\\/tags\\//}\" >> $GITHUB_ENV\n - name: Build and publish\n shell: bash -el {0}\n env:\n ANACONDA_API_TOKEN: ${{ secrets.ANACONDA_TOKEN }}\n VERSION: ${{ env.VERSION }}\n run: |\n echo \"BUILD_VERSION=${VERSION}\" >> $GITHUB_ENV\n python setup.py sdist\n mkdir conda-dist\n conda build .conda/ --output-folder conda-dist\n conda-verify conda-dist/linux-64/*conda --ignore=C1115\n anaconda upload conda-dist/linux-64/*conda\n\n conda-check:\n if: \"!github.event.release.prerelease\"\n runs-on: ubuntu-latest\n needs: conda\n steps:\n - uses: conda-incubator/setup-miniconda@v4\n with:\n auto-update-conda: true\n python-version: \"3.10\"\n - name: Install package\n shell: bash -el {0}\n run: |\n conda config --set channel_priority strict\n conda install -c conda-forge onnxruntime\n conda install -c felix92 -c pypdfium2-team -c bblanchon -c defaults -c conda-forge onnxtr\n python -c \"from importlib.metadata import version; print(version('onnxtr'))\"\n"
},
{
"path": ".github/workflows/style.yml",
"content": "name: style\n\non:\n push:\n branches: main\n pull_request:\n branches: main\n\njobs:\n ruff:\n runs-on: ${{ matrix.os }}\n strategy:\n matrix:\n os: [ubuntu-latest]\n python: [\"3.10\"]\n steps:\n - uses: actions/checkout@v6\n - name: Set up Python\n uses: actions/setup-python@v6\n with:\n python-version: ${{ matrix.python }}\n architecture: x64\n - name: Run ruff\n run: |\n pip install ruff --upgrade\n ruff --version\n ruff check --diff .\n\n mypy:\n runs-on: ${{ matrix.os }}\n strategy:\n matrix:\n os: [ubuntu-latest]\n python: [\"3.10\"]\n steps:\n - uses: actions/checkout@v6\n - name: Set up Python\n uses: actions/setup-python@v6\n with:\n python-version: ${{ matrix.python }}\n architecture: x64\n - name: Cache python modules\n uses: actions/cache@v5\n with:\n path: ~/.cache/pip\n key: ${{ runner.os }}-pkg-deps-${{ matrix.python }}-${{ hashFiles('pyproject.toml') }}\n - name: Install dependencies\n run: |\n python -m pip install --upgrade pip\n pip install -e .[dev] --upgrade\n pip install mypy --upgrade\n - name: Run mypy\n run: |\n mypy --version\n mypy\n"
},
{
"path": ".gitignore",
"content": "# Byte-compiled / optimized / DLL files\n__pycache__/\n*.py[cod]\n*$py.class\n\n# C extensions\n*.so\n\n# Distribution / packaging\n.Python\nbuild/\ndevelop-eggs/\ndist/\ndownloads/\neggs/\n.eggs/\nlib/\nlib64/\nparts/\nsdist/\nvar/\nwheels/\npip-wheel-metadata/\nshare/python-wheels/\n*.egg-info/\n.installed.cfg\n*.egg\nMANIFEST\n\n# PyInstaller\n# Usually these files are written by a python script from a template\n# before PyInstaller builds the exe, so as to inject date/other infos into it.\n*.manifest\n*.spec\n\n# Installer logs\npip-log.txt\npip-delete-this-directory.txt\n\n# Unit test / coverage reports\nhtmlcov/\n.tox/\n.nox/\n.coverage\n.coverage.*\n.cache\nnosetests.xml\ncoverage.xml\n*.cover\n*.py,cover\n.hypothesis/\n.pytest_cache/\n\n# Translations\n*.mo\n*.pot\n\n# Django stuff:\n*.log\nlocal_settings.py\ndb.sqlite3\ndb.sqlite3-journal\n\n# Flask stuff:\ninstance/\n.webassets-cache\n\n# Scrapy stuff:\n.scrapy\n\n# Sphinx documentation\ndocs/_build/\n\n# PyBuilder\ntarget/\n\n# Jupyter Notebook\n.ipynb_checkpoints\n\n# IPython\nprofile_default/\nipython_config.py\n\n# pyenv\n.python-version\n\n# pipenv\n# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.\n# However, in case of collaboration, if having platform-specific dependencies or dependencies\n# having no cross-platform support, pipenv may install dependencies that don't work, or not\n# install all needed dependencies.\n#Pipfile.lock\n\n# PEP 582; used by e.g. github.com/David-OConnor/pyflow\n__pypackages__/\n\n# Celery stuff\ncelerybeat-schedule\ncelerybeat.pid\n\n# SageMath parsed files\n*.sage.py\n\n# Environments\n.env\n.venv\nenv/\nvenv/\nENV/\nenv.bak/\nvenv.bak/\n\n# Spyder project settings\n.spyderproject\n.spyproject\n\n# Rope project settings\n.ropeproject\n\n# mkdocs documentation\n/site\n\n# mypy\n.mypy_cache/\n.dmypy.json\ndmypy.json\n\n# Pyre type checker\n.pyre/\n\n# Temp files\nonnxtr/version.py\nlogs/\nwandb/\n.idea/\n\n# Model files\n*.onnx\n.qodo\n\n# Profile files\nyappi_profile.stats\nmemray_profile.bin\nmemray_flamegraph.html\n"
},
{
"path": ".pre-commit-config.yaml",
"content": "repos:\n - repo: https://github.com/pre-commit/pre-commit-hooks\n rev: v6.0.0\n hooks:\n - id: check-ast\n - id: check-yaml\n exclude: .conda\n - id: check-toml\n - id: check-json\n - id: check-added-large-files\n exclude: docs/images/\n - id: end-of-file-fixer\n - id: trailing-whitespace\n - id: debug-statements\n - id: check-merge-conflict\n - id: no-commit-to-branch\n args: ['--branch', 'main']\n - repo: https://github.com/astral-sh/ruff-pre-commit\n rev: v0.15.0\n hooks:\n - id: ruff\n args: [ --fix ]\n - id: ruff-format\n"
},
{
"path": "CODE_OF_CONDUCT.md",
"content": "# Contributor Covenant Code of Conduct\n\n## Our Pledge\n\nWe as members, contributors, and leaders pledge to make participation in our\ncommunity a harassment-free experience for everyone, regardless of age, body\nsize, visible or invisible disability, ethnicity, sex characteristics, gender\nidentity and expression, level of experience, education, socio-economic status,\nnationality, personal appearance, race, religion, or sexual identity\nand orientation.\n\nWe pledge to act and interact in ways that contribute to an open, welcoming,\ndiverse, inclusive, and healthy community.\n\n## Our Standards\n\nExamples of behavior that contributes to a positive environment for our\ncommunity include:\n\n* Demonstrating empathy and kindness toward other people\n* Being respectful of differing opinions, viewpoints, and experiences\n* Giving and gracefully accepting constructive feedback\n* Accepting responsibility and apologizing to those affected by our mistakes,\n and learning from the experience\n* Focusing on what is best not just for us as individuals, but for the\n overall community\n\nExamples of unacceptable behavior include:\n\n* The use of sexualized language or imagery, and sexual attention or\n advances of any kind\n* Trolling, insulting or derogatory comments, and personal or political attacks\n* Public or private harassment\n* Publishing others' private information, such as a physical or email\n address, without their explicit permission\n* Other conduct which could reasonably be considered inappropriate in a\n professional setting\n\n## Enforcement Responsibilities\n\nCommunity leaders are responsible for clarifying and enforcing our standards of\nacceptable behavior and will take appropriate and fair corrective action in\nresponse to any behavior that they deem inappropriate, threatening, offensive,\nor harmful.\n\nCommunity leaders have the right and responsibility to remove, edit, or reject\ncomments, commits, code, wiki edits, issues, and other contributions that are\nnot aligned to this Code of Conduct, and will communicate reasons for moderation\ndecisions when appropriate.\n\n## Scope\n\nThis Code of Conduct applies within all community spaces, and also applies when\nan individual is officially representing the community in public spaces.\nExamples of representing our community include using an official e-mail address,\nposting via an official social media account, or acting as an appointed\nrepresentative at an online or offline event.\n\n## Enforcement\n\nInstances of abusive, harassing, or otherwise unacceptable behavior may be\nreported to the community leaders responsible for enforcement at\ncontact@mindee.com.\nAll complaints will be reviewed and investigated promptly and fairly.\n\nAll community leaders are obligated to respect the privacy and security of the\nreporter of any incident.\n\n## Enforcement Guidelines\n\nCommunity leaders will follow these Community Impact Guidelines in determining\nthe consequences for any action they deem in violation of this Code of Conduct:\n\n### 1. Correction\n\n**Community Impact**: Use of inappropriate language or other behavior deemed\nunprofessional or unwelcome in the community.\n\n**Consequence**: A private, written warning from community leaders, providing\nclarity around the nature of the violation and an explanation of why the\nbehavior was inappropriate. A public apology may be requested.\n\n### 2. Warning\n\n**Community Impact**: A violation through a single incident or series\nof actions.\n\n**Consequence**: A warning with consequences for continued behavior. No\ninteraction with the people involved, including unsolicited interaction with\nthose enforcing the Code of Conduct, for a specified period of time. This\nincludes avoiding interactions in community spaces as well as external channels\nlike social media. Violating these terms may lead to a temporary or\npermanent ban.\n\n### 3. Temporary Ban\n\n**Community Impact**: A serious violation of community standards, including\nsustained inappropriate behavior.\n\n**Consequence**: A temporary ban from any sort of interaction or public\ncommunication with the community for a specified period of time. No public or\nprivate interaction with the people involved, including unsolicited interaction\nwith those enforcing the Code of Conduct, is allowed during this period.\nViolating these terms may lead to a permanent ban.\n\n### 4. Permanent Ban\n\n**Community Impact**: Demonstrating a pattern of violation of community\nstandards, including sustained inappropriate behavior, harassment of an\nindividual, or aggression toward or disparagement of classes of individuals.\n\n**Consequence**: A permanent ban from any sort of public interaction within\nthe community.\n\n## Attribution\n\nThis Code of Conduct is adapted from the [Contributor Covenant][homepage],\nversion 2.0, available at\nhttps://www.contributor-covenant.org/version/2/0/code_of_conduct.html.\n\nCommunity Impact Guidelines were inspired by [Mozilla's code of conduct\nenforcement ladder](https://github.com/mozilla/diversity).\n\n[homepage]: https://www.contributor-covenant.org\n\nFor answers to common questions about this code of conduct, see the FAQ at\nhttps://www.contributor-covenant.org/faq. Translations are available at\nhttps://www.contributor-covenant.org/translations.\n"
},
{
"path": "Dockerfile",
"content": "ARG BASE_IMAGE\n\nFROM ${BASE_IMAGE}\n\nENV DEBIAN_FRONTEND=noninteractive\nENV LANG=C.UTF-8\nENV PYTHONUNBUFFERED=1\nENV PYTHONDONTWRITEBYTECODE=1\n\nARG SYSTEM\nARG PYTHON_VERSION\n\nRUN apt-get update && apt-get install -y --no-install-recommends \\\n # - Other packages\n build-essential \\\n pkg-config \\\n curl \\\n wget \\\n software-properties-common \\\n unzip \\\n git \\\n # - Packages to build Python\n tar make gcc zlib1g-dev libffi-dev libssl-dev liblzma-dev libbz2-dev libsqlite3-dev \\\n # - Packages for docTR\n libgl1-mesa-dev libsm6 libxext6 libxrender-dev libpangocairo-1.0-0 \\\n && apt-get clean \\\n && rm -rf /var/lib/apt/lists/* \\\nfi\n\n# Install Python\n\nRUN wget http://www.python.org/ftp/python/$PYTHON_VERSION/Python-$PYTHON_VERSION.tgz && \\\n tar -zxf Python-$PYTHON_VERSION.tgz && \\\n cd Python-$PYTHON_VERSION && \\\n mkdir /opt/python/ && \\\n ./configure --prefix=/opt/python && \\\n make && \\\n make install && \\\n cd .. && \\\n rm Python-$PYTHON_VERSION.tgz && \\\n rm -r Python-$PYTHON_VERSION\n\nENV PATH=/opt/python/bin:$PATH\n\n# Install OnnxTR\nARG ONNXTR_REPO='felixdittrich92/onnxtr'\nARG ONNXTR_VERSION=main\nRUN pip3 install -U pip setuptools wheel && \\\n pip3 install \"onnxtr[$SYSTEM,html]@git+https://github.com/$ONNXTR_REPO.git@$ONNXTR_VERSION\"\n"
},
{
"path": "LICENSE",
"content": " Apache License\n Version 2.0, January 2004\n http://www.apache.org/licenses/\n\n TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION\n\n 1. Definitions.\n\n \"License\" shall mean the terms and conditions for use, reproduction,\n and distribution as defined by Sections 1 through 9 of this document.\n\n \"Licensor\" shall mean the copyright owner or entity authorized by\n the copyright owner that is granting the License.\n\n \"Legal Entity\" shall mean the union of the acting entity and all\n other entities that control, are controlled by, or are under common\n control with that entity. For the purposes of this definition,\n \"control\" means (i) the power, direct or indirect, to cause the\n direction or management of such entity, whether by contract or\n otherwise, or (ii) ownership of fifty percent (50%) or more of the\n outstanding shares, or (iii) beneficial ownership of such entity.\n\n \"You\" (or \"Your\") shall mean an individual or Legal Entity\n exercising permissions granted by this License.\n\n \"Source\" form shall mean the preferred form for making modifications,\n including but not limited to software source code, documentation\n source, and configuration files.\n\n \"Object\" form shall mean any form resulting from mechanical\n transformation or translation of a Source form, including but\n not limited to compiled object code, generated documentation,\n and conversions to other media types.\n\n \"Work\" shall mean the work of authorship, whether in Source or\n Object form, made available under the License, as indicated by a\n copyright notice that is included in or attached to the work\n (an example is provided in the Appendix below).\n\n \"Derivative Works\" shall mean any work, whether in Source or Object\n form, that is based on (or derived from) the Work and for which the\n editorial revisions, annotations, elaborations, or other modifications\n represent, as a whole, an original work of authorship. For the purposes\n of this License, Derivative Works shall not include works that remain\n separable from, or merely link (or bind by name) to the interfaces of,\n the Work and Derivative Works thereof.\n\n \"Contribution\" shall mean any work of authorship, including\n the original version of the Work and any modifications or additions\n to that Work or Derivative Works thereof, that is intentionally\n submitted to Licensor for inclusion in the Work by the copyright owner\n or by an individual or Legal Entity authorized to submit on behalf of\n the copyright owner. For the purposes of this definition, \"submitted\"\n means any form of electronic, verbal, or written communication sent\n to the Licensor or its representatives, including but not limited to\n communication on electronic mailing lists, source code control systems,\n and issue tracking systems that are managed by, or on behalf of, the\n Licensor for the purpose of discussing and improving the Work, but\n excluding communication that is conspicuously marked or otherwise\n designated in writing by the copyright owner as \"Not a Contribution.\"\n\n \"Contributor\" shall mean Licensor and any individual or Legal Entity\n on behalf of whom a Contribution has been received by Licensor and\n subsequently incorporated within the Work.\n\n 2. Grant of Copyright License. Subject to the terms and conditions of\n this License, each Contributor hereby grants to You a perpetual,\n worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n copyright license to reproduce, prepare Derivative Works of,\n publicly display, publicly perform, sublicense, and distribute the\n Work and such Derivative Works in Source or Object form.\n\n 3. Grant of Patent License. Subject to the terms and conditions of\n this License, each Contributor hereby grants to You a perpetual,\n worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n (except as stated in this section) patent license to make, have made,\n use, offer to sell, sell, import, and otherwise transfer the Work,\n where such license applies only to those patent claims licensable\n by such Contributor that are necessarily infringed by their\n Contribution(s) alone or by combination of their Contribution(s)\n with the Work to which such Contribution(s) was submitted. If You\n institute patent litigation against any entity (including a\n cross-claim or counterclaim in a lawsuit) alleging that the Work\n or a Contribution incorporated within the Work constitutes direct\n or contributory patent infringement, then any patent licenses\n granted to You under this License for that Work shall terminate\n as of the date such litigation is filed.\n\n 4. Redistribution. You may reproduce and distribute copies of the\n Work or Derivative Works thereof in any medium, with or without\n modifications, and in Source or Object form, provided that You\n meet the following conditions:\n\n (a) You must give any other recipients of the Work or\n Derivative Works a copy of this License; and\n\n (b) You must cause any modified files to carry prominent notices\n stating that You changed the files; and\n\n (c) You must retain, in the Source form of any Derivative Works\n that You distribute, all copyright, patent, trademark, and\n attribution notices from the Source form of the Work,\n excluding those notices that do not pertain to any part of\n the Derivative Works; and\n\n (d) If the Work includes a \"NOTICE\" text file as part of its\n distribution, then any Derivative Works that You distribute must\n include a readable copy of the attribution notices contained\n within such NOTICE file, excluding those notices that do not\n pertain to any part of the Derivative Works, in at least one\n of the following places: within a NOTICE text file distributed\n as part of the Derivative Works; within the Source form or\n documentation, if provided along with the Derivative Works; or,\n within a display generated by the Derivative Works, if and\n wherever such third-party notices normally appear. The contents\n of the NOTICE file are for informational purposes only and\n do not modify the License. You may add Your own attribution\n notices within Derivative Works that You distribute, alongside\n or as an addendum to the NOTICE text from the Work, provided\n that such additional attribution notices cannot be construed\n as modifying the License.\n\n You may add Your own copyright statement to Your modifications and\n may provide additional or different license terms and conditions\n for use, reproduction, or distribution of Your modifications, or\n for any such Derivative Works as a whole, provided Your use,\n reproduction, and distribution of the Work otherwise complies with\n the conditions stated in this License.\n\n 5. Submission of Contributions. Unless You explicitly state otherwise,\n any Contribution intentionally submitted for inclusion in the Work\n by You to the Licensor shall be under the terms and conditions of\n this License, without any additional terms or conditions.\n Notwithstanding the above, nothing herein shall supersede or modify\n the terms of any separate license agreement you may have executed\n with Licensor regarding such Contributions.\n\n 6. Trademarks. This License does not grant permission to use the trade\n names, trademarks, service marks, or product names of the Licensor,\n except as required for reasonable and customary use in describing the\n origin of the Work and reproducing the content of the NOTICE file.\n\n 7. Disclaimer of Warranty. Unless required by applicable law or\n agreed to in writing, Licensor provides the Work (and each\n Contributor provides its Contributions) on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or\n implied, including, without limitation, any warranties or conditions\n of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A\n PARTICULAR PURPOSE. You are solely responsible for determining the\n appropriateness of using or redistributing the Work and assume any\n risks associated with Your exercise of permissions under this License.\n\n 8. Limitation of Liability. In no event and under no legal theory,\n whether in tort (including negligence), contract, or otherwise,\n unless required by applicable law (such as deliberate and grossly\n negligent acts) or agreed to in writing, shall any Contributor be\n liable to You for damages, including any direct, indirect, special,\n incidental, or consequential damages of any character arising as a\n result of this License or out of the use or inability to use the\n Work (including but not limited to damages for loss of goodwill,\n work stoppage, computer failure or malfunction, or any and all\n other commercial damages or losses), even if such Contributor\n has been advised of the possibility of such damages.\n\n 9. Accepting Warranty or Additional Liability. While redistributing\n the Work or Derivative Works thereof, You may choose to offer,\n and charge a fee for, acceptance of support, warranty, indemnity,\n or other liability obligations and/or rights consistent with this\n License. However, in accepting such obligations, You may act only\n on Your own behalf and on Your sole responsibility, not on behalf\n of any other Contributor, and only if You agree to indemnify,\n defend, and hold each Contributor harmless for any liability\n incurred by, or claims asserted against, such Contributor by reason\n of your accepting any such warranty or additional liability.\n\n END OF TERMS AND CONDITIONS\n\n APPENDIX: How to apply the Apache License to your work.\n\n To apply the Apache License to your work, attach the following\n boilerplate notice, with the fields enclosed by brackets \"[]\"\n replaced with your own identifying information. (Don't include\n the brackets!) The text should be enclosed in the appropriate\n comment syntax for the file format. We also recommend that a\n file or class name and description of purpose be included on the\n same \"printed page\" as the copyright notice for easier\n identification within third-party archives.\n\n Copyright [yyyy] [name of copyright owner]\n\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n\n http://www.apache.org/licenses/LICENSE-2.0\n\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n"
},
{
"path": "Makefile",
"content": ".PHONY: quality style test docs-single-version docs\n# this target runs checks on all files\nquality:\n\truff check .\n\tmypy onnxtr/\n\n# this target runs checks on all files and potentially modifies some of them\nstyle:\n\truff format .\n\truff check --fix .\n\n# Run tests for the library\ntest:\n\tcoverage run -m pytest tests/common/ -rs --memray\n\tcoverage report --fail-under=80 --show-missing\n\n# Check that docs can build\ndocs-single-version:\n\tsphinx-build docs/source docs/_build -a\n\n# Check that docs can build\ndocs:\n\tcd docs && bash build.sh"
},
{
"path": "README.md",
"content": "\n ![](\"https://github.com/felixdittrich92/OnnxTR/raw/main/docs/images/logo.jpg\")
\n
\n\n[](LICENSE)\n\n[](https://codecov.io/gh/felixdittrich92/OnnxTR)\n[](https://app.codacy.com/gh/felixdittrich92/OnnxTR/dashboard?utm_source=gh&utm_medium=referral&utm_content=&utm_campaign=Badge_grade)\n[](https://www.codefactor.io/repository/github/felixdittrich92/onnxtr)\n[](https://socket.dev/pypi/package/onnxtr/overview/0.8.1/tar-gz)\n[](https://pypi.org/project/OnnxTR/)\n[](https://github.com/felixdittrich92/OnnxTR/pkgs/container/onnxtr)\n[](https://huggingface.co/spaces/Felix92/OnnxTR-OCR)\n\n\n> :warning: Please note that this is a wrapper around the [doctr](https://github.com/mindee/doctr) library to provide a Onnx pipeline for docTR. For feature requests, which are not directly related to the Onnx pipeline, please refer to the base project.\n\n**Optical Character Recognition made seamless & accessible to anyone, powered by Onnx**\n\nWhat you can expect from this repository:\n\n- efficient ways to parse textual information (localize and identify each word) from your documents\n- a Onnx pipeline for docTR, a wrapper around the [doctr](https://github.com/mindee/doctr) library - no PyTorch or TensorFlow dependencies\n- more lightweight package with faster inference latency and less required resources\n- 8-Bit quantized models for faster inference on CPU\n\n\n\n## Installation\n\n### Prerequisites\n\nPython 3.10 (or higher) and [pip](https://pip.pypa.io/en/stable/) are required to install OnnxTR.\n\n### Latest release\n\nYou can then install the latest release of the package using [pypi](https://pypi.org/project/OnnxTR/) as follows:\n\n**NOTE:**\n\nCurrently supported execution providers by default are: CPU, CUDA (NVIDIA GPU), OpenVINO (Intel CPU | GPU), CoreML (Apple Silicon).\n\nFor GPU support please take a look at: [ONNX Runtime](https://onnxruntime.ai/getting-started).\n\n- **Prerequisites:** CUDA & cuDNN needs to be installed before [Version table](https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html).\n\n```shell\n# standard cpu support\npip install \"onnxtr[cpu]\"\npip install \"onnxtr[cpu-headless]\" # same as cpu but with opencv-headless\n# with gpu support\npip install \"onnxtr[gpu]\"\npip install \"onnxtr[gpu-headless]\" # same as gpu but with opencv-headless\n# OpenVINO cpu | gpu support for Intel CPUs | GPUs\npip install \"onnxtr[openvino]\"\npip install \"onnxtr[openvino-headless]\" # same as openvino but with opencv-headless\n# with HTML support\npip install \"onnxtr[html]\"\n# with support for visualization\npip install \"onnxtr[viz]\"\n# with support for all dependencies\npip install \"onnxtr[html, gpu, viz]\"\n```\n\n**Recommendation:**\n\nIf you have:\n\n- a NVIDIA GPU, use one of the `gpu` variants\n- an Intel CPU or GPU, use one of the `openvino` variants\n- an Apple Silicon Mac, use one of the `cpu` variants (CoreML is auto-detected)\n- otherwise, use one of the `cpu` variants\n\n**OpenVINO:**\n\nBy default OnnxTR running with the OpenVINO execution provider backend uses the `CPU` device with `FP32` precision, to change the device or for further configuaration please refer to the [ONNX Runtime OpenVINO documentation](https://onnxruntime.ai/docs/execution-providers/OpenVINO-ExecutionProvider.html#summary-of-options).\n\n### Reading files\n\nDocuments can be interpreted from PDF / Images / Webpages / Multiple page images using the following code snippet:\n\n```python\nfrom onnxtr.io import DocumentFile\n\n# PDF\npdf_doc = DocumentFile.from_pdf(\"path/to/your/doc.pdf\")\n# Image\nsingle_img_doc = DocumentFile.from_images(\"path/to/your/img.jpg\")\n# Webpage (requires `weasyprint` to be installed)\nwebpage_doc = DocumentFile.from_url(\"https://www.yoursite.com\")\n# Multiple page images\nmulti_img_doc = DocumentFile.from_images([\"path/to/page1.jpg\", \"path/to/page2.jpg\"])\n```\n\n### Putting it together\n\nLet's use the default `ocr_predictor` model for an example:\n\n```python\nfrom onnxtr.io import DocumentFile\nfrom onnxtr.models import ocr_predictor, EngineConfig\n\nmodel = ocr_predictor(\n det_arch=\"fast_base\", # detection architecture\n reco_arch=\"vitstr_base\", # recognition architecture\n det_bs=2, # detection batch size\n reco_bs=512, # recognition batch size\n # Document related parameters\n assume_straight_pages=True, # set to `False` if the pages are not straight (rotation, perspective, etc.) (default: True)\n straighten_pages=False, # set to `True` if the pages should be straightened before final processing (default: False)\n export_as_straight_boxes=False, # set to `True` if the boxes should be exported as if the pages were straight (default: False)\n # Preprocessing related parameters\n preserve_aspect_ratio=True, # set to `False` if the aspect ratio should not be preserved (default: True)\n symmetric_pad=True, # set to `False` to disable symmetric padding (default: True)\n # Additional parameters - meta information\n detect_orientation=False, # set to `True` if the orientation of the pages should be detected (default: False)\n detect_language=False, # set to `True` if the language of the pages should be detected (default: False)\n # Orientation specific parameters in combination with `assume_straight_pages=False` and/or `straighten_pages=True`\n disable_crop_orientation=False, # set to `True` if the crop orientation classification should be disabled (default: False)\n disable_page_orientation=False, # set to `True` if the general page orientation classification should be disabled (default: False)\n # DocumentBuilder specific parameters\n resolve_lines=True, # whether words should be automatically grouped into lines (default: True)\n resolve_blocks=False, # whether lines should be automatically grouped into blocks (default: False)\n paragraph_break=0.035, # relative length of the minimum space separating paragraphs (default: 0.035)\n # OnnxTR specific parameters\n # NOTE: 8-Bit quantized models are not available for FAST detection models and can in general lead to poorer accuracy\n load_in_8_bit=False, # set to `True` to load 8-bit quantized models instead of the full precision onces (default: False)\n # Advanced engine configuration options\n det_engine_cfg=EngineConfig(), # detection model engine configuration (default: internal predefined configuration)\n reco_engine_cfg=EngineConfig(), # recognition model engine configuration (default: internal predefined configuration)\n clf_engine_cfg=EngineConfig(), # classification (orientation) model engine configuration (default: internal predefined configuration)\n)\n# PDF\ndoc = DocumentFile.from_pdf(\"path/to/your/doc.pdf\")\n# Analyze\nresult = model(doc)\n# Display the result (requires matplotlib & mplcursors to be installed)\nresult.show()\n```\n\n\n\nOr even rebuild the original document from its predictions:\n\n```python\nimport matplotlib.pyplot as plt\n\nsynthetic_pages = result.synthesize()\nplt.imshow(synthetic_pages[0])\nplt.axis(\"off\")\nplt.show()\n```\n\n\n\nThe `ocr_predictor` returns a `Document` object with a nested structure (with `Page`, `Block`, `Line`, `Word`, `Artefact`).\nTo get a better understanding of the document model, check out [documentation](https://mindee.github.io/doctr/modules/io.html#document-structure):\n\nYou can also export them as a nested dict, more appropriate for JSON format / render it or export as XML (hocr format):\n\n```python\njson_output = result.export() # nested dict\ntext_output = result.render() # human-readable text\nxml_output = result.export_as_xml() # hocr format\nfor output in xml_output:\n xml_bytes_string = output[0]\n xml_element = output[1]\n```\n\n\n Advanced engine configuration options
\n\nYou can also define advanced engine configurations for the models / predictors:\n\n```python\nfrom onnxruntime import SessionOptions\n\nfrom onnxtr.models import ocr_predictor, EngineConfig\n\ngeneral_options = (\n SessionOptions()\n) # For configuartion options see: https://onnxruntime.ai/docs/api/python/api_summary.html#sessionoptions\ngeneral_options.enable_cpu_mem_arena = False\n\n# NOTE: The following would force to run only on the GPU if no GPU is available it will raise an error\n# List of strings e.g. [\"CUDAExecutionProvider\", \"CPUExecutionProvider\"] or a list of tuples with the provider and its options e.g.\n# [(\"CUDAExecutionProvider\", {\"device_id\": 0}), (\"CPUExecutionProvider\", {\"arena_extend_strategy\": \"kSameAsRequested\"})]\nproviders = [\n (\"CUDAExecutionProvider\", {\"device_id\": 0, \"cudnn_conv_algo_search\": \"DEFAULT\"})\n] # For available providers see: https://onnxruntime.ai/docs/execution-providers/\n\nengine_config = EngineConfig(session_options=general_options, providers=providers)\n# We use the default predictor with the custom engine configuration\n# NOTE: You can define differnt engine configurations for detection, recognition and classification depending on your needs\npredictor = ocr_predictor(det_engine_cfg=engine_config, reco_engine_cfg=engine_config, clf_engine_cfg=engine_config)\n```\n\nYou can also dynamically configure whether the memory arena should shrink:\n\n```python\nfrom random import random\nfrom onnxruntime import RunOptions, SessionOptions\n\nfrom onnxtr.models import ocr_predictor, EngineConfig\n\n\ndef arena_shrinkage_handler(run_options: RunOptions) -> RunOptions:\n \"\"\"\n Shrink the memory arena on 10% of inference runs.\n \"\"\"\n if random() < 0.1:\n run_options.add_run_config_entry(\"memory.enable_memory_arena_shrinkage\", \"cpu:0\")\n return run_options\n\n\nengine_config = EngineConfig(run_options_provider=arena_shrinkage_handler)\nengine_config.session_options.enable_mem_pattern = False\n\npredictor = ocr_predictor(det_engine_cfg=engine_config, reco_engine_cfg=engine_config, clf_engine_cfg=engine_config)\n```\n\n \n\n## Loading custom exported models\n\nYou can also load docTR custom exported models:\nFor exporting please take a look at the [doctr documentation](https://mindee.github.io/doctr/using_doctr/using_model_export.html#export-to-onnx).\n\n```python\nfrom onnxtr.models import ocr_predictor, linknet_resnet18, parseq\n\nreco_model = parseq(\"path_to_custom_model.onnx\", vocab=\"ABC\")\ndet_model = linknet_resnet18(\"path_to_custom_model.onnx\")\nmodel = ocr_predictor(det_arch=det_model, reco_arch=reco_model)\n```\n\n## Loading models from HuggingFace Hub\n\nYou can also load models from the HuggingFace Hub:\n\n```python\nfrom onnxtr.io import DocumentFile\nfrom onnxtr.models import ocr_predictor, from_hub\n\nimg = DocumentFile.from_images([\"\"])\n# Load your model from the hub\nmodel = from_hub(\"onnxtr/my-model\")\n\n# Pass it to the predictor\n# If your model is a recognition model:\npredictor = ocr_predictor(det_arch=\"db_mobilenet_v3_large\", reco_arch=model)\n\n# If your model is a detection model:\npredictor = ocr_predictor(det_arch=model, reco_arch=\"crnn_mobilenet_v3_small\")\n\n# Get your predictions\nres = predictor(img)\n```\n\nHF Hub search: [here](https://huggingface.co/models?search=onnxtr).\n\nCollection: [here](https://huggingface.co/collections/Felix92/onnxtr-66bf213a9f88f7346c90e842)\n\nOr push your own models to the hub:\n\n```python\nfrom onnxtr.models import parseq, push_to_hf_hub, login_to_hub\nfrom onnxtr.utils.vocabs import VOCABS\n\n# Login to the hub\nlogin_to_hub()\n\n# Recogniton model\nmodel = parseq(\"~/onnxtr-parseq-multilingual-v1.onnx\", vocab=VOCABS[\"multilingual\"])\npush_to_hf_hub(\n model,\n model_name=\"onnxtr-parseq-multilingual-v1\",\n task=\"recognition\", # The task for which the model is intended [detection, recognition, classification]\n arch=\"parseq\", # The name of the model architecture\n override=False, # Set to `True` if you want to override an existing model / repository\n)\n\n# Detection model\nmodel = linknet_resnet18(\"~/onnxtr-linknet-resnet18.onnx\")\npush_to_hf_hub(model, model_name=\"onnxtr-linknet-resnet18\", task=\"detection\", arch=\"linknet_resnet18\", override=True)\n```\n\n## Models architectures\n\nCredits where it's due: this repository provides ONNX models for the following architectures, converted from the docTR models:\n\n### Text Detection\n\n- DBNet: [Real-time Scene Text Detection with Differentiable Binarization](https://arxiv.org/pdf/1911.08947.pdf).\n- LinkNet: [LinkNet: Exploiting Encoder Representations for Efficient Semantic Segmentation](https://arxiv.org/pdf/1707.03718.pdf)\n- FAST: [FAST: Faster Arbitrarily-Shaped Text Detector with Minimalist Kernel Representation](https://arxiv.org/pdf/2111.02394.pdf)\n\n### Text Recognition\n\n- CRNN: [An End-to-End Trainable Neural Network for Image-based Sequence Recognition and Its Application to Scene Text Recognition](https://arxiv.org/pdf/1507.05717.pdf).\n- SAR: [Show, Attend and Read:A Simple and Strong Baseline for Irregular Text Recognition](https://arxiv.org/pdf/1811.00751.pdf).\n- MASTER: [MASTER: Multi-Aspect Non-local Network for Scene Text Recognition](https://arxiv.org/pdf/1910.02562.pdf).\n- ViTSTR: [Vision Transformer for Fast and Efficient Scene Text Recognition](https://arxiv.org/pdf/2105.08582.pdf).\n- PARSeq: [Scene Text Recognition with Permuted Autoregressive Sequence Models](https://arxiv.org/pdf/2207.06966).\n- VIPTR: [A Vision Permutable Extractor for Fast and Efficient Scene Text Recognition](https://arxiv.org/abs/2401.10110).\n\n```python\npredictor = ocr_predictor()\npredictor.list_archs()\n{\n \"detection archs\": [\n \"db_resnet34\",\n \"db_resnet50\",\n \"db_mobilenet_v3_large\",\n \"linknet_resnet18\",\n \"linknet_resnet34\",\n \"linknet_resnet50\",\n \"fast_tiny\", # No 8-bit support\n \"fast_small\", # No 8-bit support\n \"fast_base\", # No 8-bit support\n ],\n \"recognition archs\": [\n \"crnn_vgg16_bn\",\n \"crnn_mobilenet_v3_small\",\n \"crnn_mobilenet_v3_large\",\n \"sar_resnet31\",\n \"master\",\n \"vitstr_small\",\n \"vitstr_base\",\n \"parseqviptr_tiny\", # No 8-bit support\n ],\n}\n```\n\n### Documentation\n\nThis repository is in sync with the [doctr](https://github.com/mindee/doctr) library, which provides a high-level API to perform OCR on documents.\nThis repository stays up-to-date with the latest features and improvements from the base project.\nSo we can refer to the [doctr documentation](https://mindee.github.io/doctr/) for more detailed information.\n\nNOTE:\n\n- `pretrained` is the default in OnnxTR, and not available as a parameter.\n- docTR specific environment variables (e.g.: DOCTR_CACHE_DIR -> ONNXTR_CACHE_DIR) needs to be replaced with `ONNXTR_` prefix.\n\n### Benchmarks\n\nThe CPU benchmarks was measured on a `i7-14700K Intel CPU`.\n\nThe GPU benchmarks was measured on a `RTX 4080 Nvidia GPU`.\n\nBenchmarking performed on the FUNSD dataset and CORD dataset.\n\ndocTR / OnnxTR models used for the benchmarks are `fast_base` (full precision) | `db_resnet50` (8-bit variant) for detection and `crnn_vgg16_bn` for recognition.\n\nThe smallest combination in OnnxTR (docTR) of `db_mobilenet_v3_large` and `crnn_mobilenet_v3_small` takes as comparison `~0.17s / Page` on the FUNSD dataset and `~0.12s / Page` on the CORD dataset in **full precision** on CPU.\n\n- CPU benchmarks:\n\n|Library |FUNSD (199 pages) |CORD (900 pages) |\n|------------------------------------|-------------------------------|-------------------------------|\n|docTR (CPU) - v0.8.1 | ~1.29s / Page | ~0.60s / Page |\n|**OnnxTR (CPU)** - v0.6.0 | ~0.57s / Page | **~0.25s / Page** |\n|**OnnxTR (CPU) 8-bit** - v0.6.0 | **~0.38s / Page** | **~0.14s / Page** |\n|**OnnxTR (CPU-OpenVINO)** - v0.6.0 | **~0.15s / Page** | **~0.14s / Page** |\n|EasyOCR (CPU) - v1.7.1 | ~1.96s / Page | ~1.75s / Page |\n|**PyTesseract (CPU)** - v0.3.10 | **~0.50s / Page** | ~0.52s / Page |\n|Surya (line) (CPU) - v0.4.4 | ~48.76s / Page | ~35.49s / Page |\n|PaddleOCR (CPU) - no cls - v2.7.3 | ~1.27s / Page | ~0.38s / Page |\n\n- GPU benchmarks:\n\n|Library |FUNSD (199 pages) |CORD (900 pages) |\n|-------------------------------------|-------------------------------|-------------------------------|\n|docTR (GPU) - v0.8.1 | ~0.07s / Page | ~0.05s / Page |\n|**docTR (GPU) float16** - v0.8.1 | **~0.06s / Page** | **~0.03s / Page** |\n|OnnxTR (GPU) - v0.6.0 | **~0.06s / Page** | ~0.04s / Page |\n|**OnnxTR (GPU) float16 - v0.6.0** | **~0.05s / Page** | **~0.03s / Page** |\n|EasyOCR (GPU) - v1.7.1 | ~0.31s / Page | ~0.19s / Page |\n|Surya (GPU) float16 - v0.4.4 | ~3.70s / Page | ~2.81s / Page |\n|**PaddleOCR (GPU) - no cls - v2.7.3**| ~0.08s / Page | **~0.03s / Page** |\n\n## Citation\n\nIf you wish to cite please refer to the base project citation, feel free to use this [BibTeX](http://www.bibtex.org/) reference:\n\n```bibtex\n@misc{doctr2021,\n title={docTR: Document Text Recognition},\n author={Mindee},\n year={2021},\n publisher = {GitHub},\n howpublished = {\\url{https://github.com/mindee/doctr}}\n}\n```\n\n```bibtex\n@misc{onnxtr2024,\n title={OnnxTR: Optical Character Recognition made seamless & accessible to anyone, powered by Onnx},\n author={Felix Dittrich},\n year={2024},\n publisher = {GitHub},\n howpublished = {\\url{https://github.com/felixdittrich92/OnnxTR}}\n}\n```\n\n## License\n\nDistributed under the Apache 2.0 License. See [`LICENSE`](https://github.com/felixdittrich92/OnnxTR?tab=Apache-2.0-1-ov-file#readme) for more information.\n"
},
{
"path": "demo/README.md",
"content": "---\ntitle: OnnxTR OCR\nemoji: 🔥\ncolorFrom: red\ncolorTo: purple\nsdk: gradio\nsdk_version: 5.34.2\napp_file: app.py\npinned: false\nlicense: apache-2.0\n---\n\nCheck out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference\n\n## Run the demo locally\n\n```bash\ncd demo\npip install -r requirements.txt\npython3 app.py\n```\n"
},
{
"path": "demo/app.py",
"content": "import io\nimport os\nfrom typing import Any\n\n# NOTE: This is a fix to run the demo on the HuggingFace Zero GPU or CPU spaces\nif os.environ.get(\"SPACES_ZERO_GPU\") is not None:\n import spaces\nelse:\n\n class spaces: # noqa: N801\n @staticmethod\n def GPU(func): # noqa: N802\n def wrapper(*args, **kwargs):\n return func(*args, **kwargs)\n\n return wrapper\n\n\nimport cv2\nimport gradio as gr\nimport matplotlib.pyplot as plt\nimport numpy as np\nfrom matplotlib.figure import Figure\nfrom PIL import Image\n\nfrom onnxtr.io import DocumentFile\nfrom onnxtr.models import EngineConfig, from_hub, ocr_predictor\nfrom onnxtr.models.predictor import OCRPredictor\nfrom onnxtr.utils.visualization import visualize_page\n\nDET_ARCHS: list[str] = [\n \"fast_base\",\n \"fast_small\",\n \"fast_tiny\",\n \"db_resnet50\",\n \"db_resnet34\",\n \"db_mobilenet_v3_large\",\n \"linknet_resnet18\",\n \"linknet_resnet34\",\n \"linknet_resnet50\",\n]\nRECO_ARCHS: list[str] = [\n \"crnn_vgg16_bn\",\n \"crnn_mobilenet_v3_small\",\n \"crnn_mobilenet_v3_large\",\n \"master\",\n \"sar_resnet31\",\n \"vitstr_small\",\n \"vitstr_base\",\n \"parseq\",\n \"viptr_tiny\",\n]\n\nCUSTOM_RECO_ARCHS: list[str] = [\n \"Felix92/onnxtr-parseq-multilingual-v1\",\n]\n\n\ndef load_predictor(\n det_arch: str,\n reco_arch: str,\n use_gpu: bool,\n assume_straight_pages: bool,\n straighten_pages: bool,\n export_as_straight_boxes: bool,\n detect_language: bool,\n load_in_8_bit: bool,\n bin_thresh: float,\n box_thresh: float,\n disable_crop_orientation: bool = False,\n disable_page_orientation: bool = False,\n) -> OCRPredictor:\n \"\"\"Load a predictor from doctr.models\n\n Args:\n ----\n det_arch: detection architecture\n reco_arch: recognition architecture\n use_gpu: whether to use the GPU or not\n assume_straight_pages: whether to assume straight pages or not\n disable_crop_orientation: whether to disable crop orientation or not\n disable_page_orientation: whether to disable page orientation or not\n straighten_pages: whether to straighten rotated pages or not\n export_as_straight_boxes: whether to export straight boxes\n detect_language: whether to detect the language of the text\n load_in_8_bit: whether to load the image in 8 bit mode\n bin_thresh: binarization threshold for the segmentation map\n box_thresh: minimal objectness score to consider a box\n\n Returns:\n -------\n instance of OCRPredictor\n \"\"\"\n engine_cfg = (\n EngineConfig()\n if use_gpu\n else EngineConfig(providers=[(\"CPUExecutionProvider\", {\"arena_extend_strategy\": \"kSameAsRequested\"})])\n )\n predictor = ocr_predictor(\n det_arch=det_arch,\n reco_arch=reco_arch if reco_arch not in CUSTOM_RECO_ARCHS else from_hub(reco_arch),\n assume_straight_pages=assume_straight_pages,\n straighten_pages=straighten_pages,\n detect_language=detect_language,\n load_in_8_bit=load_in_8_bit,\n export_as_straight_boxes=export_as_straight_boxes,\n detect_orientation=not assume_straight_pages,\n disable_crop_orientation=disable_crop_orientation,\n disable_page_orientation=disable_page_orientation,\n det_engine_cfg=engine_cfg,\n reco_engine_cfg=engine_cfg,\n clf_engine_cfg=engine_cfg,\n )\n predictor.det_predictor.model.postprocessor.bin_thresh = bin_thresh\n predictor.det_predictor.model.postprocessor.box_thresh = box_thresh\n return predictor\n\n\ndef forward_image(predictor: OCRPredictor, image: np.ndarray) -> np.ndarray:\n \"\"\"Forward an image through the predictor\n\n Args:\n ----\n predictor: instance of OCRPredictor\n image: image to process\n\n Returns:\n -------\n segmentation map\n \"\"\"\n processed_batches = predictor.det_predictor.pre_processor([image])\n out = predictor.det_predictor.model(processed_batches[0], return_model_output=True)\n seg_map = out[\"out_map\"]\n\n return seg_map\n\n\ndef matplotlib_to_pil(fig: Figure | np.ndarray) -> Image.Image:\n \"\"\"Convert a matplotlib figure to a PIL image\n\n Args:\n ----\n fig: matplotlib figure or numpy array\n\n Returns:\n -------\n PIL image\n \"\"\"\n buf = io.BytesIO()\n if isinstance(fig, Figure):\n fig.savefig(buf)\n else:\n plt.imsave(buf, fig)\n buf.seek(0)\n return Image.open(buf)\n\n\n@spaces.GPU\ndef analyze_page(\n uploaded_file: Any,\n page_idx: int,\n det_arch: str,\n reco_arch: str,\n use_gpu: bool,\n assume_straight_pages: bool,\n disable_crop_orientation: bool,\n disable_page_orientation: bool,\n straighten_pages: bool,\n export_as_straight_boxes: bool,\n detect_language: bool,\n load_in_8_bit: bool,\n bin_thresh: float,\n box_thresh: float,\n):\n \"\"\"Analyze a page\n\n Args:\n ----\n uploaded_file: file to analyze\n page_idx: index of the page to analyze\n det_arch: detection architecture\n reco_arch: recognition architecture\n use_gpu: whether to use the GPU or not\n assume_straight_pages: whether to assume straight pages or not\n disable_crop_orientation: whether to disable crop orientation or not\n disable_page_orientation: whether to disable page orientation or not\n straighten_pages: whether to straighten rotated pages or not\n export_as_straight_boxes: whether to export straight boxes\n detect_language: whether to detect the language of the text\n load_in_8_bit: whether to load the image in 8 bit mode\n bin_thresh: binarization threshold for the segmentation map\n box_thresh: minimal objectness score to consider a box\n\n Returns:\n -------\n input image, segmentation heatmap, output image, OCR output, synthesized page\n \"\"\"\n if uploaded_file is None:\n return None, \"Please upload a document\", None, None, None\n\n if uploaded_file.name.endswith(\".pdf\"):\n doc = DocumentFile.from_pdf(uploaded_file)\n else:\n doc = DocumentFile.from_images(uploaded_file)\n try:\n page = doc[page_idx - 1]\n except IndexError:\n page = doc[-1]\n\n img = page\n\n predictor = load_predictor(\n det_arch=det_arch,\n reco_arch=reco_arch,\n use_gpu=use_gpu,\n assume_straight_pages=assume_straight_pages,\n straighten_pages=straighten_pages,\n export_as_straight_boxes=export_as_straight_boxes,\n detect_language=detect_language,\n load_in_8_bit=load_in_8_bit,\n bin_thresh=bin_thresh,\n box_thresh=box_thresh,\n disable_crop_orientation=disable_crop_orientation,\n disable_page_orientation=disable_page_orientation,\n )\n\n seg_map = forward_image(predictor, page)\n seg_map = np.squeeze(seg_map)\n seg_map = cv2.resize(seg_map, (img.shape[1], img.shape[0]), interpolation=cv2.INTER_LINEAR)\n seg_heatmap = matplotlib_to_pil(seg_map)\n\n out = predictor([page])\n\n page_export = out.pages[0].export()\n fig = visualize_page(out.pages[0].export(), out.pages[0].page, interactive=False, add_labels=False)\n\n out_img = matplotlib_to_pil(fig)\n\n if assume_straight_pages or (not assume_straight_pages and straighten_pages):\n synthesized_page = out.pages[0].synthesize()\n else:\n synthesized_page = None\n\n return img, seg_heatmap, out_img, page_export, synthesized_page\n\n\nwith gr.Blocks(fill_height=True) as demo:\n gr.HTML(\n \"\"\"\n \n
\n \n
\n\n
OnnxTR OCR Demo
\n\n
\n \n ![\"GitHub](\"https://img.shields.io/badge/GitHub-blue?logo=github\")
\n \n \n ![\"PyPI\"](\"https://img.shields.io/pypi/v/onnxtr?color=blue\")
\n \n
\n
To use this interactive demo for OnnxTR:
\n 1. Upload a document (PDF, JPG, or PNG)
\n 2. Select the model architectures for text detection and recognition you want to use
\n 3. Press the \"Analyze page\" button to process the uploaded document
\n \"\"\"\n )\n with gr.Row():\n with gr.Column(scale=1):\n upload = gr.File(label=\"Upload File [JPG | PNG | PDF]\", file_types=[\".pdf\", \".jpg\", \".png\"])\n page_selection = gr.Slider(minimum=1, maximum=10, step=1, value=1, label=\"Page selection\")\n det_model = gr.Dropdown(choices=DET_ARCHS, value=DET_ARCHS[0], label=\"Text detection model\")\n reco_model = gr.Dropdown(\n choices=RECO_ARCHS + CUSTOM_RECO_ARCHS, value=RECO_ARCHS[0], label=\"Text recognition model\"\n )\n use_gpu = gr.Checkbox(value=True, label=\"Use GPU\")\n assume_straight = gr.Checkbox(value=True, label=\"Assume straight pages\")\n disable_crop_orientation = gr.Checkbox(value=False, label=\"Disable crop orientation\")\n disable_page_orientation = gr.Checkbox(value=False, label=\"Disable page orientation\")\n straighten = gr.Checkbox(value=False, label=\"Straighten pages\")\n export_as_straight_boxes = gr.Checkbox(value=False, label=\"Export as straight boxes\")\n det_language = gr.Checkbox(value=False, label=\"Detect language\")\n load_in_8_bit = gr.Checkbox(value=False, label=\"Load 8-bit quantized models\")\n binarization_threshold = gr.Slider(\n minimum=0.1, maximum=0.9, value=0.3, step=0.1, label=\"Binarization threshold\"\n )\n box_threshold = gr.Slider(minimum=0.1, maximum=0.9, value=0.1, step=0.1, label=\"Box threshold\")\n analyze_button = gr.Button(\"Analyze page\")\n with gr.Column(scale=3):\n with gr.Row():\n input_image = gr.Image(label=\"Input page\", width=700, height=500)\n segmentation_heatmap = gr.Image(label=\"Segmentation heatmap\", width=700, height=500)\n output_image = gr.Image(label=\"Output page\", width=700, height=500)\n with gr.Row():\n with gr.Column(scale=3):\n ocr_output = gr.JSON(label=\"OCR output\", render=True, scale=1, height=500)\n with gr.Column(scale=3):\n synthesized_page = gr.Image(label=\"Synthesized page\", width=700, height=500)\n\n analyze_button.click(\n analyze_page,\n inputs=[\n upload,\n page_selection,\n det_model,\n reco_model,\n use_gpu,\n assume_straight,\n disable_crop_orientation,\n disable_page_orientation,\n straighten,\n export_as_straight_boxes,\n det_language,\n load_in_8_bit,\n binarization_threshold,\n box_threshold,\n ],\n outputs=[input_image, segmentation_heatmap, output_image, ocr_output, synthesized_page],\n )\n\ndemo.launch(inbrowser=True, allowed_paths=[\"./data/logo.jpg\"])\n"
},
{
"path": "demo/packages.txt",
"content": "python3-opencv\nfonts-freefont-ttf\n"
},
{
"path": "demo/requirements.txt",
"content": "-e \"onnxtr[gpu-headless,viz] @ git+https://github.com/felixdittrich92/OnnxTR.git\"\ngradio>=5.30.0,<7.0.0\nspaces>=0.37.0\n\n# Quick fix to avoid HuggingFace Spaces cudnn9.x Cuda12.x issue\n# NOTE: outdated\n# onnxruntime-gpu==1.19.0\n"
},
{
"path": "onnxtr/__init__.py",
"content": "from . import io, models, contrib, transforms, utils\nfrom .version import __version__ # noqa: F401\n"
},
{
"path": "onnxtr/contrib/__init__.py",
"content": "from .artefacts import ArtefactDetector"
},
{
"path": "onnxtr/contrib/artefacts.py",
"content": "# Copyright (C) 2021-2026, Mindee | Felix Dittrich.\n\n# This program is licensed under the Apache License 2.0.\n# See LICENSE or go to for full license details.\n\nfrom typing import Any\n\nimport cv2\nimport numpy as np\n\nfrom onnxtr.file_utils import requires_package\n\nfrom .base import _BasePredictor\n\n__all__ = [\"ArtefactDetector\"]\n\ndefault_cfgs: dict[str, dict[str, Any]] = {\n \"yolov8_artefact\": {\n \"input_shape\": (3, 1024, 1024),\n \"labels\": [\"bar_code\", \"qr_code\", \"logo\", \"photo\"],\n \"url\": \"https://github.com/felixdittrich92/OnnxTR/releases/download/v0.0.1/yolo_artefact-f9d66f14.onnx\",\n },\n}\n\n\nclass ArtefactDetector(_BasePredictor):\n \"\"\"\n A class to detect artefacts in images\n\n >>> from onnxtr.io import DocumentFile\n >>> from onnxtr.contrib.artefacts import ArtefactDetector\n >>> doc = DocumentFile.from_images([\"path/to/image.jpg\"])\n >>> detector = ArtefactDetector()\n >>> results = detector(doc)\n\n Args:\n arch: the architecture to use\n batch_size: the batch size to use\n model_path: the path to the model to use\n labels: the labels to use\n input_shape: the input shape to use\n mask_labels: the mask labels to use\n conf_threshold: the confidence threshold to use\n iou_threshold: the intersection over union threshold to use\n **kwargs: additional arguments to be passed to `download_from_url`\n \"\"\"\n\n def __init__(\n self,\n arch: str = \"yolov8_artefact\",\n batch_size: int = 2,\n model_path: str | None = None,\n labels: list[str] | None = None,\n input_shape: tuple[int, int, int] | None = None,\n conf_threshold: float = 0.5,\n iou_threshold: float = 0.5,\n **kwargs: Any,\n ) -> None:\n super().__init__(batch_size=batch_size, url=default_cfgs[arch][\"url\"], model_path=model_path, **kwargs)\n self.labels = labels or default_cfgs[arch][\"labels\"]\n self.input_shape = input_shape or default_cfgs[arch][\"input_shape\"]\n self.conf_threshold = conf_threshold\n self.iou_threshold = iou_threshold\n\n def preprocess(self, img: np.ndarray) -> np.ndarray:\n return np.transpose(cv2.resize(img, (self.input_shape[2], self.input_shape[1])), (2, 0, 1)) / np.array(255.0)\n\n def postprocess(self, output: list[np.ndarray], input_images: list[list[np.ndarray]]) -> list[list[dict[str, Any]]]:\n results = []\n\n for batch in zip(output, input_images):\n for out, img in zip(batch[0], batch[1]):\n org_height, org_width = img.shape[:2]\n width_scale, height_scale = org_width / self.input_shape[2], org_height / self.input_shape[1]\n for res in out:\n sample_results = []\n for row in np.transpose(np.squeeze(res)):\n classes_scores = row[4:]\n max_score = np.amax(classes_scores)\n if max_score >= self.conf_threshold:\n class_id = np.argmax(classes_scores)\n x, y, w, h = row[0], row[1], row[2], row[3]\n # to rescaled xmin, ymin, xmax, ymax\n xmin = int((x - w / 2) * width_scale)\n ymin = int((y - h / 2) * height_scale)\n xmax = int((x + w / 2) * width_scale)\n ymax = int((y + h / 2) * height_scale)\n\n sample_results.append({\n \"label\": self.labels[class_id],\n \"confidence\": float(max_score),\n \"box\": [xmin, ymin, xmax, ymax],\n })\n\n # Filter out overlapping boxes\n boxes = [res[\"box\"] for res in sample_results]\n scores = [res[\"confidence\"] for res in sample_results]\n keep_indices = cv2.dnn.NMSBoxes(boxes, scores, self.conf_threshold, self.iou_threshold) # type: ignore[arg-type]\n sample_results = [sample_results[i] for i in keep_indices]\n\n results.append(sample_results)\n\n self._results = results\n return results\n\n def show(self, **kwargs: Any) -> None:\n \"\"\"\n Display the results\n\n Args:\n **kwargs: additional keyword arguments to be passed to `plt.show`\n \"\"\"\n requires_package(\"matplotlib\", \"`.show()` requires matplotlib installed\")\n import matplotlib.pyplot as plt\n from matplotlib.patches import Rectangle\n\n # visualize the results with matplotlib\n if self._results and self._inputs:\n for img, res in zip(self._inputs, self._results):\n plt.figure(figsize=(10, 10))\n plt.imshow(img)\n for obj in res:\n xmin, ymin, xmax, ymax = obj[\"box\"]\n label = obj[\"label\"]\n plt.text(xmin, ymin, f\"{label} {obj['confidence']:.2f}\", color=\"red\")\n plt.gca().add_patch(\n Rectangle((xmin, ymin), xmax - xmin, ymax - ymin, fill=False, edgecolor=\"red\", linewidth=2)\n )\n plt.show(**kwargs)\n"
},
{
"path": "onnxtr/contrib/base.py",
"content": "# Copyright (C) 2021-2026, Mindee | Felix Dittrich.\n\n# This program is licensed under the Apache License 2.0.\n# See LICENSE or go to for full license details.\n\nfrom typing import Any\n\nimport numpy as np\nimport onnxruntime as ort\n\nfrom onnxtr.utils.data import download_from_url\n\n\nclass _BasePredictor:\n \"\"\"\n Base class for all predictors\n\n Args:\n batch_size: the batch size to use\n url: the url to use to download a model if needed\n model_path: the path to the model to use\n **kwargs: additional arguments to be passed to `download_from_url`\n \"\"\"\n\n def __init__(self, batch_size: int, url: str | None = None, model_path: str | None = None, **kwargs) -> None:\n self.batch_size = batch_size\n self.session = self._init_model(url, model_path, **kwargs)\n\n self._inputs: list[np.ndarray] = []\n self._results: list[Any] = []\n\n def _init_model(self, url: str | None = None, model_path: str | None = None, **kwargs: Any) -> Any:\n \"\"\"\n Download the model from the given url if needed\n\n Args:\n url: the url to use\n model_path: the path to the model to use\n **kwargs: additional arguments to be passed to `download_from_url`\n\n Returns:\n Any: the ONNX loaded model\n \"\"\"\n if not url and not model_path:\n raise ValueError(\"You must provide either a url or a model_path\")\n onnx_model_path = model_path if model_path else str(download_from_url(url, cache_subdir=\"models\", **kwargs)) # type: ignore[arg-type]\n return ort.InferenceSession(onnx_model_path, providers=[\"CUDAExecutionProvider\", \"CPUExecutionProvider\"])\n\n def preprocess(self, img: np.ndarray) -> np.ndarray:\n \"\"\"\n Preprocess the input image\n\n Args:\n img: the input image to preprocess\n\n Returns:\n np.ndarray: the preprocessed image\n \"\"\"\n raise NotImplementedError\n\n def postprocess(self, output: list[np.ndarray], input_images: list[list[np.ndarray]]) -> Any:\n \"\"\"\n Postprocess the model output\n\n Args:\n output: the model output to postprocess\n input_images: the input images used to generate the output\n\n Returns:\n Any: the postprocessed output\n \"\"\"\n raise NotImplementedError\n\n def __call__(self, inputs: list[np.ndarray]) -> Any:\n \"\"\"\n Call the model on the given inputs\n\n Args:\n inputs: the inputs to use\n\n Returns:\n Any: the postprocessed output\n \"\"\"\n self._inputs = inputs\n model_inputs = self.session.get_inputs()\n\n batched_inputs = [inputs[i : i + self.batch_size] for i in range(0, len(inputs), self.batch_size)]\n processed_batches = [\n np.array([self.preprocess(img) for img in batch], dtype=np.float32) for batch in batched_inputs\n ]\n\n outputs = [self.session.run(None, {model_inputs[0].name: batch}) for batch in processed_batches]\n return self.postprocess(outputs, batched_inputs)\n"
},
{
"path": "onnxtr/file_utils.py",
"content": "# Copyright (C) 2021-2026, Mindee | Felix Dittrich.\n\n# This program is licensed under the Apache License 2.0.\n# See LICENSE or go to for full license details.\n\nimport importlib.metadata\nimport logging\n\n__all__ = [\"requires_package\"]\n\nENV_VARS_TRUE_VALUES = {\"1\", \"ON\", \"YES\", \"TRUE\"}\nENV_VARS_TRUE_AND_AUTO_VALUES = ENV_VARS_TRUE_VALUES.union({\"AUTO\"})\n\n\ndef requires_package(name: str, extra_message: str | None = None) -> None: # pragma: no cover\n \"\"\"\n package requirement helper\n\n Args:\n name: name of the package\n extra_message: additional message to display if the package is not found\n \"\"\"\n try:\n _pkg_version = importlib.metadata.version(name)\n logging.info(f\"{name} version {_pkg_version} available.\")\n except importlib.metadata.PackageNotFoundError:\n raise ImportError(\n f\"\\n\\n{extra_message if extra_message is not None else ''} \"\n f\"\\nPlease install it with the following command: pip install {name}\\n\"\n )\n"
},
{
"path": "onnxtr/io/__init__.py",
"content": "from .elements import *\nfrom .html import *\nfrom .image import *\nfrom .pdf import *\nfrom .reader import *\n"
},
{
"path": "onnxtr/io/elements.py",
"content": "# Copyright (C) 2021-2026, Mindee | Felix Dittrich.\n\n# This program is licensed under the Apache License 2.0.\n# See LICENSE or go to for full license details.\n\nfrom typing import Any\n\nfrom defusedxml import defuse_stdlib\n\ndefuse_stdlib()\nfrom xml.etree import ElementTree as ET\nfrom xml.etree.ElementTree import Element as ETElement\nfrom xml.etree.ElementTree import SubElement\n\nimport numpy as np\n\nimport onnxtr\nfrom onnxtr.file_utils import requires_package\nfrom onnxtr.utils.common_types import BoundingBox\nfrom onnxtr.utils.geometry import resolve_enclosing_bbox, resolve_enclosing_rbbox\nfrom onnxtr.utils.reconstitution import synthesize_page\nfrom onnxtr.utils.repr import NestedObject\n\ntry: # optional dependency for visualization\n from onnxtr.utils.visualization import visualize_page\nexcept ModuleNotFoundError: # pragma: no cover\n pass\n\n__all__ = [\"Element\", \"Word\", \"Artefact\", \"Line\", \"Block\", \"Page\", \"Document\"]\n\n\nclass Element(NestedObject):\n \"\"\"Implements an abstract document element with exporting and text rendering capabilities\"\"\"\n\n _children_names: list[str] = []\n _exported_keys: list[str] = []\n\n def __init__(self, **kwargs: Any) -> None:\n for k, v in kwargs.items():\n if k in self._children_names:\n setattr(self, k, v)\n else:\n raise KeyError(f\"{self.__class__.__name__} object does not have any attribute named '{k}'\")\n\n def export(self) -> dict[str, Any]:\n \"\"\"Exports the object into a nested dict format\"\"\"\n export_dict = {k: getattr(self, k) for k in self._exported_keys}\n for children_name in self._children_names:\n export_dict[children_name] = [c.export() for c in getattr(self, children_name)]\n\n return export_dict\n\n @classmethod\n def from_dict(cls, save_dict: dict[str, Any], **kwargs):\n raise NotImplementedError\n\n def render(self) -> str:\n raise NotImplementedError\n\n\nclass Word(Element):\n \"\"\"Implements a word element\n\n Args:\n value: the text string of the word\n confidence: the confidence associated with the text prediction\n geometry: bounding box of the word in format ((xmin, ymin), (xmax, ymax)) where coordinates are relative to\n the page's size\n objectness_score: the objectness score of the detection\n crop_orientation: the general orientation of the crop in degrees and its confidence\n \"\"\"\n\n _exported_keys: list[str] = [\"value\", \"confidence\", \"geometry\", \"objectness_score\", \"crop_orientation\"]\n _children_names: list[str] = []\n\n def __init__(\n self,\n value: str,\n confidence: float,\n geometry: BoundingBox | np.ndarray,\n objectness_score: float,\n crop_orientation: dict[str, Any],\n ) -> None:\n super().__init__()\n self.value = value\n self.confidence = confidence\n self.geometry = geometry\n self.objectness_score = objectness_score\n self.crop_orientation = crop_orientation\n\n def render(self) -> str:\n \"\"\"Renders the full text of the element\"\"\"\n return self.value\n\n def extra_repr(self) -> str:\n return f\"value='{self.value}', confidence={self.confidence:.2}\"\n\n @classmethod\n def from_dict(cls, save_dict: dict[str, Any], **kwargs):\n kwargs = {k: save_dict[k] for k in cls._exported_keys}\n return cls(**kwargs)\n\n\nclass Artefact(Element):\n \"\"\"Implements a non-textual element\n\n Args:\n artefact_type: the type of artefact\n confidence: the confidence of the type prediction\n geometry: bounding box of the word in format ((xmin, ymin), (xmax, ymax)) where coordinates are relative to\n the page's size.\n \"\"\"\n\n _exported_keys: list[str] = [\"geometry\", \"type\", \"confidence\"]\n _children_names: list[str] = []\n\n def __init__(self, artefact_type: str, confidence: float, geometry: BoundingBox) -> None:\n super().__init__()\n self.geometry = geometry\n self.type = artefact_type\n self.confidence = confidence\n\n def render(self) -> str:\n \"\"\"Renders the full text of the element\"\"\"\n return f\"[{self.type.upper()}]\"\n\n def extra_repr(self) -> str:\n return f\"type='{self.type}', confidence={self.confidence:.2}\"\n\n @classmethod\n def from_dict(cls, save_dict: dict[str, Any], **kwargs):\n kwargs = {k: save_dict[k] for k in cls._exported_keys}\n return cls(**kwargs)\n\n\nclass Line(Element):\n \"\"\"Implements a line element as a collection of words\n\n Args:\n words: list of word elements\n geometry: bounding box of the word in format ((xmin, ymin), (xmax, ymax)) where coordinates are relative to\n the page's size. If not specified, it will be resolved by default to the smallest bounding box enclosing\n all words in it.\n \"\"\"\n\n _exported_keys: list[str] = [\"geometry\", \"objectness_score\"]\n _children_names: list[str] = [\"words\"]\n words: list[Word] = []\n\n def __init__(\n self,\n words: list[Word],\n geometry: BoundingBox | np.ndarray | None = None,\n objectness_score: float | None = None,\n ) -> None:\n # Compute the objectness score of the line\n if objectness_score is None:\n objectness_score = float(np.mean([w.objectness_score for w in words]))\n # Resolve the geometry using the smallest enclosing bounding box\n if geometry is None:\n # Check whether this is a rotated or straight box\n box_resolution_fn = resolve_enclosing_rbbox if len(words[0].geometry) == 4 else resolve_enclosing_bbox\n geometry = box_resolution_fn([w.geometry for w in words]) # type: ignore[misc]\n\n super().__init__(words=words)\n self.geometry = geometry\n self.objectness_score = objectness_score\n\n def render(self) -> str:\n \"\"\"Renders the full text of the element\"\"\"\n return \" \".join(w.render() for w in self.words)\n\n @classmethod\n def from_dict(cls, save_dict: dict[str, Any], **kwargs):\n kwargs = {k: save_dict[k] for k in cls._exported_keys}\n kwargs.update({\n \"words\": [Word.from_dict(_dict) for _dict in save_dict[\"words\"]],\n })\n return cls(**kwargs)\n\n\nclass Block(Element):\n \"\"\"Implements a block element as a collection of lines and artefacts\n\n Args:\n lines: list of line elements\n artefacts: list of artefacts\n geometry: bounding box of the word in format ((xmin, ymin), (xmax, ymax)) where coordinates are relative to\n the page's size. If not specified, it will be resolved by default to the smallest bounding box enclosing\n all lines and artefacts in it.\n \"\"\"\n\n _exported_keys: list[str] = [\"geometry\", \"objectness_score\"]\n _children_names: list[str] = [\"lines\", \"artefacts\"]\n lines: list[Line] = []\n artefacts: list[Artefact] = []\n\n def __init__(\n self,\n lines: list[Line] = [],\n artefacts: list[Artefact] = [],\n geometry: BoundingBox | np.ndarray | None = None,\n objectness_score: float | None = None,\n ) -> None:\n # Compute the objectness score of the line\n if objectness_score is None:\n objectness_score = float(np.mean([w.objectness_score for line in lines for w in line.words]))\n # Resolve the geometry using the smallest enclosing bounding box\n if geometry is None:\n line_boxes = [word.geometry for line in lines for word in line.words]\n artefact_boxes = [artefact.geometry for artefact in artefacts]\n box_resolution_fn = (\n resolve_enclosing_rbbox if isinstance(lines[0].geometry, np.ndarray) else resolve_enclosing_bbox\n )\n geometry = box_resolution_fn(line_boxes + artefact_boxes) # type: ignore\n\n super().__init__(lines=lines, artefacts=artefacts)\n self.geometry = geometry\n self.objectness_score = objectness_score\n\n def render(self, line_break: str = \"\\n\") -> str:\n \"\"\"Renders the full text of the element\"\"\"\n return line_break.join(line.render() for line in self.lines)\n\n @classmethod\n def from_dict(cls, save_dict: dict[str, Any], **kwargs):\n kwargs = {k: save_dict[k] for k in cls._exported_keys}\n kwargs.update({\n \"lines\": [Line.from_dict(_dict) for _dict in save_dict[\"lines\"]],\n \"artefacts\": [Artefact.from_dict(_dict) for _dict in save_dict[\"artefacts\"]],\n })\n return cls(**kwargs)\n\n\nclass Page(Element):\n \"\"\"Implements a page element as a collection of blocks\n\n Args:\n page: image encoded as a numpy array in uint8\n blocks: list of block elements\n page_idx: the index of the page in the input raw document\n dimensions: the page size in pixels in format (height, width)\n orientation: a dictionary with the value of the rotation angle in degress and confidence of the prediction\n language: a dictionary with the language value and confidence of the prediction\n \"\"\"\n\n _exported_keys: list[str] = [\"page_idx\", \"dimensions\", \"orientation\", \"language\"]\n _children_names: list[str] = [\"blocks\"]\n blocks: list[Block] = []\n\n def __init__(\n self,\n page: np.ndarray,\n blocks: list[Block],\n page_idx: int,\n dimensions: tuple[int, int],\n orientation: dict[str, Any] | None = None,\n language: dict[str, Any] | None = None,\n ) -> None:\n super().__init__(blocks=blocks)\n self.page = page\n self.page_idx = page_idx\n self.dimensions = dimensions\n self.orientation = orientation if isinstance(orientation, dict) else dict(value=None, confidence=None)\n self.language = language if isinstance(language, dict) else dict(value=None, confidence=None)\n\n def render(self, block_break: str = \"\\n\\n\") -> str:\n \"\"\"Renders the full text of the element\"\"\"\n return block_break.join(b.render() for b in self.blocks)\n\n def extra_repr(self) -> str:\n return f\"dimensions={self.dimensions}\"\n\n def show(self, interactive: bool = True, preserve_aspect_ratio: bool = False, **kwargs) -> None:\n \"\"\"Overlay the result on a given image\n\n Args:\n interactive: whether the display should be interactive\n preserve_aspect_ratio: pass True if you passed True to the predictor\n **kwargs: additional keyword arguments passed to the matplotlib.pyplot.show method\n \"\"\"\n requires_package(\"matplotlib\", \"`.show()` requires matplotlib & mplcursors installed\")\n requires_package(\"mplcursors\", \"`.show()` requires matplotlib & mplcursors installed\")\n import matplotlib.pyplot as plt\n\n visualize_page(self.export(), self.page, interactive=interactive, preserve_aspect_ratio=preserve_aspect_ratio)\n plt.show(**kwargs)\n\n def synthesize(self, **kwargs) -> np.ndarray:\n \"\"\"Synthesize the page from the predictions\n\n Args:\n **kwargs: keyword arguments passed to the `synthesize_page` method\n\n Returns\n synthesized page\n \"\"\"\n return synthesize_page(self.export(), **kwargs)\n\n def export_as_xml(self, file_title: str = \"OnnxTR - XML export (hOCR)\") -> tuple[bytes, ET.ElementTree]:\n \"\"\"Export the page as XML (hOCR-format)\n convention: https://github.com/kba/hocr-spec/blob/master/1.2/spec.md\n\n Args:\n file_title: the title of the XML file\n\n Returns:\n a tuple of the XML byte string, and its ElementTree\n \"\"\"\n p_idx = self.page_idx\n block_count: int = 1\n line_count: int = 1\n word_count: int = 1\n height, width = self.dimensions\n language = self.language if \"language\" in self.language.keys() else \"en\"\n # Create the XML root element\n page_hocr = ETElement(\"html\", attrib={\"xmlns\": \"http://www.w3.org/1999/xhtml\", \"xml:lang\": str(language)})\n # Create the header / SubElements of the root element\n head = SubElement(page_hocr, \"head\")\n SubElement(head, \"title\").text = file_title\n SubElement(head, \"meta\", attrib={\"http-equiv\": \"Content-Type\", \"content\": \"text/html; charset=utf-8\"})\n SubElement(\n head,\n \"meta\",\n attrib={\"name\": \"ocr-system\", \"content\": f\"onnxtr {onnxtr.__version__}\"}, # type: ignore[attr-defined]\n )\n SubElement(\n head,\n \"meta\",\n attrib={\"name\": \"ocr-capabilities\", \"content\": \"ocr_page ocr_carea ocr_par ocr_line ocrx_word\"},\n )\n # Create the body\n body = SubElement(page_hocr, \"body\")\n page_div = SubElement(\n body,\n \"div\",\n attrib={\n \"class\": \"ocr_page\",\n \"id\": f\"page_{p_idx + 1}\",\n \"title\": f\"image; bbox 0 0 {width} {height}; ppageno 0\",\n },\n )\n # iterate over the blocks / lines / words and create the XML elements in body line by line with the attributes\n for block in self.blocks:\n if len(block.geometry) != 2:\n raise TypeError(\"XML export is only available for straight bounding boxes for now.\")\n (xmin, ymin), (xmax, ymax) = block.geometry\n block_div = SubElement(\n page_div,\n \"div\",\n attrib={\n \"class\": \"ocr_carea\",\n \"id\": f\"block_{block_count}\",\n \"title\": f\"bbox {int(round(xmin * width))} {int(round(ymin * height))} \\\n {int(round(xmax * width))} {int(round(ymax * height))}\",\n },\n )\n paragraph = SubElement(\n block_div,\n \"p\",\n attrib={\n \"class\": \"ocr_par\",\n \"id\": f\"par_{block_count}\",\n \"title\": f\"bbox {int(round(xmin * width))} {int(round(ymin * height))} \\\n {int(round(xmax * width))} {int(round(ymax * height))}\",\n },\n )\n block_count += 1\n for line in block.lines:\n (xmin, ymin), (xmax, ymax) = line.geometry\n # NOTE: baseline, x_size, x_descenders, x_ascenders is currently initalized to 0\n line_span = SubElement(\n paragraph,\n \"span\",\n attrib={\n \"class\": \"ocr_line\",\n \"id\": f\"line_{line_count}\",\n \"title\": f\"bbox {int(round(xmin * width))} {int(round(ymin * height))} \\\n {int(round(xmax * width))} {int(round(ymax * height))}; \\\n baseline 0 0; x_size 0; x_descenders 0; x_ascenders 0\",\n },\n )\n line_count += 1\n for word in line.words:\n (xmin, ymin), (xmax, ymax) = word.geometry\n conf = word.confidence\n word_div = SubElement(\n line_span,\n \"span\",\n attrib={\n \"class\": \"ocrx_word\",\n \"id\": f\"word_{word_count}\",\n \"title\": f\"bbox {int(round(xmin * width))} {int(round(ymin * height))} \\\n {int(round(xmax * width))} {int(round(ymax * height))}; \\\n x_wconf {int(round(conf * 100))}\",\n },\n )\n # set the text\n word_div.text = word.value\n word_count += 1\n\n return (ET.tostring(page_hocr, encoding=\"utf-8\", method=\"xml\"), ET.ElementTree(page_hocr))\n\n @classmethod\n def from_dict(cls, save_dict: dict[str, Any], **kwargs):\n kwargs = {k: save_dict[k] for k in cls._exported_keys}\n kwargs.update({\"blocks\": [Block.from_dict(block_dict) for block_dict in save_dict[\"blocks\"]]})\n return cls(**kwargs)\n\n\nclass Document(Element):\n \"\"\"Implements a document element as a collection of pages\n\n Args:\n pages: list of page elements\n \"\"\"\n\n _children_names: list[str] = [\"pages\"]\n pages: list[Page] = []\n\n def __init__(\n self,\n pages: list[Page],\n ) -> None:\n super().__init__(pages=pages)\n\n def render(self, page_break: str = \"\\n\\n\\n\\n\") -> str:\n \"\"\"Renders the full text of the element\"\"\"\n return page_break.join(p.render() for p in self.pages)\n\n def show(self, **kwargs) -> None:\n \"\"\"Overlay the result on a given image\"\"\"\n for result in self.pages:\n result.show(**kwargs)\n\n def synthesize(self, **kwargs) -> list[np.ndarray]:\n \"\"\"Synthesize all pages from their predictions\n\n Args:\n **kwargs: keyword arguments passed to the `Page.synthesize` method\n\n Returns:\n list of synthesized pages\n \"\"\"\n return [page.synthesize(**kwargs) for page in self.pages]\n\n def export_as_xml(self, **kwargs) -> list[tuple[bytes, ET.ElementTree]]:\n \"\"\"Export the document as XML (hOCR-format)\n\n Args:\n **kwargs: additional keyword arguments passed to the Page.export_as_xml method\n\n Returns:\n list of tuple of (bytes, ElementTree)\n \"\"\"\n return [page.export_as_xml(**kwargs) for page in self.pages]\n\n @classmethod\n def from_dict(cls, save_dict: dict[str, Any], **kwargs):\n kwargs = {k: save_dict[k] for k in cls._exported_keys}\n kwargs.update({\"pages\": [Page.from_dict(page_dict) for page_dict in save_dict[\"pages\"]]})\n return cls(**kwargs)\n"
},
{
"path": "onnxtr/io/html.py",
"content": "# Copyright (C) 2021-2026, Mindee | Felix Dittrich.\n\n# This program is licensed under the Apache License 2.0.\n# See LICENSE or go to for full license details.\n\nfrom typing import Any\n\n__all__ = [\"read_html\"]\n\n\ndef read_html(url: str, **kwargs: Any) -> bytes:\n \"\"\"Read a PDF file and convert it into an image in numpy format\n\n >>> from onnxtr.io import read_html\n >>> doc = read_html(\"https://www.yoursite.com\")\n\n Args:\n url: URL of the target web page\n **kwargs: keyword arguments from `weasyprint.HTML`\n\n Returns:\n decoded PDF file as a bytes stream\n \"\"\"\n from weasyprint import HTML\n\n return HTML(url, **kwargs).write_pdf()\n"
},
{
"path": "onnxtr/io/image.py",
"content": "# Copyright (C) 2021-2026, Mindee | Felix Dittrich.\n\n# This program is licensed under the Apache License 2.0.\n# See LICENSE or go to for full license details.\n\nfrom pathlib import Path\n\nimport cv2\nimport numpy as np\n\nfrom onnxtr.utils.common_types import AbstractFile\n\n__all__ = [\"read_img_as_numpy\"]\n\n\ndef read_img_as_numpy(\n file: AbstractFile,\n output_size: tuple[int, int] | None = None,\n rgb_output: bool = True,\n) -> np.ndarray:\n \"\"\"Read an image file into numpy format\n\n >>> from onnxtr.io import read_img_as_numpy\n >>> page = read_img_as_numpy(\"path/to/your/doc.jpg\")\n\n Args:\n file: the path to the image file\n output_size: the expected output size of each page in format H x W\n rgb_output: whether the output ndarray channel order should be RGB instead of BGR.\n\n Returns:\n the page decoded as numpy ndarray of shape H x W x 3\n \"\"\"\n if isinstance(file, (str, Path)):\n if not Path(file).is_file():\n raise FileNotFoundError(f\"unable to access {file}\")\n img = cv2.imread(str(file), cv2.IMREAD_COLOR)\n elif isinstance(file, bytes):\n _file: np.ndarray = np.frombuffer(file, np.uint8)\n img = cv2.imdecode(_file, cv2.IMREAD_COLOR)\n else:\n raise TypeError(\"unsupported object type for argument 'file'\")\n\n # Validity check\n if img is None:\n raise ValueError(\"unable to read file.\")\n # Resizing\n if isinstance(output_size, tuple):\n img = cv2.resize(img, output_size[::-1], interpolation=cv2.INTER_LINEAR)\n # Switch the channel order\n if rgb_output:\n img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n return img\n"
},
{
"path": "onnxtr/io/pdf.py",
"content": "# Copyright (C) 2021-2026, Mindee | Felix Dittrich.\n\n# This program is licensed under the Apache License 2.0.\n# See LICENSE or go to for full license details.\n\nfrom typing import Any\n\nimport numpy as np\nimport pypdfium2 as pdfium\n\nfrom onnxtr.utils.common_types import AbstractFile\n\n__all__ = [\"read_pdf\"]\n\n\ndef read_pdf(\n file: AbstractFile,\n scale: int = 2,\n rgb_mode: bool = True,\n password: str | None = None,\n **kwargs: Any,\n) -> list[np.ndarray]:\n \"\"\"Read a PDF file and convert it into an image in numpy format\n\n >>> from onnxtr.io import read_pdf\n >>> doc = read_pdf(\"path/to/your/doc.pdf\")\n\n Args:\n file: the path to the PDF file\n scale: rendering scale (1 corresponds to 72dpi)\n rgb_mode: if True, the output will be RGB, otherwise BGR\n password: a password to unlock the document, if encrypted\n **kwargs: additional parameters to :meth:`pypdfium2.PdfPage.render`\n\n Returns:\n the list of pages decoded as numpy ndarray of shape H x W x C\n \"\"\"\n # Rasterise pages to numpy ndarrays with pypdfium2\n pdf = pdfium.PdfDocument(file, password=password)\n try:\n return [page.render(scale=scale, rev_byteorder=rgb_mode, **kwargs).to_numpy() for page in pdf]\n finally:\n pdf.close()\n"
},
{
"path": "onnxtr/io/reader.py",
"content": "# Copyright (C) 2021-2026, Mindee | Felix Dittrich.\n\n# This program is licensed under the Apache License 2.0.\n# See LICENSE or go to for full license details.\n\nfrom collections.abc import Sequence\nfrom pathlib import Path\n\nimport numpy as np\n\nfrom onnxtr.file_utils import requires_package\nfrom onnxtr.utils.common_types import AbstractFile\n\nfrom .html import read_html\nfrom .image import read_img_as_numpy\nfrom .pdf import read_pdf\n\n__all__ = [\"DocumentFile\"]\n\n\nclass DocumentFile:\n \"\"\"Read a document from multiple extensions\"\"\"\n\n @classmethod\n def from_pdf(cls, file: AbstractFile, **kwargs) -> list[np.ndarray]:\n \"\"\"Read a PDF file\n\n >>> from onnxtr.io import DocumentFile\n >>> doc = DocumentFile.from_pdf(\"path/to/your/doc.pdf\")\n\n Args:\n file: the path to the PDF file or a binary stream\n **kwargs: additional parameters to :meth:`pypdfium2.PdfPage.render`\n\n Returns:\n the list of pages decoded as numpy ndarray of shape H x W x 3\n \"\"\"\n return read_pdf(file, **kwargs)\n\n @classmethod\n def from_url(cls, url: str, **kwargs) -> list[np.ndarray]:\n \"\"\"Interpret a web page as a PDF document\n\n >>> from onnxtr.io import DocumentFile\n >>> doc = DocumentFile.from_url(\"https://www.yoursite.com\")\n\n Args:\n url: the URL of the target web page\n **kwargs: additional parameters to :meth:`pypdfium2.PdfPage.render`\n\n Returns:\n the list of pages decoded as numpy ndarray of shape H x W x 3\n \"\"\"\n requires_package(\n \"weasyprint\",\n \"`.from_url` requires weasyprint installed.\\n\"\n + \"Installation instructions: https://doc.courtbouillon.org/weasyprint/stable/first_steps.html#installation\",\n )\n pdf_stream = read_html(url)\n return cls.from_pdf(pdf_stream, **kwargs)\n\n @classmethod\n def from_images(cls, files: Sequence[AbstractFile] | AbstractFile, **kwargs) -> list[np.ndarray]:\n \"\"\"Read an image file (or a collection of image files) and convert it into an image in numpy format\n\n >>> from onnxtr.io import DocumentFile\n >>> pages = DocumentFile.from_images([\"path/to/your/page1.png\", \"path/to/your/page2.png\"])\n\n Args:\n files: the path to the image file or a binary stream, or a collection of those\n **kwargs: additional parameters to :meth:`onnxtr.io.image.read_img_as_numpy`\n\n Returns:\n the list of pages decoded as numpy ndarray of shape H x W x 3\n \"\"\"\n if isinstance(files, (str, Path, bytes)):\n files = [files]\n\n return [read_img_as_numpy(file, **kwargs) for file in files]\n"
},
{
"path": "onnxtr/models/__init__.py",
"content": "from .engine import EngineConfig\nfrom .classification import *\nfrom .detection import *\nfrom .recognition import *\nfrom .zoo import *\nfrom .factory import *\n"
},
{
"path": "onnxtr/models/_utils.py",
"content": "# Copyright (C) 2021-2026, Mindee | Felix Dittrich.\n\n# This program is licensed under the Apache License 2.0.\n# See LICENSE or go to for full license details.\n\nfrom math import floor\nfrom statistics import median_low\n\nimport cv2\nimport numpy as np\nfrom langdetect import LangDetectException, detect_langs\n\nfrom onnxtr.utils.geometry import rotate_image\n\n__all__ = [\"estimate_orientation\", \"get_language\"]\n\n\ndef get_max_width_length_ratio(contour: np.ndarray) -> float:\n \"\"\"Get the maximum shape ratio of a contour.\n\n Args:\n contour: the contour from cv2.findContour\n\n Returns:\n the maximum shape ratio\n \"\"\"\n _, (w, h), _ = cv2.minAreaRect(contour)\n if w == 0 or h == 0:\n return 0.0\n return max(w / h, h / w)\n\n\ndef estimate_orientation(\n img: np.ndarray,\n general_page_orientation: tuple[int, float] | None = None,\n n_ct: int = 70,\n ratio_threshold_for_lines: float = 3,\n min_confidence: float = 0.2,\n lower_area: int = 100,\n) -> int:\n \"\"\"Estimate the angle of the general document orientation based on the\n lines of the document and the assumption that they should be horizontal.\n\n Args:\n img: the img or bitmap to analyze (H, W, C)\n general_page_orientation: the general orientation of the page (angle [0, 90, 180, 270 (-90)], confidence)\n estimated by a model\n n_ct: the number of contours used for the orientation estimation\n ratio_threshold_for_lines: this is the ratio w/h used to discriminates lines\n min_confidence: the minimum confidence to consider the general_page_orientation\n lower_area: the minimum area of a contour to be considered\n\n Returns:\n the estimated angle of the page (clockwise, negative for left side rotation, positive for right side rotation)\n \"\"\"\n assert len(img.shape) == 3 and img.shape[-1] in [1, 3], f\"Image shape {img.shape} not supported\"\n\n # Convert image to grayscale if necessary\n if img.shape[-1] == 3:\n gray_img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)\n gray_img = cv2.medianBlur(gray_img, 5)\n thresh = cv2.threshold(gray_img, thresh=0, maxval=255, type=cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)[1]\n else:\n thresh = img.astype(np.uint8)\n\n page_orientation, orientation_confidence = general_page_orientation or (0, 0.0)\n is_confident = page_orientation is not None and orientation_confidence >= min_confidence\n base_angle = page_orientation if is_confident else 0\n\n if is_confident:\n # We rotate the image to the general orientation which improves the detection\n # No expand needed bitmap is already padded\n thresh = rotate_image(thresh, -base_angle)\n else: # That's only required if we do not work on the detection models bin map\n # try to merge words in lines\n (h, w) = img.shape[:2]\n k_x = max(1, (floor(w / 100)))\n k_y = max(1, (floor(h / 100)))\n kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (k_x, k_y))\n thresh = cv2.dilate(thresh, kernel, iterations=1)\n\n # extract contours\n contours, _ = cv2.findContours(thresh, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)\n\n # Filter & Sort contours\n contours = sorted(\n [contour for contour in contours if cv2.contourArea(contour) > lower_area],\n key=get_max_width_length_ratio,\n reverse=True,\n )\n\n angles = []\n for contour in contours[:n_ct]:\n _, (w, h), angle = cv2.minAreaRect(contour)\n\n # OpenCV version-proof normalization: force 'w' to be the long side\n # so the angle is consistently relative to the major axis.\n # https://github.com/opencv/opencv/pull/28051/changes\n if w < h:\n w, h = h, w\n angle -= 90\n\n # Normalize angle to be within [-90, 90]\n while angle <= -90:\n angle += 180\n while angle > 90:\n angle -= 180\n\n if h > 0:\n if w / h > ratio_threshold_for_lines: # select only contours with ratio like lines\n angles.append(angle)\n elif w / h < 1 / ratio_threshold_for_lines: # if lines are vertical, substract 90 degree\n angles.append(angle - 90)\n\n if len(angles) == 0:\n skew_angle = 0 # in case no angles is found\n else:\n # median_low picks a value from the data to avoid outliers\n median = -median_low(angles)\n skew_angle = -round(median) if abs(median) != 0 else 0\n\n # Resolve the 90-degree flip ambiguity.\n # If the estimation is exactly 90/-90, it's usually a vertical detection of horizontal lines.\n if abs(skew_angle) == 90:\n skew_angle = 0\n\n # combine with the general orientation and the estimated angle\n # Apply the detected skew to our base orientation\n final_angle = base_angle + skew_angle\n\n # Standardize result to [-179, 180] range to handle wrap-around cases (e.g., 180 + -31)\n while final_angle > 180:\n final_angle -= 360\n while final_angle <= -180:\n final_angle += 360\n\n if is_confident:\n # If the estimated angle is perpendicular, treat it as 0 to avoid wrong flips\n if abs(skew_angle) % 90 == 0:\n return page_orientation\n\n # special case where the estimated angle is mostly wrong:\n # case 1: - and + swapped\n # case 2: estimated angle is completely wrong\n # so in this case we prefer the general page orientation\n if abs(skew_angle) == abs(page_orientation) and page_orientation != 0:\n return page_orientation\n\n return int(\n final_angle\n ) # return the clockwise angle (negative - left side rotation, positive - right side rotation)\n\n\ndef rectify_crops(\n crops: list[np.ndarray],\n orientations: list[int],\n) -> list[np.ndarray]:\n \"\"\"Rotate each crop of the list according to the predicted orientation:\n 0: already straight, no rotation\n 1: 90 ccw, rotate 3 times ccw\n 2: 180, rotate 2 times ccw\n 3: 270 ccw, rotate 1 time ccw\n \"\"\"\n # Inverse predictions (if angle of +90 is detected, rotate by -90)\n orientations = [4 - pred if pred != 0 else 0 for pred in orientations]\n return (\n [crop if orientation == 0 else np.rot90(crop, orientation) for orientation, crop in zip(orientations, crops)]\n if len(orientations) > 0\n else []\n )\n\n\ndef rectify_loc_preds(\n page_loc_preds: np.ndarray,\n orientations: list[int],\n) -> np.ndarray | None:\n \"\"\"Orient the quadrangle (Polygon4P) according to the predicted orientation,\n so that the points are in this order: top L, top R, bot R, bot L if the crop is readable\n \"\"\"\n return (\n np.stack(\n [\n np.roll(page_loc_pred, orientation, axis=0)\n for orientation, page_loc_pred in zip(orientations, page_loc_preds)\n ],\n axis=0,\n )\n if len(orientations) > 0\n else None\n )\n\n\ndef get_language(text: str) -> tuple[str, float]:\n \"\"\"Get languages of a text using langdetect model.\n Get the language with the highest probability or no language if only a few words or a low probability\n\n Args:\n text (str): text\n\n Returns:\n The detected language in ISO 639 code and confidence score\n \"\"\"\n try:\n lang = detect_langs(text.lower())[0]\n except LangDetectException:\n return \"unknown\", 0.0\n if len(text) <= 1 or (len(text) <= 5 and lang.prob <= 0.2):\n return \"unknown\", 0.0\n return lang.lang, lang.prob\n"
},
{
"path": "onnxtr/models/builder.py",
"content": "# Copyright (C) 2021-2026, Mindee | Felix Dittrich.\n\n# This program is licensed under the Apache License 2.0.\n# See LICENSE or go to for full license details.\n\n\nfrom typing import Any\n\nimport numpy as np\nfrom scipy.cluster.hierarchy import fclusterdata\n\nfrom onnxtr.io.elements import Block, Document, Line, Page, Word\nfrom onnxtr.utils.geometry import estimate_page_angle, resolve_enclosing_bbox, resolve_enclosing_rbbox, rotate_boxes\nfrom onnxtr.utils.repr import NestedObject\n\n__all__ = [\"DocumentBuilder\"]\n\n\nclass DocumentBuilder(NestedObject):\n \"\"\"Implements a document builder\n\n Args:\n resolve_lines: whether words should be automatically grouped into lines\n resolve_blocks: whether lines should be automatically grouped into blocks\n paragraph_break: relative length of the minimum space separating paragraphs\n export_as_straight_boxes: if True, force straight boxes in the export (fit a rectangle\n box to all rotated boxes). Else, keep the boxes format unchanged, no matter what it is.\n \"\"\"\n\n def __init__(\n self,\n resolve_lines: bool = True,\n resolve_blocks: bool = False,\n paragraph_break: float = 0.035,\n export_as_straight_boxes: bool = False,\n ) -> None:\n self.resolve_lines = resolve_lines\n self.resolve_blocks = resolve_blocks\n self.paragraph_break = paragraph_break\n self.export_as_straight_boxes = export_as_straight_boxes\n\n @staticmethod\n def _sort_boxes(boxes: np.ndarray) -> tuple[np.ndarray, np.ndarray]:\n \"\"\"Sort bounding boxes from top to bottom, left to right\n\n Args:\n boxes: bounding boxes of shape (N, 4) or (N, 4, 2) (in case of rotated bbox)\n\n Returns:\n tuple: indices of ordered boxes of shape (N,), boxes\n If straight boxes are passed tpo the function, boxes are unchanged\n else: boxes returned are straight boxes fitted to the straightened rotated boxes\n so that we fit the lines afterwards to the straigthened page\n \"\"\"\n if boxes.ndim == 3:\n boxes = rotate_boxes(\n loc_preds=boxes,\n angle=-estimate_page_angle(boxes),\n orig_shape=(1024, 1024),\n min_angle=5.0,\n )\n boxes = np.concatenate((boxes.min(1), boxes.max(1)), -1)\n return (boxes[:, 0] + 2 * boxes[:, 3] / np.median(boxes[:, 3] - boxes[:, 1])).argsort(), boxes\n\n def _resolve_sub_lines(self, boxes: np.ndarray, word_idcs: list[int]) -> list[list[int]]:\n \"\"\"Split a line in sub_lines\n\n Args:\n boxes: bounding boxes of shape (N, 4)\n word_idcs: list of indexes for the words of the line\n\n Returns:\n A list of (sub-)lines computed from the original line (words)\n \"\"\"\n lines = []\n # Sort words horizontally\n word_idcs = [word_idcs[idx] for idx in boxes[word_idcs, 0].argsort().tolist()]\n\n # Eventually split line horizontally\n if len(word_idcs) < 2:\n lines.append(word_idcs)\n else:\n sub_line = [word_idcs[0]]\n for i in word_idcs[1:]:\n horiz_break = True\n\n prev_box = boxes[sub_line[-1]]\n # Compute distance between boxes\n dist = boxes[i, 0] - prev_box[2]\n # If distance between boxes is lower than paragraph break, same sub-line\n if dist < self.paragraph_break:\n horiz_break = False\n\n if horiz_break:\n lines.append(sub_line)\n sub_line = []\n\n sub_line.append(i)\n lines.append(sub_line)\n\n return lines\n\n def _resolve_lines(self, boxes: np.ndarray) -> list[list[int]]:\n \"\"\"Order boxes to group them in lines\n\n Args:\n boxes: bounding boxes of shape (N, 4) or (N, 4, 2) in case of rotated bbox\n\n Returns:\n nested list of box indices\n \"\"\"\n # Sort boxes, and straighten the boxes if they are rotated\n idxs, boxes = self._sort_boxes(boxes)\n\n # Compute median for boxes heights\n y_med = np.median(boxes[:, 3] - boxes[:, 1])\n\n lines = []\n words = [idxs[0]] # Assign the top-left word to the first line\n # Define a mean y-center for the line\n y_center_sum = boxes[idxs[0]][[1, 3]].mean()\n\n for idx in idxs[1:]:\n vert_break = True\n\n # Compute y_dist\n y_dist = abs(boxes[idx][[1, 3]].mean() - y_center_sum / len(words))\n # If y-center of the box is close enough to mean y-center of the line, same line\n if y_dist < y_med / 2:\n vert_break = False\n\n if vert_break:\n # Compute sub-lines (horizontal split)\n lines.extend(self._resolve_sub_lines(boxes, words))\n words = []\n y_center_sum = 0\n\n words.append(idx)\n y_center_sum += boxes[idx][[1, 3]].mean()\n\n # Use the remaining words to form the last(s) line(s)\n if len(words) > 0:\n # Compute sub-lines (horizontal split)\n lines.extend(self._resolve_sub_lines(boxes, words))\n\n return lines\n\n @staticmethod\n def _resolve_blocks(boxes: np.ndarray, lines: list[list[int]]) -> list[list[list[int]]]:\n \"\"\"Order lines to group them in blocks\n\n Args:\n boxes: bounding boxes of shape (N, 4) or (N, 4, 2)\n lines: list of lines, each line is a list of idx\n\n Returns:\n nested list of box indices\n \"\"\"\n # Resolve enclosing boxes of lines\n if boxes.ndim == 3:\n box_lines: np.ndarray = np.asarray([\n resolve_enclosing_rbbox([tuple(boxes[idx, :, :]) for idx in line]) # type: ignore[misc]\n for line in lines\n ])\n else:\n _box_lines = [\n resolve_enclosing_bbox([(tuple(boxes[idx, :2]), tuple(boxes[idx, 2:])) for idx in line])\n for line in lines\n ]\n box_lines = np.asarray([(x1, y1, x2, y2) for ((x1, y1), (x2, y2)) in _box_lines])\n\n # Compute geometrical features of lines to clusterize\n # Clusterizing only with box centers yield to poor results for complex documents\n if boxes.ndim == 3:\n box_features: np.ndarray = np.stack(\n (\n (box_lines[:, 0, 0] + box_lines[:, 0, 1]) / 2,\n (box_lines[:, 0, 0] + box_lines[:, 2, 0]) / 2,\n (box_lines[:, 0, 0] + box_lines[:, 2, 1]) / 2,\n (box_lines[:, 0, 1] + box_lines[:, 2, 1]) / 2,\n (box_lines[:, 0, 1] + box_lines[:, 2, 0]) / 2,\n (box_lines[:, 2, 0] + box_lines[:, 2, 1]) / 2,\n ),\n axis=-1,\n )\n else:\n box_features = np.stack(\n (\n (box_lines[:, 0] + box_lines[:, 3]) / 2,\n (box_lines[:, 1] + box_lines[:, 2]) / 2,\n (box_lines[:, 0] + box_lines[:, 2]) / 2,\n (box_lines[:, 1] + box_lines[:, 3]) / 2,\n box_lines[:, 0],\n box_lines[:, 1],\n ),\n axis=-1,\n )\n # Compute clusters\n clusters = fclusterdata(box_features, t=0.1, depth=4, criterion=\"distance\", metric=\"euclidean\")\n\n _blocks: dict[int, list[int]] = {}\n # Form clusters\n for line_idx, cluster_idx in enumerate(clusters):\n if cluster_idx in _blocks.keys():\n _blocks[cluster_idx].append(line_idx)\n else:\n _blocks[cluster_idx] = [line_idx]\n\n # Retrieve word-box level to return a fully nested structure\n blocks = [[lines[idx] for idx in block] for block in _blocks.values()]\n\n return blocks\n\n def _build_blocks(\n self,\n boxes: np.ndarray,\n objectness_scores: np.ndarray,\n word_preds: list[tuple[str, float]],\n crop_orientations: list[dict[str, Any]],\n ) -> list[Block]:\n \"\"\"Gather independent words in structured blocks\n\n Args:\n boxes: bounding boxes of all detected words of the page, of shape (N, 4) or (N, 4, 2)\n objectness_scores: objectness scores of all detected words of the page, of shape N\n word_preds: list of all detected words of the page, of shape N\n crop_orientations: list of dictoinaries containing\n the general orientation (orientations + confidences) of the crops\n\n Returns:\n list of block elements\n \"\"\"\n if boxes.shape[0] != len(word_preds):\n raise ValueError(f\"Incompatible argument lengths: {boxes.shape[0]}, {len(word_preds)}\")\n\n if boxes.shape[0] == 0:\n return []\n\n # Decide whether we try to form lines\n _boxes = boxes\n if self.resolve_lines:\n lines = self._resolve_lines(_boxes if _boxes.ndim == 3 else _boxes[:, :4])\n # Decide whether we try to form blocks\n if self.resolve_blocks and len(lines) > 1:\n _blocks = self._resolve_blocks(_boxes if _boxes.ndim == 3 else _boxes[:, :4], lines)\n else:\n _blocks = [lines]\n else:\n # Sort bounding boxes, one line for all boxes, one block for the line\n lines = [self._sort_boxes(_boxes if _boxes.ndim == 3 else _boxes[:, :4])[0]] # type: ignore[list-item]\n _blocks = [lines]\n\n blocks = [\n Block([\n Line([\n Word(\n *word_preds[idx],\n tuple(tuple(pt) for pt in boxes[idx].tolist()), # type: ignore[arg-type]\n float(objectness_scores[idx]),\n crop_orientations[idx],\n )\n if boxes.ndim == 3\n else Word(\n *word_preds[idx],\n ((boxes[idx, 0], boxes[idx, 1]), (boxes[idx, 2], boxes[idx, 3])),\n float(objectness_scores[idx]),\n crop_orientations[idx],\n )\n for idx in line\n ])\n for line in lines\n ])\n for lines in _blocks\n ]\n\n return blocks\n\n def extra_repr(self) -> str:\n return (\n f\"resolve_lines={self.resolve_lines}, resolve_blocks={self.resolve_blocks}, \"\n f\"paragraph_break={self.paragraph_break}, \"\n f\"export_as_straight_boxes={self.export_as_straight_boxes}\"\n )\n\n def __call__(\n self,\n pages: list[np.ndarray],\n boxes: list[np.ndarray],\n objectness_scores: list[np.ndarray],\n text_preds: list[list[tuple[str, float]]],\n page_shapes: list[tuple[int, int]],\n crop_orientations: list[dict[str, Any]],\n orientations: list[dict[str, Any]] | None = None,\n languages: list[dict[str, Any]] | None = None,\n ) -> Document:\n \"\"\"Re-arrange detected words into structured blocks\n\n Args:\n pages: list of N elements, where each element represents the page image\n boxes: list of N elements, where each element represents the localization predictions, of shape (*, 4)\n or (*, 4, 2) for all words for a given page\n objectness_scores: list of N elements, where each element represents the objectness scores\n text_preds: list of N elements, where each element is the list of all word prediction (text + confidence)\n page_shapes: shape of each page, of size N\n crop_orientations: list of N elements, where each element is\n a dictionary containing the general orientation (orientations + confidences) of the crops\n orientations: optional, list of N elements,\n where each element is a dictionary containing the orientation (orientation + confidence)\n languages: optional, list of N elements,\n where each element is a dictionary containing the language (language + confidence)\n\n Returns:\n document object\n \"\"\"\n if len(boxes) != len(text_preds) != len(crop_orientations) != len(objectness_scores) or len(boxes) != len(\n page_shapes\n ) != len(crop_orientations) != len(objectness_scores):\n raise ValueError(\"All arguments are expected to be lists of the same size\")\n\n _orientations = orientations if isinstance(orientations, list) else [None] * len(boxes)\n _languages = languages if isinstance(languages, list) else [None] * len(boxes)\n if self.export_as_straight_boxes and len(boxes) > 0:\n # If boxes are already straight OK, else fit a bounding rect\n if boxes[0].ndim == 3:\n # Iterate over pages and boxes\n boxes = [np.concatenate((p_boxes.min(1), p_boxes.max(1)), 1) for p_boxes in boxes]\n\n _pages = [\n Page(\n page,\n self._build_blocks(\n page_boxes,\n loc_scores,\n word_preds,\n word_crop_orientations,\n ),\n _idx,\n shape,\n orientation,\n language,\n )\n for page, _idx, shape, page_boxes, loc_scores, word_preds, word_crop_orientations, orientation, language in zip( # noqa: E501\n pages,\n range(len(boxes)),\n page_shapes,\n boxes,\n objectness_scores,\n text_preds,\n crop_orientations,\n _orientations,\n _languages,\n )\n ]\n\n return Document(_pages)\n"
},
{
"path": "onnxtr/models/classification/__init__.py",
"content": "from .models import *\nfrom .zoo import *\n"
},
{
"path": "onnxtr/models/classification/models/__init__.py",
"content": "from .mobilenet import *"
},
{
"path": "onnxtr/models/classification/models/mobilenet.py",
"content": "# Copyright (C) 2021-2026, Mindee | Felix Dittrich.\n\n# This program is licensed under the Apache License 2.0.\n# See LICENSE or go to for full license details.\n\n# Greatly inspired by https://github.com/pytorch/vision/blob/master/torchvision/models/mobilenetv3.py\n\nfrom copy import deepcopy\nfrom typing import Any\n\nimport numpy as np\n\nfrom ...engine import Engine, EngineConfig\n\n__all__ = [\n \"MobileNetV3\",\n \"mobilenet_v3_small_crop_orientation\",\n \"mobilenet_v3_small_page_orientation\",\n]\n\ndefault_cfgs: dict[str, dict[str, Any]] = {\n \"mobilenet_v3_small_crop_orientation\": {\n \"mean\": (0.694, 0.695, 0.693),\n \"std\": (0.299, 0.296, 0.301),\n \"input_shape\": (3, 256, 256),\n \"classes\": [0, -90, 180, 90],\n \"url\": \"https://github.com/felixdittrich92/OnnxTR/releases/download/v0.6.0/mobilenet_v3_small_crop_orientation-4fde60a1.onnx\",\n \"url_8_bit\": \"https://github.com/felixdittrich92/OnnxTR/releases/download/v0.6.0/mobilenet_v3_small_crop_orientation_static_8_bit-c32c7721.onnx\",\n },\n \"mobilenet_v3_small_page_orientation\": {\n \"mean\": (0.694, 0.695, 0.693),\n \"std\": (0.299, 0.296, 0.301),\n \"input_shape\": (3, 512, 512),\n \"classes\": [0, -90, 180, 90],\n \"url\": \"https://github.com/felixdittrich92/OnnxTR/releases/download/v0.6.0/mobilenet_v3_small_page_orientation-60606ce4.onnx\",\n \"url_8_bit\": \"https://github.com/felixdittrich92/OnnxTR/releases/download/v0.6.0/mobilenet_v3_small_page_orientation_static_8_bit-13b5b014.onnx\",\n },\n}\n\n\nclass MobileNetV3(Engine):\n \"\"\"MobileNetV3 Onnx loader\n\n Args:\n model_path: path or url to onnx model file\n engine_cfg: configuration for the inference engine\n cfg: configuration dictionary\n **kwargs: additional arguments to be passed to `Engine`\n \"\"\"\n\n def __init__(\n self,\n model_path: str,\n engine_cfg: EngineConfig | None = None,\n cfg: dict[str, Any] | None = None,\n **kwargs: Any,\n ) -> None:\n super().__init__(url=model_path, engine_cfg=engine_cfg, **kwargs)\n\n self.cfg = cfg\n\n def __call__(\n self,\n x: np.ndarray,\n ) -> np.ndarray:\n return self.run(x)\n\n\ndef _mobilenet_v3(\n arch: str,\n model_path: str,\n load_in_8_bit: bool = False,\n engine_cfg: EngineConfig | None = None,\n **kwargs: Any,\n) -> MobileNetV3:\n # Patch the url\n model_path = default_cfgs[arch][\"url_8_bit\"] if load_in_8_bit and \"http\" in model_path else model_path\n _cfg = deepcopy(default_cfgs[arch])\n return MobileNetV3(model_path, cfg=_cfg, engine_cfg=engine_cfg, **kwargs)\n\n\ndef mobilenet_v3_small_crop_orientation(\n model_path: str = default_cfgs[\"mobilenet_v3_small_crop_orientation\"][\"url\"],\n load_in_8_bit: bool = False,\n engine_cfg: EngineConfig | None = None,\n **kwargs: Any,\n) -> MobileNetV3:\n \"\"\"MobileNetV3-Small architecture as described in\n `\"Searching for MobileNetV3\",\n `_.\n\n >>> import numpy as np\n >>> from onnxtr.models import mobilenet_v3_small_crop_orientation\n >>> model = mobilenet_v3_small_crop_orientation()\n >>> input_tensor = np.random.rand((1, 3, 256, 256))\n >>> out = model(input_tensor)\n\n Args:\n model_path: path to onnx model file, defaults to url in default_cfgs\n load_in_8_bit: whether to load the the 8-bit quantized model, defaults to False\n engine_cfg: configuration for the inference engine\n **kwargs: keyword arguments of the MobileNetV3 architecture\n\n Returns:\n MobileNetV3\n \"\"\"\n return _mobilenet_v3(\"mobilenet_v3_small_crop_orientation\", model_path, load_in_8_bit, engine_cfg, **kwargs)\n\n\ndef mobilenet_v3_small_page_orientation(\n model_path: str = default_cfgs[\"mobilenet_v3_small_page_orientation\"][\"url\"],\n load_in_8_bit: bool = False,\n engine_cfg: EngineConfig | None = None,\n **kwargs: Any,\n) -> MobileNetV3:\n \"\"\"MobileNetV3-Small architecture as described in\n `\"Searching for MobileNetV3\",\n `_.\n\n >>> import numpy as np\n >>> from onnxtr.models import mobilenet_v3_small_page_orientation\n >>> model = mobilenet_v3_small_page_orientation()\n >>> input_tensor = np.random.rand((1, 3, 512, 512))\n >>> out = model(input_tensor)\n\n Args:\n model_path: path to onnx model file, defaults to url in default_cfgs\n load_in_8_bit: whether to load the the 8-bit quantized model, defaults to False\n engine_cfg: configuration for the inference engine\n **kwargs: keyword arguments of the MobileNetV3 architecture\n\n Returns:\n MobileNetV3\n \"\"\"\n return _mobilenet_v3(\"mobilenet_v3_small_page_orientation\", model_path, load_in_8_bit, engine_cfg, **kwargs)\n"
},
{
"path": "onnxtr/models/classification/predictor/__init__.py",
"content": "from .base import *\n"
},
{
"path": "onnxtr/models/classification/predictor/base.py",
"content": "# Copyright (C) 2021-2026, Mindee | Felix Dittrich.\n\n# This program is licensed under the Apache License 2.0.\n# See LICENSE or go to for full license details.\n\nfrom typing import Any\n\nimport numpy as np\nfrom scipy.special import softmax\n\nfrom onnxtr.models.preprocessor import PreProcessor\nfrom onnxtr.utils.repr import NestedObject\n\n__all__ = [\"OrientationPredictor\"]\n\n\nclass OrientationPredictor(NestedObject):\n \"\"\"Implements an object able to detect the reading direction of a text box or a page.\n 4 possible orientations: 0, 90, 180, 270 (-90) degrees counter clockwise.\n\n Args:\n pre_processor: transform inputs for easier batched model inference\n model: core classification architecture (backbone + classification head)\n load_in_8_bit: whether to load the the 8-bit quantized model, defaults to False\n \"\"\"\n\n _children_names: list[str] = [\"pre_processor\", \"model\"]\n\n def __init__(\n self,\n pre_processor: PreProcessor | None,\n model: Any | None,\n ) -> None:\n self.pre_processor = pre_processor if isinstance(pre_processor, PreProcessor) else None\n self.model = model\n\n def __call__(\n self,\n inputs: list[np.ndarray],\n ) -> list[list[int] | list[float]]:\n # Dimension check\n if any(input.ndim != 3 for input in inputs):\n raise ValueError(\"incorrect input shape: all inputs are expected to be multi-channel 2D images.\")\n\n if self.model is None or self.pre_processor is None:\n # predictor is disabled\n return [[0] * len(inputs), [0] * len(inputs), [1.0] * len(inputs)]\n\n processed_batches = self.pre_processor(inputs)\n predicted_batches = [self.model(batch) for batch in processed_batches]\n\n # confidence\n probs = [np.max(softmax(batch, axis=1), axis=1) for batch in predicted_batches]\n # Postprocess predictions\n predicted_batches = [np.argmax(out_batch, axis=1) for out_batch in predicted_batches]\n\n class_idxs = [int(pred) for batch in predicted_batches for pred in batch]\n classes = [int(self.model.cfg[\"classes\"][idx]) for idx in class_idxs]\n confs = [round(float(p), 2) for prob in probs for p in prob]\n\n return [class_idxs, classes, confs]\n"
},
{
"path": "onnxtr/models/classification/zoo.py",
"content": "# Copyright (C) 2021-2026, Mindee | Felix Dittrich.\n\n# This program is licensed under the Apache License 2.0.\n# See LICENSE or go to for full license details.\n\nfrom typing import Any\n\nfrom onnxtr.models.engine import EngineConfig\n\nfrom .. import classification\nfrom ..preprocessor import PreProcessor\nfrom .predictor import OrientationPredictor\n\n__all__ = [\"crop_orientation_predictor\", \"page_orientation_predictor\"]\n\nORIENTATION_ARCHS: list[str] = [\"mobilenet_v3_small_crop_orientation\", \"mobilenet_v3_small_page_orientation\"]\n\n\ndef _orientation_predictor(\n arch: Any,\n model_type: str,\n load_in_8_bit: bool = False,\n engine_cfg: EngineConfig | None = None,\n disabled: bool = False,\n **kwargs: Any,\n) -> OrientationPredictor:\n if disabled:\n # Case where the orientation predictor is disabled\n return OrientationPredictor(None, None)\n\n if isinstance(arch, str):\n if arch not in ORIENTATION_ARCHS:\n raise ValueError(f\"unknown architecture '{arch}'\")\n # Load directly classifier from backbone\n _model = classification.__dict__[arch](load_in_8_bit=load_in_8_bit, engine_cfg=engine_cfg)\n else:\n if not isinstance(arch, classification.MobileNetV3):\n raise ValueError(f\"unknown architecture: {type(arch)}\")\n _model = arch\n\n kwargs[\"mean\"] = kwargs.get(\"mean\", _model.cfg[\"mean\"])\n kwargs[\"std\"] = kwargs.get(\"std\", _model.cfg[\"std\"])\n kwargs[\"batch_size\"] = kwargs.get(\"batch_size\", 512 if model_type == \"crop\" else 2)\n input_shape = _model.cfg[\"input_shape\"][1:]\n predictor = OrientationPredictor(\n PreProcessor(input_shape, preserve_aspect_ratio=True, symmetric_pad=True, **kwargs),\n _model,\n )\n return predictor\n\n\ndef crop_orientation_predictor(\n arch: Any = \"mobilenet_v3_small_crop_orientation\",\n batch_size: int = 512,\n load_in_8_bit: bool = False,\n engine_cfg: EngineConfig | None = None,\n **kwargs: Any,\n) -> OrientationPredictor:\n \"\"\"Crop orientation classification architecture.\n\n >>> import numpy as np\n >>> from onnxtr.models import crop_orientation_predictor\n >>> model = crop_orientation_predictor(arch='mobilenet_v3_small_crop_orientation')\n >>> input_crop = (255 * np.random.rand(256, 256, 3)).astype(np.uint8)\n >>> out = model([input_crop])\n\n Args:\n arch: name of the architecture to use (e.g. 'mobilenet_v3_small_crop_orientation')\n batch_size: number of samples the model processes in parallel\n load_in_8_bit: load the 8-bit quantized version of the model\n engine_cfg: configuration of inference engine\n **kwargs: keyword arguments to be passed to the OrientationPredictor\n\n Returns:\n OrientationPredictor\n \"\"\"\n model_type = \"crop\"\n return _orientation_predictor(\n arch=arch,\n batch_size=batch_size,\n model_type=model_type,\n load_in_8_bit=load_in_8_bit,\n engine_cfg=engine_cfg,\n **kwargs,\n )\n\n\ndef page_orientation_predictor(\n arch: Any = \"mobilenet_v3_small_page_orientation\",\n batch_size: int = 2,\n load_in_8_bit: bool = False,\n engine_cfg: EngineConfig | None = None,\n **kwargs: Any,\n) -> OrientationPredictor:\n \"\"\"Page orientation classification architecture.\n\n >>> import numpy as np\n >>> from onnxtr.models import page_orientation_predictor\n >>> model = page_orientation_predictor(arch='mobilenet_v3_small_page_orientation')\n >>> input_page = (255 * np.random.rand(512, 512, 3)).astype(np.uint8)\n >>> out = model([input_page])\n\n Args:\n arch: name of the architecture to use (e.g. 'mobilenet_v3_small_page_orientation')\n batch_size: number of samples the model processes in parallel\n load_in_8_bit: whether to load the the 8-bit quantized model, defaults to False\n engine_cfg: configuration for the inference engine\n **kwargs: keyword arguments to be passed to the OrientationPredictor\n\n Returns:\n OrientationPredictor\n \"\"\"\n model_type = \"page\"\n return _orientation_predictor(\n arch=arch,\n batch_size=batch_size,\n model_type=model_type,\n load_in_8_bit=load_in_8_bit,\n engine_cfg=engine_cfg,\n **kwargs,\n )\n"
},
{
"path": "onnxtr/models/detection/__init__.py",
"content": "from .models import *\nfrom .zoo import *\n"
},
{
"path": "onnxtr/models/detection/_utils/__init__.py",
"content": "from . base import *"
},
{
"path": "onnxtr/models/detection/_utils/base.py",
"content": "# Copyright (C) 2021-2026, Mindee | Felix Dittrich.\n\n# This program is licensed under the Apache License 2.0.\n# See LICENSE or go to for full license details.\n\n\nimport numpy as np\n\n__all__ = [\"_remove_padding\"]\n\n\ndef _remove_padding(\n pages: list[np.ndarray],\n loc_preds: list[np.ndarray],\n preserve_aspect_ratio: bool,\n symmetric_pad: bool,\n assume_straight_pages: bool,\n) -> list[np.ndarray]:\n \"\"\"Remove padding from the localization predictions\n\n Args:\n pages: list of pages\n loc_preds: list of localization predictions\n preserve_aspect_ratio: whether the aspect ratio was preserved during padding\n symmetric_pad: whether the padding was symmetric\n assume_straight_pages: whether the pages are assumed to be straight\n\n Returns:\n list of unpaded localization predictions\n \"\"\"\n if preserve_aspect_ratio:\n # Rectify loc_preds to remove padding\n rectified_preds = []\n for page, loc_pred in zip(pages, loc_preds):\n h, w = page.shape[0], page.shape[1]\n if h > w:\n # y unchanged, dilate x coord\n if symmetric_pad:\n if assume_straight_pages:\n loc_pred[:, [0, 2]] = (loc_pred[:, [0, 2]] - 0.5) * h / w + 0.5\n else:\n loc_pred[:, :, 0] = (loc_pred[:, :, 0] - 0.5) * h / w + 0.5\n else:\n if assume_straight_pages:\n loc_pred[:, [0, 2]] *= h / w\n else:\n loc_pred[:, :, 0] *= h / w\n elif w > h:\n # x unchanged, dilate y coord\n if symmetric_pad:\n if assume_straight_pages:\n loc_pred[:, [1, 3]] = (loc_pred[:, [1, 3]] - 0.5) * w / h + 0.5\n else:\n loc_pred[:, :, 1] = (loc_pred[:, :, 1] - 0.5) * w / h + 0.5\n else:\n if assume_straight_pages:\n loc_pred[:, [1, 3]] *= w / h\n else:\n loc_pred[:, :, 1] *= w / h\n rectified_preds.append(np.clip(loc_pred, 0, 1))\n return rectified_preds\n return loc_preds\n"
},
{
"path": "onnxtr/models/detection/core.py",
"content": "# Copyright (C) 2021-2026, Mindee | Felix Dittrich.\n\n# This program is licensed under the Apache License 2.0.\n# See LICENSE or go to for full license details.\n\n\nimport cv2\nimport numpy as np\n\nfrom onnxtr.utils.repr import NestedObject\n\n__all__ = [\"DetectionPostProcessor\"]\n\n\nclass DetectionPostProcessor(NestedObject):\n \"\"\"Abstract class to postprocess the raw output of the model\n\n Args:\n box_thresh (float): minimal objectness score to consider a box\n bin_thresh (float): threshold to apply to segmentation raw heatmap\n assume straight_pages (bool): if True, fit straight boxes only\n \"\"\"\n\n def __init__(self, box_thresh: float = 0.5, bin_thresh: float = 0.5, assume_straight_pages: bool = True) -> None:\n self.box_thresh = box_thresh\n self.bin_thresh = bin_thresh\n self.assume_straight_pages = assume_straight_pages\n self._opening_kernel: np.ndarray = np.ones((3, 3), dtype=np.uint8)\n\n def extra_repr(self) -> str:\n return f\"bin_thresh={self.bin_thresh}, box_thresh={self.box_thresh}\"\n\n @staticmethod\n def box_score(pred: np.ndarray, points: np.ndarray, assume_straight_pages: bool = True) -> float:\n \"\"\"Compute the confidence score for a polygon : mean of the p values on the polygon\n\n Args:\n pred (np.ndarray): p map returned by the model\n points: coordinates of the polygon\n assume_straight_pages: if True, fit straight boxes only\n\n Returns:\n polygon objectness\n \"\"\"\n h, w = pred.shape[:2]\n\n if assume_straight_pages:\n xmin = np.clip(np.floor(points[:, 0].min()).astype(np.int32), 0, w - 1)\n xmax = np.clip(np.ceil(points[:, 0].max()).astype(np.int32), 0, w - 1)\n ymin = np.clip(np.floor(points[:, 1].min()).astype(np.int32), 0, h - 1)\n ymax = np.clip(np.ceil(points[:, 1].max()).astype(np.int32), 0, h - 1)\n return pred[ymin : ymax + 1, xmin : xmax + 1].mean()\n\n else:\n mask: np.ndarray = np.zeros((h, w), np.int32)\n cv2.fillPoly(mask, [points.astype(np.int32)], 1.0)\n product = pred * mask\n return np.sum(product) / np.count_nonzero(product)\n\n def bitmap_to_boxes(\n self,\n pred: np.ndarray,\n bitmap: np.ndarray,\n ) -> np.ndarray:\n raise NotImplementedError\n\n def __call__(\n self,\n proba_map,\n ) -> list[list[np.ndarray]]:\n \"\"\"Performs postprocessing for a list of model outputs\n\n Args:\n proba_map: probability map of shape (N, H, W, C)\n\n Returns:\n list of N class predictions (for each input sample), where each class predictions is a list of C tensors\n of shape (*, 5) or (*, 6)\n \"\"\"\n if proba_map.ndim != 4:\n raise AssertionError(f\"arg `proba_map` is expected to be 4-dimensional, got {proba_map.ndim}.\")\n\n # Erosion + dilation on the binary map\n bin_map = [\n [\n cv2.morphologyEx(bmap[..., idx], cv2.MORPH_OPEN, self._opening_kernel)\n for idx in range(proba_map.shape[-1])\n ]\n for bmap in (proba_map >= self.bin_thresh).astype(np.uint8)\n ]\n\n return [\n [self.bitmap_to_boxes(pmaps[..., idx], bmaps[idx]) for idx in range(proba_map.shape[-1])]\n for pmaps, bmaps in zip(proba_map, bin_map)\n ]\n"
},
{
"path": "onnxtr/models/detection/models/__init__.py",
"content": "from .fast import *\nfrom .differentiable_binarization import *\nfrom .linknet import *"
},
{
"path": "onnxtr/models/detection/models/differentiable_binarization.py",
"content": "# Copyright (C) 2021-2026, Mindee | Felix Dittrich.\n\n# This program is licensed under the Apache License 2.0.\n# See LICENSE or go to for full license details.\n\nfrom typing import Any\n\nimport numpy as np\nfrom scipy.special import expit\n\nfrom ...engine import Engine, EngineConfig\nfrom ..postprocessor.base import GeneralDetectionPostProcessor\n\n__all__ = [\"DBNet\", \"db_resnet50\", \"db_resnet34\", \"db_mobilenet_v3_large\"]\n\n\ndefault_cfgs: dict[str, dict[str, Any]] = {\n \"db_resnet50\": {\n \"input_shape\": (3, 1024, 1024),\n \"mean\": (0.798, 0.785, 0.772),\n \"std\": (0.264, 0.2749, 0.287),\n \"url\": \"https://github.com/felixdittrich92/OnnxTR/releases/download/v0.0.1/db_resnet50-69ba0015.onnx\",\n \"url_8_bit\": \"https://github.com/felixdittrich92/OnnxTR/releases/download/v0.1.2/db_resnet50_static_8_bit-09a6104f.onnx\",\n },\n \"db_resnet34\": {\n \"input_shape\": (3, 1024, 1024),\n \"mean\": (0.798, 0.785, 0.772),\n \"std\": (0.264, 0.2749, 0.287),\n \"url\": \"https://github.com/felixdittrich92/OnnxTR/releases/download/v0.0.1/db_resnet34-b4873198.onnx\",\n \"url_8_bit\": \"https://github.com/felixdittrich92/OnnxTR/releases/download/v0.1.2/db_resnet34_static_8_bit-027e2c7f.onnx\",\n },\n \"db_mobilenet_v3_large\": {\n \"input_shape\": (3, 1024, 1024),\n \"mean\": (0.798, 0.785, 0.772),\n \"std\": (0.264, 0.2749, 0.287),\n \"url\": \"https://github.com/felixdittrich92/OnnxTR/releases/download/v0.2.0/db_mobilenet_v3_large-4987e7bd.onnx\",\n \"url_8_bit\": \"https://github.com/felixdittrich92/OnnxTR/releases/download/v0.2.0/db_mobilenet_v3_large_static_8_bit-535a6f25.onnx\",\n },\n}\n\n\nclass DBNet(Engine):\n \"\"\"DBNet Onnx loader\n\n Args:\n model_path: path or url to onnx model file\n engine_cfg: configuration for the inference engine\n bin_thresh: threshold for binarization of the output feature map\n box_thresh: minimal objectness score to consider a box\n assume_straight_pages: if True, fit straight bounding boxes only\n cfg: the configuration dict of the model\n **kwargs: additional arguments to be passed to `Engine`\n \"\"\"\n\n def __init__(\n self,\n model_path: str,\n engine_cfg: EngineConfig | None = None,\n bin_thresh: float = 0.3,\n box_thresh: float = 0.1,\n assume_straight_pages: bool = True,\n cfg: dict[str, Any] | None = None,\n **kwargs: Any,\n ) -> None:\n super().__init__(url=model_path, engine_cfg=engine_cfg, **kwargs)\n\n self.cfg = cfg\n self.assume_straight_pages = assume_straight_pages\n\n self.postprocessor = GeneralDetectionPostProcessor(\n assume_straight_pages=self.assume_straight_pages, bin_thresh=bin_thresh, box_thresh=box_thresh\n )\n\n def __call__(\n self,\n x: np.ndarray,\n return_model_output: bool = False,\n **kwargs: Any,\n ) -> dict[str, Any]:\n logits = self.run(x)\n\n out: dict[str, Any] = {}\n\n prob_map = expit(logits)\n if return_model_output:\n out[\"out_map\"] = prob_map\n\n out[\"preds\"] = self.postprocessor(prob_map)\n\n return out\n\n\ndef _dbnet(\n arch: str,\n model_path: str,\n load_in_8_bit: bool = False,\n engine_cfg: EngineConfig | None = None,\n **kwargs: Any,\n) -> DBNet:\n # Patch the url\n model_path = default_cfgs[arch][\"url_8_bit\"] if load_in_8_bit and \"http\" in model_path else model_path\n # Build the model\n return DBNet(model_path, cfg=default_cfgs[arch], engine_cfg=engine_cfg, **kwargs)\n\n\ndef db_resnet34(\n model_path: str = default_cfgs[\"db_resnet34\"][\"url\"],\n load_in_8_bit: bool = False,\n engine_cfg: EngineConfig | None = None,\n **kwargs: Any,\n) -> DBNet:\n \"\"\"DBNet as described in `\"Real-time Scene Text Detection with Differentiable Binarization\"\n `_, using a ResNet-34 backbone.\n\n >>> import numpy as np\n >>> from onnxtr.models import db_resnet34\n >>> model = db_resnet34()\n >>> input_tensor = np.random.rand(1, 3, 1024, 1024)\n >>> out = model(input_tensor)\n\n Args:\n model_path: path to onnx model file, defaults to url in default_cfgs\n load_in_8_bit: whether to load the the 8-bit quantized model, defaults to False\n engine_cfg: configuration for the inference engine\n **kwargs: keyword arguments of the DBNet architecture\n\n Returns:\n text detection architecture\n \"\"\"\n return _dbnet(\"db_resnet34\", model_path, load_in_8_bit, engine_cfg, **kwargs)\n\n\ndef db_resnet50(\n model_path: str = default_cfgs[\"db_resnet50\"][\"url\"],\n load_in_8_bit: bool = False,\n engine_cfg: EngineConfig | None = None,\n **kwargs: Any,\n) -> DBNet:\n \"\"\"DBNet as described in `\"Real-time Scene Text Detection with Differentiable Binarization\"\n `_, using a ResNet-50 backbone.\n\n >>> import numpy as np\n >>> from onnxtr.models import db_resnet50\n >>> model = db_resnet50()\n >>> input_tensor = np.random.rand(1, 3, 1024, 1024)\n >>> out = model(input_tensor)\n\n Args:\n model_path: path to onnx model file, defaults to url in default_cfgs\n load_in_8_bit: whether to load the the 8-bit quantized model, defaults to False\n engine_cfg: configuration for the inference engine\n **kwargs: keyword arguments of the DBNet architecture\n\n Returns:\n text detection architecture\n \"\"\"\n return _dbnet(\"db_resnet50\", model_path, load_in_8_bit, engine_cfg, **kwargs)\n\n\ndef db_mobilenet_v3_large(\n model_path: str = default_cfgs[\"db_mobilenet_v3_large\"][\"url\"],\n load_in_8_bit: bool = False,\n engine_cfg: EngineConfig | None = None,\n **kwargs: Any,\n) -> DBNet:\n \"\"\"DBNet as described in `\"Real-time Scene Text Detection with Differentiable Binarization\"\n `_, using a MobileNet V3 Large backbone.\n\n >>> import numpy as np\n >>> from onnxtr.models import db_mobilenet_v3_large\n >>> model = db_mobilenet_v3_large()\n >>> input_tensor = np.random.rand(1, 3, 1024, 1024)\n >>> out = model(input_tensor)\n\n Args:\n model_path: path to onnx model file, defaults to url in default_cfgs\n load_in_8_bit: whether to load the the 8-bit quantized model, defaults to False\n engine_cfg: configuration for the inference engine\n **kwargs: keyword arguments of the DBNet architecture\n\n Returns:\n text detection architecture\n \"\"\"\n return _dbnet(\"db_mobilenet_v3_large\", model_path, load_in_8_bit, engine_cfg, **kwargs)\n"
},
{
"path": "onnxtr/models/detection/models/fast.py",
"content": "# Copyright (C) 2021-2026, Mindee | Felix Dittrich.\n\n# This program is licensed under the Apache License 2.0.\n# See LICENSE or go to for full license details.\n\nimport logging\nfrom typing import Any\n\nimport numpy as np\nfrom scipy.special import expit\n\nfrom ...engine import Engine, EngineConfig\nfrom ..postprocessor.base import GeneralDetectionPostProcessor\n\n__all__ = [\"FAST\", \"fast_tiny\", \"fast_small\", \"fast_base\"]\n\n\ndefault_cfgs: dict[str, dict[str, Any]] = {\n \"fast_tiny\": {\n \"input_shape\": (3, 1024, 1024),\n \"mean\": (0.798, 0.785, 0.772),\n \"std\": (0.264, 0.2749, 0.287),\n \"url\": \"https://github.com/felixdittrich92/OnnxTR/releases/download/v0.0.1/rep_fast_tiny-28867779.onnx\",\n },\n \"fast_small\": {\n \"input_shape\": (3, 1024, 1024),\n \"mean\": (0.798, 0.785, 0.772),\n \"std\": (0.264, 0.2749, 0.287),\n \"url\": \"https://github.com/felixdittrich92/OnnxTR/releases/download/v0.0.1/rep_fast_small-10428b70.onnx\",\n },\n \"fast_base\": {\n \"input_shape\": (3, 1024, 1024),\n \"mean\": (0.798, 0.785, 0.772),\n \"std\": (0.264, 0.2749, 0.287),\n \"url\": \"https://github.com/felixdittrich92/OnnxTR/releases/download/v0.0.1/rep_fast_base-1b89ebf9.onnx\",\n },\n}\n\n\nclass FAST(Engine):\n \"\"\"FAST Onnx loader\n\n Args:\n model_path: path or url to onnx model file\n engine_cfg: configuration for the inference engine\n bin_thresh: threshold for binarization of the output feature map\n box_thresh: minimal objectness score to consider a box\n assume_straight_pages: if True, fit straight bounding boxes only\n cfg: the configuration dict of the model\n **kwargs: additional arguments to be passed to `Engine`\n \"\"\"\n\n def __init__(\n self,\n model_path: str,\n engine_cfg: EngineConfig | None = None,\n bin_thresh: float = 0.1,\n box_thresh: float = 0.1,\n assume_straight_pages: bool = True,\n cfg: dict[str, Any] | None = None,\n **kwargs: Any,\n ) -> None:\n super().__init__(url=model_path, engine_cfg=engine_cfg, **kwargs)\n\n self.cfg = cfg\n self.assume_straight_pages = assume_straight_pages\n\n self.postprocessor = GeneralDetectionPostProcessor(\n assume_straight_pages=self.assume_straight_pages, bin_thresh=bin_thresh, box_thresh=box_thresh\n )\n\n def __call__(\n self,\n x: np.ndarray,\n return_model_output: bool = False,\n **kwargs: Any,\n ) -> dict[str, Any]:\n logits = self.run(x)\n\n out: dict[str, Any] = {}\n\n prob_map = expit(logits)\n if return_model_output:\n out[\"out_map\"] = prob_map\n\n out[\"preds\"] = self.postprocessor(prob_map)\n\n return out\n\n\ndef _fast(\n arch: str,\n model_path: str,\n load_in_8_bit: bool = False,\n engine_cfg: EngineConfig | None = None,\n **kwargs: Any,\n) -> FAST:\n if load_in_8_bit:\n logging.warning(\"FAST models do not support 8-bit quantization yet. Loading full precision model...\")\n # Build the model\n return FAST(model_path, cfg=default_cfgs[arch], engine_cfg=engine_cfg, **kwargs)\n\n\ndef fast_tiny(\n model_path: str = default_cfgs[\"fast_tiny\"][\"url\"],\n load_in_8_bit: bool = False,\n engine_cfg: EngineConfig | None = None,\n **kwargs: Any,\n) -> FAST:\n \"\"\"FAST as described in `\"FAST: Faster Arbitrarily-Shaped Text Detector with Minimalist Kernel Representation\"\n `_, using a tiny TextNet backbone.\n\n >>> import numpy as np\n >>> from onnxtr.models import fast_tiny\n >>> model = fast_tiny()\n >>> input_tensor = np.random.rand(1, 3, 1024, 1024)\n >>> out = model(input_tensor)\n\n Args:\n model_path: path to onnx model file, defaults to url in default_cfgs\n load_in_8_bit: whether to load the the 8-bit quantized model, defaults to False\n engine_cfg: configuration for the inference engine\n **kwargs: keyword arguments of the DBNet architecture\n\n Returns:\n text detection architecture\n \"\"\"\n return _fast(\"fast_tiny\", model_path, load_in_8_bit, engine_cfg, **kwargs)\n\n\ndef fast_small(\n model_path: str = default_cfgs[\"fast_small\"][\"url\"],\n load_in_8_bit: bool = False,\n engine_cfg: EngineConfig | None = None,\n **kwargs: Any,\n) -> FAST:\n \"\"\"FAST as described in `\"FAST: Faster Arbitrarily-Shaped Text Detector with Minimalist Kernel Representation\"\n `_, using a small TextNet backbone.\n\n >>> import numpy as np\n >>> from onnxtr.models import fast_small\n >>> model = fast_small()\n >>> input_tensor = np.random.rand(1, 3, 1024, 1024)\n >>> out = model(input_tensor)\n\n Args:\n model_path: path to onnx model file, defaults to url in default_cfgs\n load_in_8_bit: whether to load the the 8-bit quantized model, defaults to False\n engine_cfg: configuration for the inference engine\n **kwargs: keyword arguments of the DBNet architecture\n\n Returns:\n text detection architecture\n \"\"\"\n return _fast(\"fast_small\", model_path, load_in_8_bit, engine_cfg, **kwargs)\n\n\ndef fast_base(\n model_path: str = default_cfgs[\"fast_base\"][\"url\"],\n load_in_8_bit: bool = False,\n engine_cfg: EngineConfig | None = None,\n **kwargs: Any,\n) -> FAST:\n \"\"\"FAST as described in `\"FAST: Faster Arbitrarily-Shaped Text Detector with Minimalist Kernel Representation\"\n `_, using a base TextNet backbone.\n\n >>> import numpy as np\n >>> from onnxtr.models import fast_base\n >>> model = fast_base()\n >>> input_tensor = np.random.rand(1, 3, 1024, 1024)\n >>> out = model(input_tensor)\n\n Args:\n model_path: path to onnx model file, defaults to url in default_cfgs\n load_in_8_bit: whether to load the the 8-bit quantized model, defaults to False\n engine_cfg: configuration for the inference engine\n **kwargs: keyword arguments of the DBNet architecture\n\n Returns:\n text detection architecture\n \"\"\"\n return _fast(\"fast_base\", model_path, load_in_8_bit, engine_cfg, **kwargs)\n"
},
{
"path": "onnxtr/models/detection/models/linknet.py",
"content": "# Copyright (C) 2021-2026, Mindee | Felix Dittrich.\n\n# This program is licensed under the Apache License 2.0.\n# See LICENSE or go to for full license details.\n\nfrom typing import Any\n\nimport numpy as np\nfrom scipy.special import expit\n\nfrom ...engine import Engine, EngineConfig\nfrom ..postprocessor.base import GeneralDetectionPostProcessor\n\n__all__ = [\"LinkNet\", \"linknet_resnet18\", \"linknet_resnet34\", \"linknet_resnet50\"]\n\n\ndefault_cfgs: dict[str, dict[str, Any]] = {\n \"linknet_resnet18\": {\n \"input_shape\": (3, 1024, 1024),\n \"mean\": (0.798, 0.785, 0.772),\n \"std\": (0.264, 0.2749, 0.287),\n \"url\": \"https://github.com/felixdittrich92/OnnxTR/releases/download/v0.0.1/linknet_resnet18-e0e0b9dc.onnx\",\n \"url_8_bit\": \"https://github.com/felixdittrich92/OnnxTR/releases/download/v0.1.2/linknet_resnet18_static_8_bit-3b3a37dd.onnx\",\n },\n \"linknet_resnet34\": {\n \"input_shape\": (3, 1024, 1024),\n \"mean\": (0.798, 0.785, 0.772),\n \"std\": (0.264, 0.2749, 0.287),\n \"url\": \"https://github.com/felixdittrich92/OnnxTR/releases/download/v0.0.1/linknet_resnet34-93e39a39.onnx\",\n \"url_8_bit\": \"https://github.com/felixdittrich92/OnnxTR/releases/download/v0.1.2/linknet_resnet34_static_8_bit-2824329d.onnx\",\n },\n \"linknet_resnet50\": {\n \"input_shape\": (3, 1024, 1024),\n \"mean\": (0.798, 0.785, 0.772),\n \"std\": (0.264, 0.2749, 0.287),\n \"url\": \"https://github.com/felixdittrich92/OnnxTR/releases/download/v0.0.1/linknet_resnet50-15d8c4ec.onnx\",\n \"url_8_bit\": \"https://github.com/felixdittrich92/OnnxTR/releases/download/v0.1.2/linknet_resnet50_static_8_bit-65d6b0b8.onnx\",\n },\n}\n\n\nclass LinkNet(Engine):\n \"\"\"LinkNet Onnx loader\n\n Args:\n model_path: path or url to onnx model file\n engine_cfg: configuration for the inference engine\n bin_thresh: threshold for binarization of the output feature map\n box_thresh: minimal objectness score to consider a box\n assume_straight_pages: if True, fit straight bounding boxes only\n cfg: the configuration dict of the model\n **kwargs: additional arguments to be passed to `Engine`\n \"\"\"\n\n def __init__(\n self,\n model_path: str,\n engine_cfg: EngineConfig | None = None,\n bin_thresh: float = 0.1,\n box_thresh: float = 0.1,\n assume_straight_pages: bool = True,\n cfg: dict[str, Any] | None = None,\n **kwargs: Any,\n ) -> None:\n super().__init__(url=model_path, engine_cfg=engine_cfg, **kwargs)\n\n self.cfg = cfg\n self.assume_straight_pages = assume_straight_pages\n\n self.postprocessor = GeneralDetectionPostProcessor(\n assume_straight_pages=self.assume_straight_pages, bin_thresh=bin_thresh, box_thresh=box_thresh\n )\n\n def __call__(\n self,\n x: np.ndarray,\n return_model_output: bool = False,\n **kwargs: Any,\n ) -> dict[str, Any]:\n logits = self.run(x)\n\n out: dict[str, Any] = {}\n\n prob_map = expit(logits)\n if return_model_output:\n out[\"out_map\"] = prob_map\n\n out[\"preds\"] = self.postprocessor(prob_map)\n\n return out\n\n\ndef _linknet(\n arch: str,\n model_path: str,\n load_in_8_bit: bool = False,\n engine_cfg: EngineConfig | None = None,\n **kwargs: Any,\n) -> LinkNet:\n # Patch the url\n model_path = default_cfgs[arch][\"url_8_bit\"] if load_in_8_bit and \"http\" in model_path else model_path\n # Build the model\n return LinkNet(model_path, cfg=default_cfgs[arch], engine_cfg=engine_cfg, **kwargs)\n\n\ndef linknet_resnet18(\n model_path: str = default_cfgs[\"linknet_resnet18\"][\"url\"],\n load_in_8_bit: bool = False,\n engine_cfg: EngineConfig | None = None,\n **kwargs: Any,\n) -> LinkNet:\n \"\"\"LinkNet as described in `\"LinkNet: Exploiting Encoder Representations for Efficient Semantic Segmentation\"\n `_.\n\n >>> import numpy as np\n >>> from onnxtr.models import linknet_resnet18\n >>> model = linknet_resnet18()\n >>> input_tensor = np.random.rand(1, 3, 1024, 1024)\n >>> out = model(input_tensor)\n\n Args:\n model_path: path to onnx model file, defaults to url in default_cfgs\n load_in_8_bit: whether to load the the 8-bit quantized model, defaults to False\n engine_cfg: configuration for the inference engine\n **kwargs: keyword arguments of the LinkNet architecture\n\n Returns:\n text detection architecture\n \"\"\"\n return _linknet(\"linknet_resnet18\", model_path, load_in_8_bit, engine_cfg, **kwargs)\n\n\ndef linknet_resnet34(\n model_path: str = default_cfgs[\"linknet_resnet34\"][\"url\"],\n load_in_8_bit: bool = False,\n engine_cfg: EngineConfig | None = None,\n **kwargs: Any,\n) -> LinkNet:\n \"\"\"LinkNet as described in `\"LinkNet: Exploiting Encoder Representations for Efficient Semantic Segmentation\"\n `_.\n\n >>> import numpy as np\n >>> from onnxtr.models import linknet_resnet34\n >>> model = linknet_resnet34()\n >>> input_tensor = np.random.rand(1, 3, 1024, 1024)\n >>> out = model(input_tensor)\n\n Args:\n model_path: path to onnx model file, defaults to url in default_cfgs\n load_in_8_bit: whether to load the the 8-bit quantized model, defaults to False\n engine_cfg: configuration for the inference engine\n **kwargs: keyword arguments of the LinkNet architecture\n\n Returns:\n text detection architecture\n \"\"\"\n return _linknet(\"linknet_resnet34\", model_path, load_in_8_bit, engine_cfg, **kwargs)\n\n\ndef linknet_resnet50(\n model_path: str = default_cfgs[\"linknet_resnet50\"][\"url\"],\n load_in_8_bit: bool = False,\n engine_cfg: EngineConfig | None = None,\n **kwargs: Any,\n) -> LinkNet:\n \"\"\"LinkNet as described in `\"LinkNet: Exploiting Encoder Representations for Efficient Semantic Segmentation\"\n `_.\n\n >>> import numpy as np\n >>> from onnxtr.models import linknet_resnet50\n >>> model = linknet_resnet50()\n >>> input_tensor = np.random.rand(1, 3, 1024, 1024)\n >>> out = model(input_tensor)\n\n Args:\n model_path: path to onnx model file, defaults to url in default_cfgs\n load_in_8_bit: whether to load the the 8-bit quantized model, defaults to False\n engine_cfg: configuration for the inference engine\n **kwargs: keyword arguments of the LinkNet architecture\n\n Returns:\n text detection architecture\n \"\"\"\n return _linknet(\"linknet_resnet50\", model_path, load_in_8_bit, engine_cfg, **kwargs)\n"
},
{
"path": "onnxtr/models/detection/postprocessor/__init__.py",
"content": ""
},
{
"path": "onnxtr/models/detection/postprocessor/base.py",
"content": "# Copyright (C) 2021-2026, Mindee | Felix Dittrich.\n\n# This program is licensed under the Apache License 2.0.\n# See LICENSE or go to for full license details.\n\n# Credits: post-processing adapted from https://github.com/xuannianz/DifferentiableBinarization\n\n\nimport cv2\nimport numpy as np\nimport pyclipper\n\nfrom onnxtr.utils import order_points\n\nfrom ..core import DetectionPostProcessor\n\n__all__ = [\"GeneralDetectionPostProcessor\"]\n\n\nclass GeneralDetectionPostProcessor(DetectionPostProcessor):\n \"\"\"Implements a post processor for FAST model.\n\n Args:\n bin_thresh: threshold used to binzarized p_map at inference time\n box_thresh: minimal objectness score to consider a box\n assume_straight_pages: whether the inputs were expected to have horizontal text elements\n \"\"\"\n\n def __init__(\n self,\n bin_thresh: float = 0.1,\n box_thresh: float = 0.1,\n assume_straight_pages: bool = True,\n ) -> None:\n super().__init__(box_thresh, bin_thresh, assume_straight_pages)\n self.unclip_ratio = 1.5\n\n def polygon_to_box(\n self,\n points: np.ndarray,\n ) -> np.ndarray:\n \"\"\"Expand a polygon (points) by a factor unclip_ratio, and returns a polygon\n\n Args:\n points: The first parameter.\n\n Returns:\n a box in absolute coordinates (xmin, ymin, xmax, ymax) or (4, 2) array (quadrangle)\n \"\"\"\n if not self.assume_straight_pages:\n # Compute the rectangle polygon enclosing the raw polygon\n rect = cv2.minAreaRect(points)\n points = cv2.boxPoints(rect)\n # Add 1 pixel to correct cv2 approx\n area = (rect[1][0] + 1) * (1 + rect[1][1])\n length = 2 * (rect[1][0] + rect[1][1]) + 2\n else:\n area = cv2.contourArea(points)\n length = cv2.arcLength(points, closed=True)\n distance = area * self.unclip_ratio / length # compute distance to expand polygon\n offset = pyclipper.PyclipperOffset()\n offset.AddPath(points, pyclipper.JT_ROUND, pyclipper.ET_CLOSEDPOLYGON)\n _points = offset.Execute(distance)\n # Take biggest stack of points\n idx = 0\n if len(_points) > 1:\n max_size = 0\n for _idx, p in enumerate(_points):\n if len(p) > max_size:\n idx = _idx\n max_size = len(p)\n # We ensure that _points can be correctly casted to a ndarray\n _points = [_points[idx]]\n expanded_points: np.ndarray = np.asarray(_points) # expand polygon\n if len(expanded_points) < 1:\n return None # type: ignore[return-value]\n return (\n cv2.boundingRect(expanded_points) # type: ignore[return-value]\n if self.assume_straight_pages\n else order_points(cv2.boxPoints(cv2.minAreaRect(expanded_points)))\n )\n\n def bitmap_to_boxes(\n self,\n pred: np.ndarray,\n bitmap: np.ndarray,\n ) -> np.ndarray:\n \"\"\"Compute boxes from a bitmap/pred_map: find connected components then filter boxes\n\n Args:\n pred: Pred map from differentiable linknet output\n bitmap: Bitmap map computed from pred (binarized)\n angle_tol: Comparison tolerance of the angle with the median angle across the page\n ratio_tol: Under this limit aspect ratio, we cannot resolve the direction of the crop\n\n Returns:\n np tensor boxes for the bitmap, each box is a 6-element list\n containing x, y, w, h, alpha, score for the box\n \"\"\"\n height, width = bitmap.shape[:2]\n boxes: list[np.ndarray | list[float]] = []\n # get contours from connected components on the bitmap\n contours, _ = cv2.findContours(bitmap.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)\n for contour in contours:\n # Check whether smallest enclosing bounding box is not too small\n if np.any(contour[:, 0].max(axis=0) - contour[:, 0].min(axis=0) < 2):\n continue\n # Compute objectness\n if self.assume_straight_pages:\n x, y, w, h = cv2.boundingRect(contour)\n points: np.ndarray = np.array([[x, y], [x, y + h], [x + w, y + h], [x + w, y]])\n score = self.box_score(pred, points, assume_straight_pages=True)\n else:\n score = self.box_score(pred, contour, assume_straight_pages=False)\n\n if score < self.box_thresh: # remove polygons with a weak objectness\n continue\n\n if self.assume_straight_pages:\n _box = self.polygon_to_box(points)\n else:\n _box = self.polygon_to_box(np.squeeze(contour))\n\n if self.assume_straight_pages:\n # compute relative polygon to get rid of img shape\n x, y, w, h = _box\n xmin, ymin, xmax, ymax = x / width, y / height, (x + w) / width, (y + h) / height\n boxes.append([xmin, ymin, xmax, ymax, score])\n else:\n # compute relative box to get rid of img shape\n _box[:, 0] /= width\n _box[:, 1] /= height\n # Add score to box as (0, score)\n boxes.append(np.vstack([_box, np.array([0.0, score])]))\n\n if not self.assume_straight_pages:\n return np.clip(np.asarray(boxes), 0, 1) if len(boxes) > 0 else np.zeros((0, 5, 2), dtype=pred.dtype)\n else:\n return np.clip(np.asarray(boxes), 0, 1) if len(boxes) > 0 else np.zeros((0, 5), dtype=pred.dtype)\n"
},
{
"path": "onnxtr/models/detection/predictor/__init__.py",
"content": "from .base import *\n"
},
{
"path": "onnxtr/models/detection/predictor/base.py",
"content": "# Copyright (C) 2021-2026, Mindee | Felix Dittrich.\n\n# This program is licensed under the Apache License 2.0.\n# See LICENSE or go to for full license details.\n\nfrom typing import Any\n\nimport numpy as np\n\nfrom onnxtr.models.detection._utils import _remove_padding\nfrom onnxtr.models.preprocessor import PreProcessor\nfrom onnxtr.utils.repr import NestedObject\n\n__all__ = [\"DetectionPredictor\"]\n\n\nclass DetectionPredictor(NestedObject):\n \"\"\"Implements an object able to localize text elements in a document\n\n Args:\n pre_processor: transform inputs for easier batched model inference\n model: core detection architecture\n \"\"\"\n\n _children_names: list[str] = [\"pre_processor\", \"model\"]\n\n def __init__(\n self,\n pre_processor: PreProcessor,\n model: Any,\n ) -> None:\n self.pre_processor = pre_processor\n self.model = model\n\n def __call__(\n self,\n pages: list[np.ndarray],\n return_maps: bool = False,\n **kwargs: Any,\n ) -> list[np.ndarray] | tuple[list[np.ndarray], list[np.ndarray]]:\n # Extract parameters from the preprocessor\n preserve_aspect_ratio = self.pre_processor.resize.preserve_aspect_ratio\n symmetric_pad = self.pre_processor.resize.symmetric_pad\n assume_straight_pages = self.model.assume_straight_pages\n\n # Dimension check\n if any(page.ndim != 3 for page in pages):\n raise ValueError(\"incorrect input shape: all pages are expected to be multi-channel 2D images.\")\n\n processed_batches = self.pre_processor(pages)\n predicted_batches = [\n self.model(batch, return_preds=True, return_model_output=True, **kwargs) for batch in processed_batches\n ]\n\n # Remove padding from loc predictions\n preds = _remove_padding(\n pages,\n [pred[0] for batch in predicted_batches for pred in batch[\"preds\"]],\n preserve_aspect_ratio=preserve_aspect_ratio,\n symmetric_pad=symmetric_pad,\n assume_straight_pages=assume_straight_pages,\n )\n\n if return_maps:\n seg_maps = [pred for batch in predicted_batches for pred in batch[\"out_map\"]]\n return preds, seg_maps\n return preds\n"
},
{
"path": "onnxtr/models/detection/zoo.py",
"content": "# Copyright (C) 2021-2026, Mindee | Felix Dittrich.\n\n# This program is licensed under the Apache License 2.0.\n# See LICENSE or go to for full license details.\n\nfrom typing import Any\n\nfrom .. import detection\nfrom ..engine import EngineConfig\nfrom ..preprocessor import PreProcessor\nfrom .predictor import DetectionPredictor\n\n__all__ = [\"detection_predictor\"]\n\nARCHS = [\n \"db_resnet34\",\n \"db_resnet50\",\n \"db_mobilenet_v3_large\",\n \"linknet_resnet18\",\n \"linknet_resnet34\",\n \"linknet_resnet50\",\n \"fast_tiny\",\n \"fast_small\",\n \"fast_base\",\n]\n\n\ndef _predictor(\n arch: Any,\n assume_straight_pages: bool = True,\n load_in_8_bit: bool = False,\n engine_cfg: EngineConfig | None = None,\n **kwargs: Any,\n) -> DetectionPredictor:\n if isinstance(arch, str):\n if arch not in ARCHS:\n raise ValueError(f\"unknown architecture '{arch}'\")\n\n _model = detection.__dict__[arch](\n assume_straight_pages=assume_straight_pages, load_in_8_bit=load_in_8_bit, engine_cfg=engine_cfg\n )\n else:\n if not isinstance(arch, (detection.DBNet, detection.LinkNet, detection.FAST)):\n raise ValueError(f\"unknown architecture: {type(arch)}\")\n\n _model = arch\n _model.assume_straight_pages = assume_straight_pages\n _model.postprocessor.assume_straight_pages = assume_straight_pages\n\n kwargs[\"mean\"] = kwargs.get(\"mean\", _model.cfg[\"mean\"])\n kwargs[\"std\"] = kwargs.get(\"std\", _model.cfg[\"std\"])\n kwargs[\"batch_size\"] = kwargs.get(\"batch_size\", 2)\n predictor = DetectionPredictor(\n PreProcessor(_model.cfg[\"input_shape\"][1:], **kwargs),\n _model,\n )\n return predictor\n\n\ndef detection_predictor(\n arch: Any = \"fast_base\",\n assume_straight_pages: bool = True,\n preserve_aspect_ratio: bool = True,\n symmetric_pad: bool = True,\n batch_size: int = 2,\n load_in_8_bit: bool = False,\n engine_cfg: EngineConfig | None = None,\n **kwargs: Any,\n) -> DetectionPredictor:\n \"\"\"Text detection architecture.\n\n >>> import numpy as np\n >>> from onnxtr.models import detection_predictor\n >>> model = detection_predictor(arch='db_resnet50')\n >>> input_page = (255 * np.random.rand(600, 800, 3)).astype(np.uint8)\n >>> out = model([input_page])\n\n Args:\n arch: name of the architecture or model itself to use (e.g. 'db_resnet50')\n assume_straight_pages: If True, fit straight boxes to the page\n preserve_aspect_ratio: If True, pad the input document image to preserve the aspect ratio before\n running the detection model on it\n symmetric_pad: if True, pad the image symmetrically instead of padding at the bottom-right\n batch_size: number of samples the model processes in parallel\n load_in_8_bit: whether to load the the 8-bit quantized model, defaults to False\n engine_cfg: configuration for the inference engine\n **kwargs: optional keyword arguments passed to the architecture\n\n Returns:\n Detection predictor\n \"\"\"\n return _predictor(\n arch=arch,\n assume_straight_pages=assume_straight_pages,\n preserve_aspect_ratio=preserve_aspect_ratio,\n symmetric_pad=symmetric_pad,\n batch_size=batch_size,\n load_in_8_bit=load_in_8_bit,\n engine_cfg=engine_cfg,\n **kwargs,\n )\n"
},
{
"path": "onnxtr/models/engine.py",
"content": "# Copyright (C) 2021-2026, Mindee | Felix Dittrich.\n\n# This program is licensed under the Apache License 2.0.\n# See LICENSE or go to for full license details.\n\nimport logging\nimport os\nfrom collections.abc import Callable\nfrom typing import Any, TypeAlias\n\nimport numpy as np\nfrom onnxruntime import (\n ExecutionMode,\n GraphOptimizationLevel,\n InferenceSession,\n RunOptions,\n SessionOptions,\n get_available_providers,\n get_device,\n)\nfrom onnxruntime.capi._pybind_state import set_default_logger_severity\n\nset_default_logger_severity(int(os.getenv(\"ORT_LOG_SEVERITY_LEVEL\", 4)))\n\nfrom onnxtr.utils.data import download_from_url\nfrom onnxtr.utils.geometry import shape_translate\n\n__all__ = [\"EngineConfig\", \"RunOptionsProvider\"]\n\nRunOptionsProvider: TypeAlias = Callable[[RunOptions], RunOptions]\n\n\nclass EngineConfig:\n \"\"\"Implements a configuration class for the engine of a model\n\n Args:\n providers: list of providers to use for inference ref.: https://onnxruntime.ai/docs/execution-providers/\n session_options: configuration for the inference session ref.: https://onnxruntime.ai/docs/api/python/api_summary.html#sessionoptions\n \"\"\"\n\n def __init__(\n self,\n providers: list[tuple[str, dict[str, Any]]] | list[str] | None = None,\n session_options: SessionOptions | None = None,\n run_options_provider: RunOptionsProvider | None = None,\n ):\n self._providers = providers or self._init_providers()\n self._session_options = session_options or self._init_sess_opts()\n self.run_options_provider = run_options_provider\n\n def _init_providers(self) -> list[tuple[str, dict[str, Any]]]:\n providers: Any = [(\"CPUExecutionProvider\", {\"arena_extend_strategy\": \"kSameAsRequested\"})]\n available_providers = get_available_providers()\n logging.info(f\"Available providers: {available_providers}\")\n if \"CUDAExecutionProvider\" in available_providers and get_device() == \"GPU\": # pragma: no cover\n providers.insert(\n 0,\n (\n \"CUDAExecutionProvider\",\n {\n \"device_id\": 0,\n \"arena_extend_strategy\": \"kNextPowerOfTwo\",\n \"cudnn_conv_algo_search\": \"DEFAULT\",\n \"do_copy_in_default_stream\": True,\n },\n ),\n )\n elif \"CoreMLExecutionProvider\" in available_providers: # pragma: no cover\n providers.insert(0, (\"CoreMLExecutionProvider\", {}))\n return providers\n\n def _init_sess_opts(self) -> SessionOptions:\n session_options = SessionOptions()\n session_options.enable_cpu_mem_arena = True\n session_options.execution_mode = ExecutionMode.ORT_SEQUENTIAL\n session_options.graph_optimization_level = GraphOptimizationLevel.ORT_ENABLE_ALL\n session_options.intra_op_num_threads = -1\n session_options.inter_op_num_threads = -1\n return session_options\n\n @property\n def providers(self) -> list[tuple[str, dict[str, Any]]] | list[str]:\n return self._providers\n\n @property\n def session_options(self) -> SessionOptions:\n return self._session_options\n\n def __repr__(self) -> str:\n return f\"EngineConfig(providers={self.providers})\"\n\n\nclass Engine:\n \"\"\"Implements an abstract class for the engine of a model\n\n Args:\n url: the url to use to download a model if needed\n engine_cfg: the configuration of the engine\n **kwargs: additional arguments to be passed to `download_from_url`\n \"\"\"\n\n def __init__(self, url: str, engine_cfg: EngineConfig | None = None, **kwargs: Any) -> None:\n engine_cfg = engine_cfg if isinstance(engine_cfg, EngineConfig) else EngineConfig()\n archive_path = download_from_url(url, cache_subdir=\"models\", **kwargs) if \"http\" in url else url\n # NOTE: older onnxruntime versions require a string path for windows\n archive_path = rf\"{archive_path}\"\n # Store model path for each model\n self.model_path = archive_path\n self.session_options = engine_cfg.session_options\n self.providers = engine_cfg.providers\n self.run_options_provider = engine_cfg.run_options_provider\n self.runtime = InferenceSession(archive_path, providers=self.providers, sess_options=self.session_options)\n self.runtime_inputs = self.runtime.get_inputs()[0]\n self.tf_exported = int(self.runtime_inputs.shape[-1]) == 3\n self.fixed_batch_size: int | str = self.runtime_inputs.shape[\n 0\n ] # mostly possible with tensorflow exported models\n self.output_name = [output.name for output in self.runtime.get_outputs()]\n\n def run(self, inputs: np.ndarray) -> np.ndarray:\n run_options = RunOptions()\n if self.run_options_provider is not None:\n run_options = self.run_options_provider(run_options)\n if self.tf_exported:\n inputs = shape_translate(inputs, format=\"BHWC\") # sanity check\n else:\n inputs = shape_translate(inputs, format=\"BCHW\")\n if isinstance(self.fixed_batch_size, int) and self.fixed_batch_size != 0: # dynamic batch size is a string\n inputs = np.broadcast_to(inputs, (self.fixed_batch_size, *inputs.shape))\n # combine the results\n logits = np.concatenate(\n [\n self.runtime.run(self.output_name, {self.runtime_inputs.name: batch}, run_options=run_options)[0]\n for batch in inputs\n ],\n axis=0,\n )\n else:\n logits = self.runtime.run(self.output_name, {self.runtime_inputs.name: inputs}, run_options=run_options)[0]\n return shape_translate(logits, format=\"BHWC\")\n"
},
{
"path": "onnxtr/models/factory/__init__.py",
"content": "from .hub import *\n"
},
{
"path": "onnxtr/models/factory/hub.py",
"content": "# Copyright (C) 2021-2026, Mindee | Felix Dittrich.\n\n# This program is licensed under the Apache License 2.0.\n# See LICENSE or go to for full license details.\n\n# Inspired by: https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/hub.py\n\nimport json\nimport logging\nimport shutil\nimport subprocess\nimport tempfile\nimport textwrap\nfrom pathlib import Path\nfrom typing import Any\n\nfrom huggingface_hub import (\n HfApi,\n get_token,\n hf_hub_download,\n login,\n)\n\nfrom onnxtr import models\nfrom onnxtr.models.engine import EngineConfig\n\n__all__ = [\"login_to_hub\", \"push_to_hf_hub\", \"from_hub\", \"_save_model_and_config_for_hf_hub\"]\n\n\nAVAILABLE_ARCHS = {\n \"classification\": models.classification.zoo.ORIENTATION_ARCHS,\n \"detection\": models.detection.zoo.ARCHS,\n \"recognition\": models.recognition.zoo.ARCHS,\n}\n\n\ndef login_to_hub() -> None: # pragma: no cover\n \"\"\"Login to huggingface hub\"\"\"\n access_token = get_token()\n if access_token is not None:\n logging.info(\"Huggingface Hub token found and valid\")\n login(token=access_token)\n else:\n login()\n # check if git lfs is installed\n try:\n subprocess.call([\"git\", \"lfs\", \"version\"])\n except FileNotFoundError:\n raise OSError(\n \"Looks like you do not have git-lfs installed, please install. \\\n You can install from https://git-lfs.github.com/. \\\n Then run `git lfs install` (you only have to do this once).\"\n )\n\n\ndef _save_model_and_config_for_hf_hub(model: Any, save_dir: str, arch: str, task: str) -> None:\n \"\"\"Save model and config to disk for pushing to huggingface hub\n\n Args:\n model: Onnx model to be saved\n save_dir: directory to save model and config\n arch: architecture name\n task: task name\n \"\"\"\n save_directory = Path(save_dir)\n shutil.copy2(model.model_path, save_directory / \"model.onnx\")\n\n config_path = save_directory / \"config.json\"\n\n # add model configuration\n model_config = model.cfg\n model_config[\"arch\"] = arch\n model_config[\"task\"] = task\n\n with config_path.open(\"w\") as f:\n json.dump(model_config, f, indent=2, ensure_ascii=False)\n\n\ndef push_to_hf_hub(\n model: Any, model_name: str, task: str, override: bool = False, **kwargs\n) -> None: # pragma: no cover\n \"\"\"Save model and its configuration on HF hub\n\n >>> from onnxtr.models import login_to_hub, push_to_hf_hub\n >>> from onnxtr.models.recognition import crnn_mobilenet_v3_small\n >>> login_to_hub()\n >>> model = crnn_mobilenet_v3_small()\n >>> push_to_hf_hub(model, 'my-model', 'recognition', arch='crnn_mobilenet_v3_small')\n\n Args:\n model: Onnx model to be saved\n model_name: name of the model which is also the repository name\n task: task name\n override: whether to override the existing model / repo on HF hub\n **kwargs: keyword arguments for push_to_hf_hub\n \"\"\"\n run_config = kwargs.get(\"run_config\", None)\n arch = kwargs.get(\"arch\", None)\n\n if run_config is None and arch is None:\n raise ValueError(\"run_config or arch must be specified\")\n if task not in [\"classification\", \"detection\", \"recognition\"]:\n raise ValueError(\"task must be one of classification, detection, recognition\")\n\n # default readme\n readme = textwrap.dedent(\n f\"\"\"\n ---\n language:\n - en\n - fr\n license: apache-2.0\n ---\n\n \n \n
\n\n **Optical Character Recognition made seamless & accessible to anyone, powered by Onnxruntime**\n\n ## Task: {task}\n\n https://github.com/felixdittrich92/OnnxTR\n\n ### Example usage:\n\n ```python\n >>> from onnxtr.io import DocumentFile\n >>> from onnxtr.models import ocr_predictor, from_hub\n\n >>> img = DocumentFile.from_images([''])\n >>> # Load your model from the hub\n >>> model = from_hub('onnxtr/my-model')\n\n >>> # Pass it to the predictor\n >>> # If your model is a recognition model:\n >>> predictor = ocr_predictor(det_arch='db_mobilenet_v3_large',\n >>> reco_arch=model)\n\n >>> # If your model is a detection model:\n >>> predictor = ocr_predictor(det_arch=model,\n >>> reco_arch='crnn_mobilenet_v3_small')\n\n >>> # Get your predictions\n >>> res = predictor(img)\n ```\n \"\"\"\n )\n\n # add run configuration to readme if available\n if run_config is not None:\n arch = run_config.arch\n readme += textwrap.dedent(\n f\"\"\"### Run Configuration\n \\n{json.dumps(vars(run_config), indent=2, ensure_ascii=False)}\"\"\"\n )\n\n if arch not in AVAILABLE_ARCHS[task]:\n raise ValueError(\n f\"Architecture: {arch} for task: {task} not found.\\\n \\nAvailable architectures: {AVAILABLE_ARCHS}\"\n )\n\n commit_message = f\"Add {model_name} model\"\n\n # Create repository\n api = HfApi()\n api.create_repo(model_name, token=get_token(), exist_ok=False)\n\n # Save model files to a temporary directory\n with tempfile.TemporaryDirectory() as tmp_dir:\n _save_model_and_config_for_hf_hub(model, tmp_dir, arch=arch, task=task)\n readme_path = Path(tmp_dir) / \"README.md\"\n readme_path.write_text(readme)\n\n # Upload all files to the hub\n api.upload_folder(\n folder_path=tmp_dir,\n repo_id=model_name,\n commit_message=commit_message,\n token=get_token(),\n )\n\n\ndef from_hub(repo_id: str, engine_cfg: EngineConfig | None = None, **kwargs: Any):\n \"\"\"Instantiate & load a pretrained model from HF hub.\n\n >>> from onnxtr.models import from_hub\n >>> model = from_hub(\"onnxtr/my-model\")\n\n Args:\n repo_id: HuggingFace model hub repo\n engine_cfg: configuration for the inference engine (optional)\n **kwargs: kwargs of `hf_hub_download`\n\n Returns:\n Model loaded with the checkpoint\n \"\"\"\n # Get the config\n with open(hf_hub_download(repo_id, filename=\"config.json\", **kwargs), \"rb\") as f:\n cfg = json.load(f)\n model_path = hf_hub_download(repo_id, filename=\"model.onnx\", **kwargs)\n\n arch = cfg[\"arch\"]\n task = cfg[\"task\"]\n cfg.pop(\"arch\")\n cfg.pop(\"task\")\n\n if task == \"classification\":\n model = models.classification.__dict__[arch](model_path, classes=cfg[\"classes\"], engine_cfg=engine_cfg)\n elif task == \"detection\":\n model = models.detection.__dict__[arch](model_path, engine_cfg=engine_cfg)\n elif task == \"recognition\":\n model = models.recognition.__dict__[arch](\n model_path, input_shape=cfg[\"input_shape\"], vocab=cfg[\"vocab\"], engine_cfg=engine_cfg\n )\n\n # convert all values which are lists to tuples\n for key, value in cfg.items():\n if isinstance(value, list):\n cfg[key] = tuple(value)\n # update model cfg\n model.cfg = cfg\n\n return model\n"
},
{
"path": "onnxtr/models/predictor/__init__.py",
"content": "from .predictor import *\n"
},
{
"path": "onnxtr/models/predictor/base.py",
"content": "# Copyright (C) 2021-2026, Mindee | Felix Dittrich.\n\n# This program is licensed under the Apache License 2.0.\n# See LICENSE or go to for full license details.\n\nfrom collections.abc import Callable\nfrom typing import Any\n\nimport numpy as np\n\nfrom onnxtr.models.builder import DocumentBuilder\nfrom onnxtr.models.engine import EngineConfig\nfrom onnxtr.utils.geometry import extract_crops, extract_rcrops, remove_image_padding, rotate_image\n\nfrom .._utils import estimate_orientation, rectify_crops, rectify_loc_preds\nfrom ..classification import crop_orientation_predictor, page_orientation_predictor\nfrom ..classification.predictor import OrientationPredictor\nfrom ..detection.zoo import ARCHS as DETECTION_ARCHS\nfrom ..recognition.zoo import ARCHS as RECOGNITION_ARCHS\n\n__all__ = [\"_OCRPredictor\"]\n\n\nclass _OCRPredictor:\n \"\"\"Implements an object able to localize and identify text elements in a set of documents\n\n Args:\n assume_straight_pages: if True, speeds up the inference by assuming you only pass straight pages\n without rotated textual elements.\n straighten_pages: if True, estimates the page general orientation based on the median line orientation.\n Then, rotates page before passing it to the deep learning modules. The final predictions will be remapped\n accordingly. Doing so will improve performances for documents with page-uniform rotations.\n preserve_aspect_ratio: if True, resize preserving the aspect ratio (with padding)\n symmetric_pad: if True and preserve_aspect_ratio is True, pas the image symmetrically.\n detect_orientation: if True, the estimated general page orientation will be added to the predictions for each\n page. Doing so will slightly deteriorate the overall latency.\n load_in_8_bit: whether to load the the 8-bit quantized model, defaults to False\n clf_engine_cfg: configuration of the orientation classification engine\n **kwargs: keyword args of `DocumentBuilder`\n \"\"\"\n\n crop_orientation_predictor: OrientationPredictor | None\n page_orientation_predictor: OrientationPredictor | None\n\n def __init__(\n self,\n assume_straight_pages: bool = True,\n straighten_pages: bool = False,\n preserve_aspect_ratio: bool = True,\n symmetric_pad: bool = True,\n detect_orientation: bool = False,\n load_in_8_bit: bool = False,\n clf_engine_cfg: EngineConfig | None = None,\n **kwargs: Any,\n ) -> None:\n self.assume_straight_pages = assume_straight_pages\n self.straighten_pages = straighten_pages\n self._page_orientation_disabled = kwargs.pop(\"disable_page_orientation\", False)\n self._crop_orientation_disabled = kwargs.pop(\"disable_crop_orientation\", False)\n self.crop_orientation_predictor = (\n None\n if assume_straight_pages\n else crop_orientation_predictor(\n load_in_8_bit=load_in_8_bit, engine_cfg=clf_engine_cfg, disabled=self._crop_orientation_disabled\n )\n )\n self.page_orientation_predictor = (\n page_orientation_predictor(\n load_in_8_bit=load_in_8_bit, engine_cfg=clf_engine_cfg, disabled=self._crop_orientation_disabled\n )\n if detect_orientation or straighten_pages or not assume_straight_pages\n else None\n )\n self.doc_builder = DocumentBuilder(**kwargs)\n self.preserve_aspect_ratio = preserve_aspect_ratio\n self.symmetric_pad = symmetric_pad\n self.hooks: list[Callable] = []\n\n def _general_page_orientations(\n self,\n pages: list[np.ndarray],\n ) -> list[tuple[int, float]]:\n _, classes, probs = zip(self.page_orientation_predictor(pages)) # type: ignore[misc]\n # Flatten to list of tuples with (value, confidence)\n page_orientations = [\n (orientation, prob)\n for page_classes, page_probs in zip(classes, probs)\n for orientation, prob in zip(page_classes, page_probs)\n ]\n return page_orientations\n\n def _get_orientations(\n self, pages: list[np.ndarray], seg_maps: list[np.ndarray]\n ) -> tuple[list[tuple[int, float]], list[int]]:\n general_pages_orientations = self._general_page_orientations(pages)\n origin_page_orientations = [\n estimate_orientation(seq_map, general_orientation)\n for seq_map, general_orientation in zip(seg_maps, general_pages_orientations)\n ]\n return general_pages_orientations, origin_page_orientations\n\n def _straighten_pages(\n self,\n pages: list[np.ndarray],\n seg_maps: list[np.ndarray],\n general_pages_orientations: list[tuple[int, float]] | None = None,\n origin_pages_orientations: list[int] | None = None,\n ) -> list[np.ndarray]:\n general_pages_orientations = (\n general_pages_orientations if general_pages_orientations else self._general_page_orientations(pages)\n )\n origin_pages_orientations = (\n origin_pages_orientations\n if origin_pages_orientations\n else [\n estimate_orientation(seq_map, general_orientation)\n for seq_map, general_orientation in zip(seg_maps, general_pages_orientations)\n ]\n )\n return [\n # expand if height and width are not equal, afterwards remove padding\n remove_image_padding(rotate_image(page, angle, expand=page.shape[0] != page.shape[1]))\n for page, angle in zip(pages, origin_pages_orientations)\n ]\n\n @staticmethod\n def _generate_crops(\n pages: list[np.ndarray],\n loc_preds: list[np.ndarray],\n channels_last: bool,\n assume_straight_pages: bool = False,\n assume_horizontal: bool = False,\n ) -> list[list[np.ndarray]]:\n if assume_straight_pages:\n crops = [\n extract_crops(page, _boxes[:, :4], channels_last=channels_last)\n for page, _boxes in zip(pages, loc_preds)\n ]\n else:\n crops = [\n extract_rcrops(page, _boxes[:, :4], channels_last=channels_last, assume_horizontal=assume_horizontal)\n for page, _boxes in zip(pages, loc_preds)\n ]\n return crops\n\n @staticmethod\n def _prepare_crops(\n pages: list[np.ndarray],\n loc_preds: list[np.ndarray],\n channels_last: bool,\n assume_straight_pages: bool = False,\n assume_horizontal: bool = False,\n ) -> tuple[list[list[np.ndarray]], list[np.ndarray]]:\n crops = _OCRPredictor._generate_crops(pages, loc_preds, channels_last, assume_straight_pages, assume_horizontal)\n\n # Avoid sending zero-sized crops\n is_kept = [[all(s > 0 for s in crop.shape) for crop in page_crops] for page_crops in crops]\n crops = [\n [crop for crop, _kept in zip(page_crops, page_kept) if _kept]\n for page_crops, page_kept in zip(crops, is_kept)\n ]\n loc_preds = [_boxes[_kept] for _boxes, _kept in zip(loc_preds, is_kept)]\n\n return crops, loc_preds\n\n def _rectify_crops(\n self,\n crops: list[list[np.ndarray]],\n loc_preds: list[np.ndarray],\n ) -> tuple[list[list[np.ndarray]], list[np.ndarray], list[tuple[int, float]]]:\n # Work at a page level\n orientations, classes, probs = zip(*[self.crop_orientation_predictor(page_crops) for page_crops in crops]) # type: ignore[misc]\n rect_crops = [rectify_crops(page_crops, orientation) for page_crops, orientation in zip(crops, orientations)]\n rect_loc_preds = [\n rectify_loc_preds(page_loc_preds, orientation) if len(page_loc_preds) > 0 else page_loc_preds\n for page_loc_preds, orientation in zip(loc_preds, orientations)\n ]\n # Flatten to list of tuples with (value, confidence)\n crop_orientations = [\n (orientation, prob)\n for page_classes, page_probs in zip(classes, probs)\n for orientation, prob in zip(page_classes, page_probs)\n ]\n return rect_crops, rect_loc_preds, crop_orientations # type: ignore[return-value]\n\n @staticmethod\n def _process_predictions(\n loc_preds: list[np.ndarray],\n word_preds: list[tuple[str, float]],\n crop_orientations: list[dict[str, Any]],\n ) -> tuple[list[np.ndarray], list[list[tuple[str, float]]], list[list[dict[str, Any]]]]:\n text_preds = []\n crop_orientation_preds = []\n if len(loc_preds) > 0:\n # Text & crop orientation predictions at page level\n _idx = 0\n for page_boxes in loc_preds:\n text_preds.append(word_preds[_idx : _idx + page_boxes.shape[0]])\n crop_orientation_preds.append(crop_orientations[_idx : _idx + page_boxes.shape[0]])\n _idx += page_boxes.shape[0]\n\n return loc_preds, text_preds, crop_orientation_preds\n\n def add_hook(self, hook: Callable) -> None:\n \"\"\"Add a hook to the predictor\n\n Args:\n hook: a callable that takes as input the `loc_preds` and returns the modified `loc_preds`\n \"\"\"\n self.hooks.append(hook)\n\n def list_archs(self) -> dict[str, list[str]]:\n return {\"detection_archs\": DETECTION_ARCHS, \"recognition_archs\": RECOGNITION_ARCHS}\n"
},
{
"path": "onnxtr/models/predictor/predictor.py",
"content": "# Copyright (C) 2021-2026, Mindee | Felix Dittrich.\n\n# This program is licensed under the Apache License 2.0.\n# See LICENSE or go to for full license details.\n\nfrom typing import Any\n\nimport numpy as np\n\nfrom onnxtr.io.elements import Document\nfrom onnxtr.models._utils import get_language\nfrom onnxtr.models.detection.predictor import DetectionPredictor\nfrom onnxtr.models.engine import EngineConfig\nfrom onnxtr.models.recognition.predictor import RecognitionPredictor\nfrom onnxtr.utils.geometry import detach_scores\nfrom onnxtr.utils.repr import NestedObject\n\nfrom .base import _OCRPredictor\n\n__all__ = [\"OCRPredictor\"]\n\n\nclass OCRPredictor(NestedObject, _OCRPredictor):\n \"\"\"Implements an object able to localize and identify text elements in a set of documents\n\n Args:\n det_predictor: detection module\n reco_predictor: recognition module\n assume_straight_pages: if True, speeds up the inference by assuming you only pass straight pages\n without rotated textual elements.\n straighten_pages: if True, estimates the page general orientation based on the median line orientation.\n Then, rotates page before passing it to the deep learning modules. The final predictions will be remapped\n accordingly. Doing so will improve performances for documents with page-uniform rotations.\n detect_orientation: if True, the estimated general page orientation will be added to the predictions for each\n page. Doing so will slightly deteriorate the overall latency.\n detect_language: if True, the language prediction will be added to the predictions for each\n page. Doing so will slightly deteriorate the overall latency.\n clf_engine_cfg: configuration of the orientation classification engine\n **kwargs: keyword args of `DocumentBuilder`\n \"\"\"\n\n _children_names = [\"det_predictor\", \"reco_predictor\", \"doc_builder\"]\n\n def __init__(\n self,\n det_predictor: DetectionPredictor,\n reco_predictor: RecognitionPredictor,\n assume_straight_pages: bool = True,\n straighten_pages: bool = False,\n preserve_aspect_ratio: bool = True,\n symmetric_pad: bool = True,\n detect_orientation: bool = False,\n detect_language: bool = False,\n clf_engine_cfg: EngineConfig | None = None,\n **kwargs: Any,\n ) -> None:\n self.det_predictor = det_predictor\n self.reco_predictor = reco_predictor\n _OCRPredictor.__init__(\n self,\n assume_straight_pages,\n straighten_pages,\n preserve_aspect_ratio,\n symmetric_pad,\n detect_orientation,\n clf_engine_cfg=clf_engine_cfg,\n **kwargs,\n )\n self.detect_orientation = detect_orientation\n self.detect_language = detect_language\n\n def __call__(\n self,\n pages: list[np.ndarray],\n **kwargs: Any,\n ) -> Document:\n # Dimension check\n if any(page.ndim != 3 for page in pages):\n raise ValueError(\"incorrect input shape: all pages are expected to be multi-channel 2D images.\")\n\n origin_page_shapes = [page.shape[:2] for page in pages]\n\n # Localize text elements\n loc_preds, out_maps = self.det_predictor(pages, return_maps=True, **kwargs)\n\n # Detect document rotation and rotate pages\n seg_maps = [\n np.where(out_map > getattr(self.det_predictor.model.postprocessor, \"bin_thresh\"), 255, 0).astype(np.uint8)\n for out_map in out_maps\n ]\n if self.detect_orientation:\n general_pages_orientations, origin_pages_orientations = self._get_orientations(pages, seg_maps)\n orientations = [\n {\"value\": orientation_page, \"confidence\": None} for orientation_page in origin_pages_orientations\n ]\n else:\n orientations = None\n general_pages_orientations = None\n origin_pages_orientations = None\n if self.straighten_pages:\n pages = self._straighten_pages(pages, seg_maps, general_pages_orientations, origin_pages_orientations)\n # update page shapes after straightening\n origin_page_shapes = [page.shape[:2] for page in pages]\n\n # forward again to get predictions on straight pages\n loc_preds = self.det_predictor(pages, **kwargs) # type: ignore[assignment]\n\n # Detach objectness scores from loc_preds\n loc_preds, objectness_scores = detach_scores(loc_preds) # type: ignore[arg-type]\n\n # Apply hooks to loc_preds if any\n for hook in self.hooks:\n loc_preds = hook(loc_preds)\n\n # Crop images\n crops, loc_preds = self._prepare_crops(\n pages,\n loc_preds,\n channels_last=True,\n assume_straight_pages=self.assume_straight_pages,\n assume_horizontal=self._page_orientation_disabled,\n )\n # Rectify crop orientation and get crop orientation predictions\n crop_orientations: Any = []\n if not self.assume_straight_pages:\n crops, loc_preds, _crop_orientations = self._rectify_crops(crops, loc_preds)\n crop_orientations = [\n {\"value\": orientation[0], \"confidence\": orientation[1]} for orientation in _crop_orientations\n ]\n\n # Identify character sequences\n word_preds = self.reco_predictor([crop for page_crops in crops for crop in page_crops], **kwargs)\n if not crop_orientations:\n crop_orientations = [{\"value\": 0, \"confidence\": None} for _ in word_preds]\n\n boxes, text_preds, crop_orientations = self._process_predictions(loc_preds, word_preds, crop_orientations)\n\n if self.detect_language:\n languages = [get_language(\" \".join([item[0] for item in text_pred])) for text_pred in text_preds]\n languages_dict = [{\"value\": lang[0], \"confidence\": lang[1]} for lang in languages]\n else:\n languages_dict = None\n\n out = self.doc_builder(\n pages,\n boxes,\n objectness_scores,\n text_preds,\n origin_page_shapes,\n crop_orientations,\n orientations,\n languages_dict,\n )\n return out\n"
},
{
"path": "onnxtr/models/preprocessor/__init__.py",
"content": "from .base import *\n"
},
{
"path": "onnxtr/models/preprocessor/base.py",
"content": "# Copyright (C) 2021-2026, Mindee | Felix Dittrich.\n\n# This program is licensed under the Apache License 2.0.\n# See LICENSE or go to for full license details.\n\nimport math\nfrom typing import Any\n\nimport numpy as np\n\nfrom onnxtr.transforms import Normalize, Resize\nfrom onnxtr.utils.geometry import shape_translate\nfrom onnxtr.utils.multithreading import multithread_exec\nfrom onnxtr.utils.repr import NestedObject\n\n__all__ = [\"PreProcessor\"]\n\n\nclass PreProcessor(NestedObject):\n \"\"\"Implements an abstract preprocessor object which performs casting, resizing, batching and normalization.\n\n Args:\n output_size: expected size of each page in format (H, W)\n batch_size: the size of page batches\n mean: mean value of the training distribution by channel\n std: standard deviation of the training distribution by channel\n **kwargs: additional arguments for the resizing operation\n \"\"\"\n\n _children_names: list[str] = [\"resize\", \"normalize\"]\n\n def __init__(\n self,\n output_size: tuple[int, int],\n batch_size: int,\n mean: tuple[float, float, float] = (0.5, 0.5, 0.5),\n std: tuple[float, float, float] = (1.0, 1.0, 1.0),\n **kwargs: Any,\n ) -> None:\n self.batch_size = batch_size\n self.resize = Resize(output_size, **kwargs)\n self.normalize = Normalize(mean, std)\n\n def batch_inputs(self, samples: list[np.ndarray]) -> list[np.ndarray]:\n \"\"\"Gather samples into batches for inference purposes\n\n Args:\n samples: list of samples (tf.Tensor)\n\n Returns:\n list of batched samples\n \"\"\"\n num_batches = int(math.ceil(len(samples) / self.batch_size))\n batches = [\n np.stack(samples[idx * self.batch_size : min((idx + 1) * self.batch_size, len(samples))], axis=0)\n for idx in range(int(num_batches))\n ]\n\n return batches\n\n def sample_transforms(self, x: np.ndarray) -> np.ndarray:\n if x.ndim != 3:\n raise AssertionError(\"expected list of 3D Tensors\")\n if isinstance(x, np.ndarray):\n if x.dtype not in (np.uint8, np.float32):\n raise TypeError(\"unsupported data type for numpy.ndarray\")\n x = shape_translate(x, \"HWC\")\n\n # Resizing\n x = self.resize(x)\n # Data type & 255 division\n if x.dtype == np.uint8:\n x = x.astype(np.float32) / 255.0\n\n return x\n\n def __call__(self, x: np.ndarray | list[np.ndarray]) -> list[np.ndarray]:\n \"\"\"Prepare document data for model forwarding\n\n Args:\n x: list of images (np.array) or tensors (already resized and batched)\n\n Returns:\n list of page batches\n \"\"\"\n # Input type check\n if isinstance(x, np.ndarray):\n if x.ndim != 4:\n raise AssertionError(\"expected 4D Tensor\")\n if isinstance(x, np.ndarray):\n if x.dtype not in (np.uint8, np.float32):\n raise TypeError(\"unsupported data type for numpy.ndarray\")\n x = shape_translate(x, \"BHWC\")\n\n # Resizing\n if (x.shape[1], x.shape[2]) != self.resize.output_size:\n x = np.array([self.resize(sample) for sample in x])\n # Data type & 255 division\n if x.dtype == np.uint8:\n x = x.astype(np.float32) / 255.0\n batches = [x]\n\n elif isinstance(x, list) and all(isinstance(sample, np.ndarray) for sample in x):\n # Sample transform (to tensor, resize)\n samples = list(multithread_exec(self.sample_transforms, x))\n # Batching\n batches = self.batch_inputs(samples)\n else:\n raise TypeError(f\"invalid input type: {type(x)}\")\n\n # Batch transforms (normalize)\n batches = list(multithread_exec(self.normalize, batches))\n\n return batches\n"
},
{
"path": "onnxtr/models/recognition/__init__.py",
"content": "from .models import *\nfrom .zoo import *\n"
},
{
"path": "onnxtr/models/recognition/core.py",
"content": "# Copyright (C) 2021-2026, Mindee | Felix Dittrich.\n\n# This program is licensed under the Apache License 2.0.\n# See LICENSE or go to for full license details.\n\n\nfrom onnxtr.utils.repr import NestedObject\n\n__all__ = [\"RecognitionPostProcessor\"]\n\n\nclass RecognitionPostProcessor(NestedObject):\n \"\"\"Abstract class to postprocess the raw output of the model\n\n Args:\n vocab: string containing the ordered sequence of supported characters\n \"\"\"\n\n def __init__(\n self,\n vocab: str,\n ) -> None:\n self.vocab = vocab\n self._embedding = list(self.vocab) + [\"\"]\n\n def extra_repr(self) -> str:\n return f\"vocab_size={len(self.vocab)}\"\n"
},
{
"path": "onnxtr/models/recognition/models/__init__.py",
"content": "from .crnn import *\nfrom .sar import *\nfrom .master import *\nfrom .vitstr import *\nfrom .parseq import *\nfrom .viptr import *\n"
},
{
"path": "onnxtr/models/recognition/models/crnn.py",
"content": "# Copyright (C) 2021-2026, Mindee | Felix Dittrich.\n\n# This program is licensed under the Apache License 2.0.\n# See LICENSE or go to for full license details.\n\nfrom copy import deepcopy\nfrom itertools import groupby\nfrom typing import Any\n\nimport numpy as np\nfrom scipy.special import softmax\n\nfrom onnxtr.utils import VOCABS\n\nfrom ...engine import Engine, EngineConfig\nfrom ..core import RecognitionPostProcessor\n\n__all__ = [\"CRNN\", \"crnn_vgg16_bn\", \"crnn_mobilenet_v3_small\", \"crnn_mobilenet_v3_large\"]\n\ndefault_cfgs: dict[str, dict[str, Any]] = {\n \"crnn_vgg16_bn\": {\n \"mean\": (0.694, 0.695, 0.693),\n \"std\": (0.299, 0.296, 0.301),\n \"input_shape\": (3, 32, 128),\n \"vocab\": VOCABS[\"french\"],\n \"url\": \"https://github.com/felixdittrich92/OnnxTR/releases/download/v0.7.1/crnn_vgg16_bn-743599aa.onnx\",\n \"url_8_bit\": \"https://github.com/felixdittrich92/OnnxTR/releases/download/v0.7.1/crnn_vgg16_bn_static_8_bit-df1b594d.onnx\",\n },\n \"crnn_mobilenet_v3_small\": {\n \"mean\": (0.694, 0.695, 0.693),\n \"std\": (0.299, 0.296, 0.301),\n \"input_shape\": (3, 32, 128),\n \"vocab\": VOCABS[\"french\"],\n \"url\": \"https://github.com/felixdittrich92/OnnxTR/releases/download/v0.0.1/crnn_mobilenet_v3_small-bded4d49.onnx\",\n \"url_8_bit\": \"https://github.com/felixdittrich92/OnnxTR/releases/download/v0.1.2/crnn_mobilenet_v3_small_static_8_bit-4949006f.onnx\",\n },\n \"crnn_mobilenet_v3_large\": {\n \"mean\": (0.694, 0.695, 0.693),\n \"std\": (0.299, 0.296, 0.301),\n \"input_shape\": (3, 32, 128),\n \"vocab\": VOCABS[\"french\"],\n \"url\": \"https://github.com/felixdittrich92/OnnxTR/releases/download/v0.0.1/crnn_mobilenet_v3_large-d42e8185.onnx\",\n \"url_8_bit\": \"https://github.com/felixdittrich92/OnnxTR/releases/download/v0.1.2/crnn_mobilenet_v3_large_static_8_bit-459e856d.onnx\",\n },\n}\n\n\nclass CRNNPostProcessor(RecognitionPostProcessor):\n \"\"\"Postprocess raw prediction of the model (logits) to a list of words using CTC decoding\n\n Args:\n vocab: string containing the ordered sequence of supported characters\n \"\"\"\n\n def __init__(self, vocab):\n self.vocab = vocab\n\n def decode_sequence(self, sequence, vocab):\n return \"\".join([vocab[int(char)] for char in sequence])\n\n def ctc_best_path(\n self,\n logits,\n vocab,\n blank=0,\n ):\n \"\"\"Implements best path decoding as shown by Graves (Dissertation, p63), highly inspired from\n `_.\n\n Args:\n logits: model output, shape: N x T x C\n vocab: vocabulary to use\n blank: index of blank label\n\n Returns:\n A list of tuples: (word, confidence)\n \"\"\"\n # Gather the most confident characters, and assign the smallest conf among those to the sequence prob\n probs = softmax(logits, axis=-1).max(axis=-1).min(axis=1)\n\n # collapse best path (using itertools.groupby), map to chars, join char list to string\n words = [\n self.decode_sequence([k for k, _ in groupby(seq.tolist()) if k != blank], vocab)\n for seq in np.argmax(logits, axis=-1)\n ]\n\n return list(zip(words, probs.astype(float).tolist()))\n\n def __call__(self, logits):\n \"\"\"Performs decoding of raw output with CTC and decoding of CTC predictions\n with label_to_idx mapping dictionnary\n\n Args:\n logits: raw output of the model, shape (N, C + 1, seq_len)\n\n Returns:\n A tuple of 2 lists: a list of str (words) and a list of float (probs)\n\n \"\"\"\n # Decode CTC\n return self.ctc_best_path(logits=logits, vocab=self.vocab, blank=len(self.vocab))\n\n\nclass CRNN(Engine):\n \"\"\"CRNN Onnx loader\n\n Args:\n model_path: path or url to onnx model file\n vocab: vocabulary used for encoding\n engine_cfg: configuration for the inference engine\n cfg: configuration dictionary\n **kwargs: additional arguments to be passed to `Engine`\n \"\"\"\n\n _children_names: list[str] = [\"postprocessor\"]\n\n def __init__(\n self,\n model_path: str,\n vocab: str,\n engine_cfg: EngineConfig | None = None,\n cfg: dict[str, Any] | None = None,\n **kwargs: Any,\n ) -> None:\n super().__init__(url=model_path, engine_cfg=engine_cfg, **kwargs)\n\n self.vocab = vocab\n self.cfg = cfg\n\n self.postprocessor = CRNNPostProcessor(self.vocab)\n\n def __call__(\n self,\n x: np.ndarray,\n return_model_output: bool = False,\n ) -> dict[str, Any]:\n logits = self.run(x)\n\n out: dict[str, Any] = {}\n if return_model_output:\n out[\"out_map\"] = logits\n\n # Post-process\n out[\"preds\"] = self.postprocessor(logits)\n\n return out\n\n\ndef _crnn(\n arch: str,\n model_path: str,\n load_in_8_bit: bool = False,\n engine_cfg: EngineConfig | None = None,\n **kwargs: Any,\n) -> CRNN:\n kwargs[\"vocab\"] = kwargs.get(\"vocab\", default_cfgs[arch][\"vocab\"])\n\n _cfg = deepcopy(default_cfgs[arch])\n _cfg[\"vocab\"] = kwargs[\"vocab\"]\n _cfg[\"input_shape\"] = kwargs.get(\"input_shape\", default_cfgs[arch][\"input_shape\"])\n # Patch the url\n model_path = default_cfgs[arch][\"url_8_bit\"] if load_in_8_bit and \"http\" in model_path else model_path\n\n # Build the model\n return CRNN(model_path, cfg=_cfg, engine_cfg=engine_cfg, **kwargs)\n\n\ndef crnn_vgg16_bn(\n model_path: str = default_cfgs[\"crnn_vgg16_bn\"][\"url\"],\n load_in_8_bit: bool = False,\n engine_cfg: EngineConfig | None = None,\n **kwargs: Any,\n) -> CRNN:\n \"\"\"CRNN with a VGG-16 backbone as described in `\"An End-to-End Trainable Neural Network for Image-based\n Sequence Recognition and Its Application to Scene Text Recognition\" `_.\n\n >>> import numpy as np\n >>> from onnxtr.models import crnn_vgg16_bn\n >>> model = crnn_vgg16_bn()\n >>> input_tensor = np.random.rand(1, 3, 32, 128)\n >>> out = model(input_tensor)\n\n Args:\n model_path: path to onnx model file, defaults to url in default_cfgs\n load_in_8_bit: whether to load the the 8-bit quantized model, defaults to False\n engine_cfg: configuration for the inference engine\n **kwargs: keyword arguments of the CRNN architecture\n\n Returns:\n text recognition architecture\n \"\"\"\n return _crnn(\"crnn_vgg16_bn\", model_path, load_in_8_bit, engine_cfg, **kwargs)\n\n\ndef crnn_mobilenet_v3_small(\n model_path: str = default_cfgs[\"crnn_mobilenet_v3_small\"][\"url\"],\n load_in_8_bit: bool = False,\n engine_cfg: EngineConfig | None = None,\n **kwargs: Any,\n) -> CRNN:\n \"\"\"CRNN with a MobileNet V3 Small backbone as described in `\"An End-to-End Trainable Neural Network for Image-based\n Sequence Recognition and Its Application to Scene Text Recognition\" `_.\n\n >>> import numpy as np\n >>> from onnxtr.models import crnn_mobilenet_v3_small\n >>> model = crnn_mobilenet_v3_small()\n >>> input_tensor = np.random.rand(1, 3, 32, 128)\n >>> out = model(input_tensor)\n\n Args:\n model_path: path to onnx model file, defaults to url in default_cfgs\n load_in_8_bit: whether to load the the 8-bit quantized model, defaults to False\n engine_cfg: configuration for the inference engine\n **kwargs: keyword arguments of the CRNN architecture\n\n Returns:\n text recognition architecture\n \"\"\"\n return _crnn(\"crnn_mobilenet_v3_small\", model_path, load_in_8_bit, engine_cfg, **kwargs)\n\n\ndef crnn_mobilenet_v3_large(\n model_path: str = default_cfgs[\"crnn_mobilenet_v3_large\"][\"url\"],\n load_in_8_bit: bool = False,\n engine_cfg: EngineConfig | None = None,\n **kwargs: Any,\n) -> CRNN:\n \"\"\"CRNN with a MobileNet V3 Large backbone as described in `\"An End-to-End Trainable Neural Network for Image-based\n Sequence Recognition and Its Application to Scene Text Recognition\" `_.\n\n >>> import numpy as np\n >>> from onnxtr.models import crnn_mobilenet_v3_large\n >>> model = crnn_mobilenet_v3_large()\n >>> input_tensor = np.random.rand(1, 3, 32, 128)\n >>> out = model(input_tensor)\n\n Args:\n model_path: path to onnx model file, defaults to url in default_cfgs\n load_in_8_bit: whether to load the the 8-bit quantized model, defaults to False\n engine_cfg: configuration for the inference engine\n **kwargs: keyword arguments of the CRNN architecture\n\n Returns:\n text recognition architecture\n \"\"\"\n return _crnn(\"crnn_mobilenet_v3_large\", model_path, load_in_8_bit, engine_cfg, **kwargs)\n"
},
{
"path": "onnxtr/models/recognition/models/master.py",
"content": "# Copyright (C) 2021-2026, Mindee | Felix Dittrich.\n\n# This program is licensed under the Apache License 2.0.\n# See LICENSE or go to for full license details.\n\nfrom copy import deepcopy\nfrom typing import Any\n\nimport numpy as np\nfrom scipy.special import softmax\n\nfrom onnxtr.utils import VOCABS\n\nfrom ...engine import Engine, EngineConfig\nfrom ..core import RecognitionPostProcessor\n\n__all__ = [\"MASTER\", \"master\"]\n\n\ndefault_cfgs: dict[str, dict[str, Any]] = {\n \"master\": {\n \"mean\": (0.694, 0.695, 0.693),\n \"std\": (0.299, 0.296, 0.301),\n \"input_shape\": (3, 32, 128),\n \"vocab\": VOCABS[\"french\"],\n \"url\": \"https://github.com/felixdittrich92/OnnxTR/releases/download/v0.0.1/master-b1287fcd.onnx\",\n \"url_8_bit\": \"https://github.com/felixdittrich92/OnnxTR/releases/download/v0.1.2/master_dynamic_8_bit-d8bd8206.onnx\",\n },\n}\n\n\nclass MASTER(Engine):\n \"\"\"MASTER Onnx loader\n\n Args:\n model_path: path or url to onnx model file\n vocab: vocabulary, (without EOS, SOS, PAD)\n engine_cfg: configuration for the inference engine\n cfg: dictionary containing information about the model\n **kwargs: additional arguments to be passed to `Engine`\n \"\"\"\n\n def __init__(\n self,\n model_path: str,\n vocab: str,\n engine_cfg: EngineConfig | None = None,\n cfg: dict[str, Any] | None = None,\n **kwargs: Any,\n ) -> None:\n super().__init__(url=model_path, engine_cfg=engine_cfg, **kwargs)\n\n self.vocab = vocab\n self.cfg = cfg\n\n self.postprocessor = MASTERPostProcessor(vocab=self.vocab)\n\n def __call__(\n self,\n x: np.ndarray,\n return_model_output: bool = False,\n ) -> dict[str, Any]:\n \"\"\"Call function\n\n Args:\n x: images\n return_model_output: if True, return logits\n\n Returns:\n A dictionnary containing eventually logits and predictions.\n \"\"\"\n logits = self.run(x)\n out: dict[str, Any] = {}\n\n if return_model_output:\n out[\"out_map\"] = logits\n\n out[\"preds\"] = self.postprocessor(logits)\n\n return out\n\n\nclass MASTERPostProcessor(RecognitionPostProcessor):\n \"\"\"Post-processor for the MASTER model\n\n Args:\n vocab: string containing the ordered sequence of supported characters\n \"\"\"\n\n def __init__(\n self,\n vocab: str,\n ) -> None:\n super().__init__(vocab)\n self._embedding = list(vocab) + [\"\"] + [\"\"] + [\"\"]\n\n def __call__(self, logits: np.ndarray) -> list[tuple[str, float]]:\n # compute pred with argmax for attention models\n out_idxs = np.argmax(logits, axis=-1)\n # N x L\n probs = np.take_along_axis(softmax(logits, axis=-1), out_idxs[..., None], axis=-1).squeeze(-1)\n # Take the minimum confidence of the sequence\n probs = np.min(probs, axis=1)\n\n word_values = [\n \"\".join(self._embedding[idx] for idx in encoded_seq).split(\"\")[0] for encoded_seq in out_idxs\n ]\n\n return list(zip(word_values, np.clip(probs, 0, 1).astype(float).tolist()))\n\n\ndef _master(\n arch: str,\n model_path: str,\n load_in_8_bit: bool = False,\n engine_cfg: EngineConfig | None = None,\n **kwargs: Any,\n) -> MASTER:\n # Patch the config\n _cfg = deepcopy(default_cfgs[arch])\n _cfg[\"input_shape\"] = kwargs.get(\"input_shape\", _cfg[\"input_shape\"])\n _cfg[\"vocab\"] = kwargs.get(\"vocab\", _cfg[\"vocab\"])\n\n kwargs[\"vocab\"] = _cfg[\"vocab\"]\n # Patch the url\n model_path = default_cfgs[arch][\"url_8_bit\"] if load_in_8_bit and \"http\" in model_path else model_path\n\n return MASTER(model_path, cfg=_cfg, engine_cfg=engine_cfg, **kwargs)\n\n\ndef master(\n model_path: str = default_cfgs[\"master\"][\"url\"],\n load_in_8_bit: bool = False,\n engine_cfg: EngineConfig | None = None,\n **kwargs: Any,\n) -> MASTER:\n \"\"\"MASTER as described in paper: `_.\n\n >>> import numpy as np\n >>> from onnxtr.models import master\n >>> model = master()\n >>> input_tensor = np.random.rand(1, 3, 32, 128)\n >>> out = model(input_tensor)\n\n Args:\n model_path: path to onnx model file, defaults to url in default_cfgs\n load_in_8_bit: whether to load the the 8-bit quantized model, defaults to False\n engine_cfg: configuration for the inference engine\n **kwargs: keywoard arguments passed to the MASTER architecture\n\n Returns:\n text recognition architecture\n \"\"\"\n return _master(\"master\", model_path, load_in_8_bit, engine_cfg, **kwargs)\n"
},
{
"path": "onnxtr/models/recognition/models/parseq.py",
"content": "# Copyright (C) 2021-2026, Mindee | Felix Dittrich.\n\n# This program is licensed under the Apache License 2.0.\n# See LICENSE or go to for full license details.\n\nfrom copy import deepcopy\nfrom typing import Any\n\nimport numpy as np\nfrom scipy.special import softmax\n\nfrom onnxtr.utils import VOCABS\n\nfrom ...engine import Engine, EngineConfig\nfrom ..core import RecognitionPostProcessor\n\n__all__ = [\"PARSeq\", \"parseq\"]\n\ndefault_cfgs: dict[str, dict[str, Any]] = {\n \"parseq\": {\n \"mean\": (0.694, 0.695, 0.693),\n \"std\": (0.299, 0.296, 0.301),\n \"input_shape\": (3, 32, 128),\n \"vocab\": VOCABS[\"french\"],\n \"url\": \"https://github.com/felixdittrich92/OnnxTR/releases/download/v0.0.1/parseq-00b40714.onnx\",\n \"url_8_bit\": \"https://github.com/felixdittrich92/OnnxTR/releases/download/v0.1.2/parseq_dynamic_8_bit-5b04d9f7.onnx\",\n },\n}\n\n\nclass PARSeq(Engine):\n \"\"\"PARSeq Onnx loader\n\n Args:\n model_path: path to onnx model file\n vocab: vocabulary used for encoding\n engine_cfg: configuration for the inference engine\n cfg: dictionary containing information about the model\n **kwargs: additional arguments to be passed to `Engine`\n \"\"\"\n\n def __init__(\n self,\n model_path: str,\n vocab: str,\n engine_cfg: EngineConfig | None = None,\n cfg: dict[str, Any] | None = None,\n **kwargs: Any,\n ) -> None:\n super().__init__(url=model_path, engine_cfg=engine_cfg, **kwargs)\n\n self.vocab = vocab\n self.cfg = cfg\n\n self.postprocessor = PARSeqPostProcessor(vocab=self.vocab)\n\n def __call__(\n self,\n x: np.ndarray,\n return_model_output: bool = False,\n ) -> dict[str, Any]:\n logits = self.run(x)\n out: dict[str, Any] = {}\n\n if return_model_output:\n out[\"out_map\"] = logits\n\n out[\"preds\"] = self.postprocessor(logits)\n return out\n\n\nclass PARSeqPostProcessor(RecognitionPostProcessor):\n \"\"\"Post processor for PARSeq architecture\n\n Args:\n vocab: string containing the ordered sequence of supported characters\n \"\"\"\n\n def __init__(\n self,\n vocab: str,\n ) -> None:\n super().__init__(vocab)\n self._embedding = list(vocab) + [\"\", \"\", \"\"]\n\n def __call__(self, logits):\n # compute pred with argmax for attention models\n out_idxs = np.argmax(logits, axis=-1)\n preds_prob = softmax(logits, axis=-1).max(axis=-1)\n\n word_values = [\n \"\".join(self._embedding[idx] for idx in encoded_seq).split(\"\")[0] for encoded_seq in out_idxs\n ]\n # compute probabilties for each word up to the EOS token\n probs = [\n preds_prob[i, : len(word)].clip(0, 1).mean().astype(float) if word else 0.0\n for i, word in enumerate(word_values)\n ]\n\n return list(zip(word_values, probs))\n\n\ndef _parseq(\n arch: str,\n model_path: str,\n load_in_8_bit: bool = False,\n engine_cfg: EngineConfig | None = None,\n **kwargs: Any,\n) -> PARSeq:\n # Patch the config\n _cfg = deepcopy(default_cfgs[arch])\n _cfg[\"vocab\"] = kwargs.get(\"vocab\", _cfg[\"vocab\"])\n _cfg[\"input_shape\"] = kwargs.get(\"input_shape\", _cfg[\"input_shape\"])\n\n kwargs[\"vocab\"] = _cfg[\"vocab\"]\n # Patch the url\n model_path = default_cfgs[arch][\"url_8_bit\"] if load_in_8_bit and \"http\" in model_path else model_path\n\n # Build the model\n return PARSeq(model_path, cfg=_cfg, engine_cfg=engine_cfg, **kwargs)\n\n\ndef parseq(\n model_path: str = default_cfgs[\"parseq\"][\"url\"],\n load_in_8_bit: bool = False,\n engine_cfg: EngineConfig | None = None,\n **kwargs: Any,\n) -> PARSeq:\n \"\"\"PARSeq architecture from\n `\"Scene Text Recognition with Permuted Autoregressive Sequence Models\" `_.\n\n >>> import numpy as np\n >>> from onnxtr.models import parseq\n >>> model = parseq()\n >>> input_tensor = np.random.rand(1, 3, 32, 128)\n >>> out = model(input_tensor)\n\n Args:\n model_path: path to onnx model file, defaults to url in default_cfgs\n load_in_8_bit: whether to load the the 8-bit quantized model, defaults to False\n engine_cfg: configuration for the inference engine\n **kwargs: keyword arguments of the PARSeq architecture\n\n Returns:\n text recognition architecture\n \"\"\"\n return _parseq(\"parseq\", model_path, load_in_8_bit, engine_cfg, **kwargs)\n"
},
{
"path": "onnxtr/models/recognition/models/sar.py",
"content": "# Copyright (C) 2021-2026, Mindee | Felix Dittrich.\n\n# This program is licensed under the Apache License 2.0.\n# See LICENSE or go to for full license details.\n\nfrom copy import deepcopy\nfrom typing import Any\n\nimport numpy as np\nfrom scipy.special import softmax\n\nfrom onnxtr.utils import VOCABS\n\nfrom ...engine import Engine, EngineConfig\nfrom ..core import RecognitionPostProcessor\n\n__all__ = [\"SAR\", \"sar_resnet31\"]\n\ndefault_cfgs: dict[str, dict[str, Any]] = {\n \"sar_resnet31\": {\n \"mean\": (0.694, 0.695, 0.693),\n \"std\": (0.299, 0.296, 0.301),\n \"input_shape\": (3, 32, 128),\n \"vocab\": VOCABS[\"french\"],\n \"url\": \"https://github.com/felixdittrich92/OnnxTR/releases/download/v0.0.1/sar_resnet31-395f8005.onnx\",\n \"url_8_bit\": \"https://github.com/felixdittrich92/OnnxTR/releases/download/v0.1.2/sar_resnet31_static_8_bit-c07316bc.onnx\",\n },\n}\n\n\nclass SAR(Engine):\n \"\"\"SAR Onnx loader\n\n Args:\n model_path: path to onnx model file\n vocab: vocabulary used for encoding\n engine_cfg: configuration for the inference engine\n cfg: dictionary containing information about the model\n **kwargs: additional arguments to be passed to `Engine`\n \"\"\"\n\n def __init__(\n self,\n model_path: str,\n vocab: str,\n engine_cfg: EngineConfig | None = None,\n cfg: dict[str, Any] | None = None,\n **kwargs: Any,\n ) -> None:\n super().__init__(url=model_path, engine_cfg=engine_cfg, **kwargs)\n\n self.vocab = vocab\n self.cfg = cfg\n\n self.postprocessor = SARPostProcessor(self.vocab)\n\n def __call__(\n self,\n x: np.ndarray,\n return_model_output: bool = False,\n ) -> dict[str, Any]:\n logits = self.run(x)\n\n out: dict[str, Any] = {}\n if return_model_output:\n out[\"out_map\"] = logits\n\n out[\"preds\"] = self.postprocessor(logits)\n\n return out\n\n\nclass SARPostProcessor(RecognitionPostProcessor):\n \"\"\"Post processor for SAR architectures\n\n Args:\n embedding: string containing the ordered sequence of supported characters\n \"\"\"\n\n def __init__(\n self,\n vocab: str,\n ) -> None:\n super().__init__(vocab)\n self._embedding = list(self.vocab) + [\"\"]\n\n def __call__(self, logits):\n # compute pred with argmax for attention models\n out_idxs = np.argmax(logits, axis=-1)\n # N x L\n probs = np.take_along_axis(softmax(logits, axis=-1), out_idxs[..., None], axis=-1).squeeze(-1)\n # Take the minimum confidence of the sequence\n probs = np.min(probs, axis=1)\n\n word_values = [\n \"\".join(self._embedding[idx] for idx in encoded_seq).split(\"\")[0] for encoded_seq in out_idxs\n ]\n\n return list(zip(word_values, np.clip(probs, 0, 1).astype(float).tolist()))\n\n\ndef _sar(\n arch: str,\n model_path: str,\n load_in_8_bit: bool = False,\n engine_cfg: EngineConfig | None = None,\n **kwargs: Any,\n) -> SAR:\n # Patch the config\n _cfg = deepcopy(default_cfgs[arch])\n _cfg[\"vocab\"] = kwargs.get(\"vocab\", _cfg[\"vocab\"])\n _cfg[\"input_shape\"] = kwargs.get(\"input_shape\", _cfg[\"input_shape\"])\n\n kwargs[\"vocab\"] = _cfg[\"vocab\"]\n # Patch the url\n model_path = default_cfgs[arch][\"url_8_bit\"] if load_in_8_bit and \"http\" in model_path else model_path\n\n # Build the model\n return SAR(model_path, cfg=_cfg, engine_cfg=engine_cfg, **kwargs)\n\n\ndef sar_resnet31(\n model_path: str = default_cfgs[\"sar_resnet31\"][\"url\"],\n load_in_8_bit: bool = False,\n engine_cfg: EngineConfig | None = None,\n **kwargs: Any,\n) -> SAR:\n \"\"\"SAR with a resnet-31 feature extractor as described in `\"Show, Attend and Read:A Simple and Strong\n Baseline for Irregular Text Recognition\" `_.\n\n >>> import numpy as np\n >>> from onnxtr.models import sar_resnet31\n >>> model = sar_resnet31()\n >>> input_tensor = np.random.rand(1, 3, 32, 128)\n >>> out = model(input_tensor)\n\n Args:\n model_path: path to onnx model file, defaults to url in default_cfgs\n load_in_8_bit: whether to load the the 8-bit quantized model, defaults to False\n engine_cfg: configuration for the inference engine\n **kwargs: keyword arguments of the SAR architecture\n\n Returns:\n text recognition architecture\n \"\"\"\n return _sar(\"sar_resnet31\", model_path, load_in_8_bit, engine_cfg, **kwargs)\n"
},
{
"path": "onnxtr/models/recognition/models/viptr.py",
"content": "# Copyright (C) 2021-2026, Mindee | Felix Dittrich.\n\n# This program is licensed under the Apache License 2.0.\n# See LICENSE or go to for full license details.\n\nimport logging\nfrom copy import deepcopy\nfrom itertools import groupby\nfrom typing import Any\n\nimport numpy as np\nfrom scipy.special import softmax\n\nfrom onnxtr.utils import VOCABS\n\nfrom ...engine import Engine, EngineConfig\nfrom ..core import RecognitionPostProcessor\n\n__all__ = [\"VIPTR\", \"viptr_tiny\"]\n\ndefault_cfgs: dict[str, dict[str, Any]] = {\n \"viptr_tiny\": {\n \"mean\": (0.694, 0.695, 0.693),\n \"std\": (0.299, 0.296, 0.301),\n \"input_shape\": (3, 32, 128),\n \"vocab\": VOCABS[\"french\"],\n \"url\": \"https://github.com/felixdittrich92/OnnxTR/releases/download/v0.6.3/viptr_tiny-499b8015.onnx\",\n \"url_8_bit\": \"https://github.com/felixdittrich92/OnnxTR/releases/download/v0.6.3/viptr_tiny-499b8015.onnx\",\n },\n}\n\n\nclass VIPTRPostProcessor(RecognitionPostProcessor):\n \"\"\"Postprocess raw prediction of the model (logits) to a list of words using CTC decoding\n\n Args:\n vocab: string containing the ordered sequence of supported characters\n \"\"\"\n\n def __init__(self, vocab):\n self.vocab = vocab\n\n def decode_sequence(self, sequence, vocab):\n return \"\".join([vocab[int(char)] for char in sequence])\n\n def ctc_best_path(\n self,\n logits,\n vocab,\n blank=0,\n ):\n \"\"\"Implements best path decoding as shown by Graves (Dissertation, p63), highly inspired from\n `_.\n\n Args:\n logits: model output, shape: N x T x C\n vocab: vocabulary to use\n blank: index of blank label\n\n Returns:\n A list of tuples: (word, confidence)\n \"\"\"\n # Gather the most confident characters, and assign the smallest conf among those to the sequence prob\n probs = softmax(logits, axis=-1).max(axis=-1).min(axis=1)\n\n # collapse best path (using itertools.groupby), map to chars, join char list to string\n words = [\n self.decode_sequence([k for k, _ in groupby(seq.tolist()) if k != blank], vocab)\n for seq in np.argmax(logits, axis=-1)\n ]\n\n return list(zip(words, probs.astype(float).tolist()))\n\n def __call__(self, logits):\n \"\"\"Performs decoding of raw output with CTC and decoding of CTC predictions\n with label_to_idx mapping dictionnary\n\n Args:\n logits: raw output of the model, shape (N, C + 1, seq_len)\n\n Returns:\n A tuple of 2 lists: a list of str (words) and a list of float (probs)\n\n \"\"\"\n # Decode CTC\n return self.ctc_best_path(logits=logits, vocab=self.vocab, blank=len(self.vocab))\n\n\nclass VIPTR(Engine):\n \"\"\"VIPTR Onnx loader\n\n Args:\n model_path: path or url to onnx model file\n vocab: vocabulary used for encoding\n engine_cfg: configuration for the inference engine\n cfg: configuration dictionary\n **kwargs: additional arguments to be passed to `Engine`\n \"\"\"\n\n _children_names: list[str] = [\"postprocessor\"]\n\n def __init__(\n self,\n model_path: str,\n vocab: str,\n engine_cfg: EngineConfig | None = None,\n cfg: dict[str, Any] | None = None,\n **kwargs: Any,\n ) -> None:\n super().__init__(url=model_path, engine_cfg=engine_cfg, **kwargs)\n\n self.vocab = vocab\n self.cfg = cfg\n\n self.postprocessor = VIPTRPostProcessor(self.vocab)\n\n def __call__(\n self,\n x: np.ndarray,\n return_model_output: bool = False,\n ) -> dict[str, Any]:\n logits = self.run(x)\n\n out: dict[str, Any] = {}\n if return_model_output:\n out[\"out_map\"] = logits\n\n # Post-process\n out[\"preds\"] = self.postprocessor(logits)\n\n return out\n\n\ndef _viptr(\n arch: str,\n model_path: str,\n load_in_8_bit: bool = False,\n engine_cfg: EngineConfig | None = None,\n **kwargs: Any,\n) -> VIPTR:\n if load_in_8_bit:\n logging.warning(\"VIPTR models do not support 8-bit quantization yet. Loading full precision model...\")\n kwargs[\"vocab\"] = kwargs.get(\"vocab\", default_cfgs[arch][\"vocab\"])\n\n _cfg = deepcopy(default_cfgs[arch])\n _cfg[\"vocab\"] = kwargs[\"vocab\"]\n _cfg[\"input_shape\"] = kwargs.get(\"input_shape\", default_cfgs[arch][\"input_shape\"])\n # Patch the url\n model_path = default_cfgs[arch][\"url_8_bit\"] if load_in_8_bit and \"http\" in model_path else model_path\n\n # Build the model\n return VIPTR(model_path, cfg=_cfg, engine_cfg=engine_cfg, **kwargs)\n\n\ndef viptr_tiny(\n model_path: str = default_cfgs[\"viptr_tiny\"][\"url\"],\n load_in_8_bit: bool = False,\n engine_cfg: EngineConfig | None = None,\n **kwargs: Any,\n) -> VIPTR:\n \"\"\"VIPTR as described in `\"A Vision Permutable Extractor for Fast and Efficient\n Scene Text Recognition\" `_.\n\n >>> import numpy as np\n >>> from onnxtr.models import viptr_tiny\n >>> model = viptr_tiny()\n >>> input_tensor = np.random.rand(1, 3, 32, 128)\n >>> out = model(input_tensor)\n\n Args:\n model_path: path to onnx model file, defaults to url in default_cfgs\n load_in_8_bit: whether to load the the 8-bit quantized model, defaults to False\n engine_cfg: configuration for the inference engine\n **kwargs: keyword arguments of the VIPTR architecture\n\n Returns:\n text recognition architecture\n \"\"\"\n return _viptr(\"viptr_tiny\", model_path, load_in_8_bit, engine_cfg, **kwargs)\n"
},
{
"path": "onnxtr/models/recognition/models/vitstr.py",
"content": "# Copyright (C) 2021-2026, Mindee | Felix Dittrich.\n\n# This program is licensed under the Apache License 2.0.\n# See LICENSE or go to for full license details.\n\nfrom copy import deepcopy\nfrom typing import Any\n\nimport numpy as np\nfrom scipy.special import softmax\n\nfrom onnxtr.utils import VOCABS\n\nfrom ...engine import Engine, EngineConfig\nfrom ..core import RecognitionPostProcessor\n\n__all__ = [\"ViTSTR\", \"vitstr_small\", \"vitstr_base\"]\n\ndefault_cfgs: dict[str, dict[str, Any]] = {\n \"vitstr_small\": {\n \"mean\": (0.694, 0.695, 0.693),\n \"std\": (0.299, 0.296, 0.301),\n \"input_shape\": (3, 32, 128),\n \"vocab\": VOCABS[\"french\"],\n \"url\": \"https://github.com/felixdittrich92/OnnxTR/releases/download/v0.0.1/vitstr_small-3ff9c500.onnx\",\n \"url_8_bit\": \"https://github.com/felixdittrich92/OnnxTR/releases/download/v0.1.2/vitstr_small_dynamic_8_bit-bec6c796.onnx\",\n },\n \"vitstr_base\": {\n \"mean\": (0.694, 0.695, 0.693),\n \"std\": (0.299, 0.296, 0.301),\n \"input_shape\": (3, 32, 128),\n \"vocab\": VOCABS[\"french\"],\n \"url\": \"https://github.com/felixdittrich92/OnnxTR/releases/download/v0.0.1/vitstr_base-ff62f5be.onnx\",\n \"url_8_bit\": \"https://github.com/felixdittrich92/OnnxTR/releases/download/v0.1.2/vitstr_base_dynamic_8_bit-976c7cd6.onnx\",\n },\n}\n\n\nclass ViTSTR(Engine):\n \"\"\"ViTSTR Onnx loader\n\n Args:\n model_path: path to onnx model file\n vocab: vocabulary used for encoding\n engine_cfg: configuration for the inference engine\n cfg: dictionary containing information about the model\n **kwargs: additional arguments to be passed to `Engine`\n \"\"\"\n\n def __init__(\n self,\n model_path: str,\n vocab: str,\n engine_cfg: EngineConfig | None = None,\n cfg: dict[str, Any] | None = None,\n **kwargs: Any,\n ) -> None:\n super().__init__(url=model_path, engine_cfg=engine_cfg, **kwargs)\n\n self.vocab = vocab\n self.cfg = cfg\n\n self.postprocessor = ViTSTRPostProcessor(vocab=self.vocab)\n\n def __call__(\n self,\n x: np.ndarray,\n return_model_output: bool = False,\n ) -> dict[str, Any]:\n logits = self.run(x)\n\n out: dict[str, Any] = {}\n if return_model_output:\n out[\"out_map\"] = logits\n\n out[\"preds\"] = self.postprocessor(logits)\n\n return out\n\n\nclass ViTSTRPostProcessor(RecognitionPostProcessor):\n \"\"\"Post processor for ViTSTR architecture\n\n Args:\n vocab: string containing the ordered sequence of supported characters\n \"\"\"\n\n def __init__(\n self,\n vocab: str,\n ) -> None:\n super().__init__(vocab)\n self._embedding = list(vocab) + [\"\", \"\"]\n\n def __call__(self, logits):\n # compute pred with argmax for attention models\n out_idxs = np.argmax(logits, axis=-1)\n preds_prob = softmax(logits, axis=-1).max(axis=-1)\n\n word_values = [\n \"\".join(self._embedding[idx] for idx in encoded_seq).split(\"\")[0] for encoded_seq in out_idxs\n ]\n # compute probabilties for each word up to the EOS token\n probs = [\n preds_prob[i, : len(word)].clip(0, 1).mean().astype(float) if word else 0.0\n for i, word in enumerate(word_values)\n ]\n\n return list(zip(word_values, probs))\n\n\ndef _vitstr(\n arch: str,\n model_path: str,\n load_in_8_bit: bool = False,\n engine_cfg: EngineConfig | None = None,\n **kwargs: Any,\n) -> ViTSTR:\n # Patch the config\n _cfg = deepcopy(default_cfgs[arch])\n _cfg[\"vocab\"] = kwargs.get(\"vocab\", _cfg[\"vocab\"])\n _cfg[\"input_shape\"] = kwargs.get(\"input_shape\", _cfg[\"input_shape\"])\n\n kwargs[\"vocab\"] = _cfg[\"vocab\"]\n # Patch the url\n model_path = default_cfgs[arch][\"url_8_bit\"] if load_in_8_bit and \"http\" in model_path else model_path\n\n # Build the model\n return ViTSTR(model_path, cfg=_cfg, engine_cfg=engine_cfg, **kwargs)\n\n\ndef vitstr_small(\n model_path: str = default_cfgs[\"vitstr_small\"][\"url\"],\n load_in_8_bit: bool = False,\n engine_cfg: EngineConfig | None = None,\n **kwargs: Any,\n) -> ViTSTR:\n \"\"\"ViTSTR-Small as described in `\"Vision Transformer for Fast and Efficient Scene Text Recognition\"\n `_.\n\n >>> import numpy as np\n >>> from onnxtr.models import vitstr_small\n >>> model = vitstr_small()\n >>> input_tensor = np.random.rand(1, 3, 32, 128)\n >>> out = model(input_tensor)\n\n Args:\n model_path: path to onnx model file, defaults to url in default_cfgs\n load_in_8_bit: whether to load the the 8-bit quantized model, defaults to False\n engine_cfg: configuration for the inference engine\n **kwargs: keyword arguments of the ViTSTR architecture\n\n Returns:\n text recognition architecture\n \"\"\"\n return _vitstr(\"vitstr_small\", model_path, load_in_8_bit, engine_cfg, **kwargs)\n\n\ndef vitstr_base(\n model_path: str = default_cfgs[\"vitstr_base\"][\"url\"],\n load_in_8_bit: bool = False,\n engine_cfg: EngineConfig | None = None,\n **kwargs: Any,\n) -> ViTSTR:\n \"\"\"ViTSTR-Base as described in `\"Vision Transformer for Fast and Efficient Scene Text Recognition\"\n `_.\n\n >>> import numpy as np\n >>> from onnxtr.models import vitstr_base\n >>> model = vitstr_base()\n >>> input_tensor = np.random.rand(1, 3, 32, 128)\n >>> out = model(input_tensor)\n\n Args:\n model_path: path to onnx model file, defaults to url in default_cfgs\n load_in_8_bit: whether to load the the 8-bit quantized model, defaults to False\n engine_cfg: configuration for the inference engine\n **kwargs: keyword arguments of the ViTSTR architecture\n\n Returns:\n text recognition architecture\n \"\"\"\n return _vitstr(\"vitstr_base\", model_path, load_in_8_bit, engine_cfg, **kwargs)\n"
},
{
"path": "onnxtr/models/recognition/predictor/__init__.py",
"content": "from .base import *\n"
},
{
"path": "onnxtr/models/recognition/predictor/_utils.py",
"content": "# Copyright (C) 2021-2026, Mindee | Felix Dittrich.\n\n# This program is licensed under the Apache License 2.0.\n# See LICENSE or go to for full license details.\n\n\nimport math\n\nimport numpy as np\n\nfrom ..utils import merge_multi_strings\n\n__all__ = [\"split_crops\", \"remap_preds\"]\n\n\ndef split_crops(\n crops: list[np.ndarray],\n max_ratio: float,\n target_ratio: int,\n split_overlap_ratio: float,\n channels_last: bool = True,\n) -> tuple[list[np.ndarray], list[int | tuple[int, int, float]], bool]:\n \"\"\"\n Split crops horizontally if they exceed a given aspect ratio.\n\n Args:\n crops: List of image crops (H, W, C) if channels_last else (C, H, W).\n max_ratio: Aspect ratio threshold above which crops are split.\n target_ratio: Target aspect ratio after splitting (e.g., 4 for 128x32).\n split_overlap_ratio: Desired overlap between splits (as a fraction of split width).\n channels_last: Whether the crops are in channels-last format.\n\n Returns:\n A tuple containing:\n - The new list of crops (possibly with splits),\n - A mapping indicating how to reassemble predictions,\n - A boolean indicating whether remapping is required.\n \"\"\"\n if split_overlap_ratio <= 0.0 or split_overlap_ratio >= 1.0:\n raise ValueError(f\"Valid range for split_overlap_ratio is (0.0, 1.0), but is: {split_overlap_ratio}\")\n\n remap_required = False\n new_crops: list[np.ndarray] = []\n crop_map: list[int | tuple[int, int, float]] = []\n\n for crop in crops:\n h, w = crop.shape[:2] if channels_last else crop.shape[-2:]\n aspect_ratio = w / h\n\n if aspect_ratio > max_ratio:\n split_width = max(1, math.ceil(h * target_ratio))\n overlap_width = max(0, math.floor(split_width * split_overlap_ratio))\n\n splits, last_overlap = _split_horizontally(crop, split_width, overlap_width, channels_last)\n\n # Remove any empty splits\n splits = [s for s in splits if all(dim > 0 for dim in s.shape)]\n if splits:\n crop_map.append((len(new_crops), len(new_crops) + len(splits), last_overlap))\n new_crops.extend(splits)\n remap_required = True\n else:\n # Fallback: treat it as a single crop\n crop_map.append(len(new_crops))\n new_crops.append(crop)\n else:\n crop_map.append(len(new_crops))\n new_crops.append(crop)\n\n return new_crops, crop_map, remap_required\n\n\ndef _split_horizontally(\n image: np.ndarray, split_width: int, overlap_width: int, channels_last: bool\n) -> tuple[list[np.ndarray], float]:\n \"\"\"\n Horizontally split a single image with overlapping regions.\n\n Args:\n image: The image to split (H, W, C) if channels_last else (C, H, W).\n split_width: Width of each split.\n overlap_width: Width of the overlapping region.\n channels_last: Whether the image is in channels-last format.\n\n Returns:\n - A list of horizontal image slices.\n - The actual overlap ratio of the last split.\n \"\"\"\n image_width = image.shape[1] if channels_last else image.shape[-1]\n if image_width <= split_width:\n return [image], 0.0\n\n # Compute start columns for each split\n step = split_width - overlap_width\n starts = list(range(0, image_width - split_width + 1, step))\n\n # Ensure the last patch reaches the end of the image\n if starts[-1] + split_width < image_width:\n starts.append(image_width - split_width)\n\n splits = []\n for start_col in starts:\n end_col = start_col + split_width\n if channels_last:\n split = image[:, start_col:end_col, :]\n else:\n split = image[:, :, start_col:end_col]\n splits.append(split)\n\n # Calculate the last overlap ratio, if only one split no overlap\n last_overlap = 0\n if len(starts) > 1:\n last_overlap = (starts[-2] + split_width) - starts[-1]\n last_overlap_ratio = last_overlap / split_width if split_width else 0.0\n\n return splits, last_overlap_ratio\n\n\ndef remap_preds(\n preds: list[tuple[str, float]],\n crop_map: list[int | tuple[int, int, float]],\n overlap_ratio: float,\n) -> list[tuple[str, float]]:\n \"\"\"\n Reconstruct predictions from possibly split crops.\n\n Args:\n preds: List of (text, confidence) tuples from each crop.\n crop_map: Map returned by `split_crops`.\n overlap_ratio: Overlap ratio used during splitting.\n\n Returns:\n List of merged (text, confidence) tuples corresponding to original crops.\n \"\"\"\n remapped = []\n for item in crop_map:\n if isinstance(item, int):\n remapped.append(preds[item])\n else:\n start_idx, end_idx, last_overlap = item\n text_parts, confidences = zip(*preds[start_idx:end_idx])\n merged_text = merge_multi_strings(list(text_parts), overlap_ratio, last_overlap)\n merged_conf = sum(confidences) / len(confidences) # average confidence\n remapped.append((merged_text, merged_conf))\n return remapped\n"
},
{
"path": "onnxtr/models/recognition/predictor/base.py",
"content": "# Copyright (C) 2021-2026, Mindee | Felix Dittrich.\n\n# This program is licensed under the Apache License 2.0.\n# See LICENSE or go to for full license details.\n\nfrom collections.abc import Sequence\nfrom typing import Any\n\nimport numpy as np\n\nfrom onnxtr.models.preprocessor import PreProcessor\nfrom onnxtr.utils.repr import NestedObject\n\nfrom ._utils import remap_preds, split_crops\n\n__all__ = [\"RecognitionPredictor\"]\n\n\nclass RecognitionPredictor(NestedObject):\n \"\"\"Implements an object able to identify character sequences in images\n\n Args:\n pre_processor: transform inputs for easier batched model inference\n model: core recognition architecture\n split_wide_crops: wether to use crop splitting for high aspect ratio crops\n \"\"\"\n\n def __init__(\n self,\n pre_processor: PreProcessor,\n model: Any,\n split_wide_crops: bool = True,\n ) -> None:\n super().__init__()\n self.pre_processor = pre_processor\n self.model = model\n self.split_wide_crops = split_wide_crops\n self.critical_ar = 8 # Critical aspect ratio\n self.overlap_ratio = 0.5 # Ratio of overlap between neighboring crops\n self.target_ar = 6 # Target aspect ratio\n\n def __call__(\n self,\n crops: Sequence[np.ndarray],\n **kwargs: Any,\n ) -> list[tuple[str, float]]:\n if len(crops) == 0:\n return []\n # Dimension check\n if any(crop.ndim != 3 for crop in crops):\n raise ValueError(\"incorrect input shape: all crops are expected to be multi-channel 2D images.\")\n\n # Split crops that are too wide\n remapped = False\n if self.split_wide_crops:\n new_crops, crop_map, remapped = split_crops(\n crops, # type: ignore[arg-type]\n self.critical_ar,\n self.target_ar,\n self.overlap_ratio,\n True,\n )\n if remapped:\n crops = new_crops\n\n # Resize & batch them\n processed_batches = self.pre_processor(crops) # type: ignore[arg-type]\n\n # Forward it\n raw = [self.model(batch, **kwargs)[\"preds\"] for batch in processed_batches]\n\n # Process outputs\n out = [charseq for batch in raw for charseq in batch]\n\n # Remap crops\n if self.split_wide_crops and remapped:\n out = remap_preds(out, crop_map, self.overlap_ratio)\n\n return out\n"
},
{
"path": "onnxtr/models/recognition/utils.py",
"content": "# Copyright (C) 2021-2026, Mindee | Felix Dittrich.\n\n# This program is licensed under the Apache License 2.0.\n# See LICENSE or go to for full license details.\n\n\nfrom rapidfuzz.distance import Hamming\n\n__all__ = [\"merge_strings\", \"merge_multi_strings\"]\n\n\ndef merge_strings(a: str, b: str, overlap_ratio: float) -> str:\n \"\"\"Merges 2 character sequences in the best way to maximize the alignment of their overlapping characters.\n\n Args:\n a: first char seq, suffix should be similar to b's prefix.\n b: second char seq, prefix should be similar to a's suffix.\n overlap_ratio: estimated ratio of overlapping characters.\n\n Returns:\n A merged character sequence.\n\n Example::\n >>> from doctr.models.recognition.utils import merge_strings\n >>> merge_strings('abcd', 'cdefgh', 0.5)\n 'abcdefgh'\n >>> merge_strings('abcdi', 'cdefgh', 0.5)\n 'abcdefgh'\n \"\"\"\n seq_len = min(len(a), len(b))\n if seq_len <= 1: # One sequence is empty or will be after cropping in next step, return both to keep data\n return a + b\n\n a_crop, b_crop = a[:-1], b[1:] # Remove last letter of \"a\" and first of \"b\", because they might be cut off\n max_overlap = min(len(a_crop), len(b_crop))\n\n # Compute Hamming distances for all possible overlaps\n scores = [Hamming.distance(a_crop[-i:], b_crop[:i], processor=None) for i in range(1, max_overlap + 1)]\n\n # Find zero-score matches\n zero_matches = [i for i, score in enumerate(scores) if score == 0]\n\n expected_overlap = round(len(b) * overlap_ratio) - 3 # adjust for cropping and index\n\n # Case 1: One perfect match - exactly one zero score - just merge there\n if len(zero_matches) == 1:\n i = zero_matches[0]\n return a_crop + b_crop[i + 1 :]\n\n # Case 2: Multiple perfect matches - likely due to repeated characters.\n # Use the estimated overlap length to choose the match closest to the expected alignment.\n elif len(zero_matches) > 1:\n best_i = min(zero_matches, key=lambda x: abs(x - expected_overlap))\n return a_crop + b_crop[best_i + 1 :]\n\n # Case 3: Absence of zero scores indicates that the same character in the image was recognized differently OR that\n # the overlap was too small and we just need to merge the crops fully\n if expected_overlap < -1:\n return a + b\n elif expected_overlap < 0:\n return a_crop + b_crop\n\n # Find best overlap by minimizing Hamming distance + distance from expected overlap size\n combined_scores = [score + abs(i - expected_overlap) for i, score in enumerate(scores)]\n best_i = combined_scores.index(min(combined_scores))\n return a_crop + b_crop[best_i + 1 :]\n\n\ndef merge_multi_strings(seq_list: list[str], overlap_ratio: float, last_overlap_ratio: float) -> str:\n \"\"\"\n Merges consecutive string sequences with overlapping characters.\n\n Args:\n seq_list: list of sequences to merge. Sequences need to be ordered from left to right.\n overlap_ratio: Estimated ratio of overlapping letters between neighboring strings.\n last_overlap_ratio: Estimated ratio of overlapping letters for the last element in seq_list.\n\n Returns:\n A merged character sequence\n\n Example::\n >>> from doctr.models.recognition.utils import merge_multi_strings\n >>> merge_multi_strings(['abc', 'bcdef', 'difghi', 'aijkl'], 0.5, 0.1)\n 'abcdefghijkl'\n \"\"\"\n if not seq_list:\n return \"\"\n result = seq_list[0]\n for i in range(1, len(seq_list)):\n text_b = seq_list[i]\n ratio = last_overlap_ratio if i == len(seq_list) - 1 else overlap_ratio\n result = merge_strings(result, text_b, ratio)\n return result\n"
},
{
"path": "onnxtr/models/recognition/zoo.py",
"content": "# Copyright (C) 2021-2026, Mindee | Felix Dittrich.\n\n# This program is licensed under the Apache License 2.0.\n# See LICENSE or go to for full license details.\n\nfrom typing import Any\n\nfrom .. import recognition\nfrom ..engine import EngineConfig\nfrom ..preprocessor import PreProcessor\nfrom .predictor import RecognitionPredictor\n\n__all__ = [\"recognition_predictor\"]\n\n\nARCHS: list[str] = [\n \"crnn_vgg16_bn\",\n \"crnn_mobilenet_v3_small\",\n \"crnn_mobilenet_v3_large\",\n \"sar_resnet31\",\n \"master\",\n \"vitstr_small\",\n \"vitstr_base\",\n \"parseq\",\n \"viptr_tiny\",\n]\n\n\ndef _predictor(\n arch: Any, load_in_8_bit: bool = False, engine_cfg: EngineConfig | None = None, **kwargs: Any\n) -> RecognitionPredictor:\n if isinstance(arch, str):\n if arch not in ARCHS:\n raise ValueError(f\"unknown architecture '{arch}'\")\n\n _model = recognition.__dict__[arch](load_in_8_bit=load_in_8_bit, engine_cfg=engine_cfg)\n else:\n if not isinstance(\n arch,\n (\n recognition.CRNN,\n recognition.SAR,\n recognition.MASTER,\n recognition.ViTSTR,\n recognition.PARSeq,\n recognition.VIPTR,\n ),\n ):\n raise ValueError(f\"unknown architecture: {type(arch)}\")\n _model = arch\n\n kwargs[\"mean\"] = kwargs.get(\"mean\", _model.cfg[\"mean\"])\n kwargs[\"std\"] = kwargs.get(\"std\", _model.cfg[\"std\"])\n kwargs[\"batch_size\"] = kwargs.get(\"batch_size\", 1024)\n input_shape = _model.cfg[\"input_shape\"][1:]\n predictor = RecognitionPredictor(PreProcessor(input_shape, preserve_aspect_ratio=True, **kwargs), _model)\n\n return predictor\n\n\ndef recognition_predictor(\n arch: Any = \"crnn_vgg16_bn\",\n symmetric_pad: bool = False,\n batch_size: int = 128,\n load_in_8_bit: bool = False,\n engine_cfg: EngineConfig | None = None,\n **kwargs: Any,\n) -> RecognitionPredictor:\n \"\"\"Text recognition architecture.\n\n Example::\n >>> import numpy as np\n >>> from onnxtr.models import recognition_predictor\n >>> model = recognition_predictor()\n >>> input_page = (255 * np.random.rand(32, 128, 3)).astype(np.uint8)\n >>> out = model([input_page])\n\n Args:\n arch: name of the architecture or model itself to use (e.g. 'crnn_vgg16_bn')\n symmetric_pad: if True, pad the image symmetrically instead of padding at the bottom-right\n batch_size: number of samples the model processes in parallel\n load_in_8_bit: whether to load the the 8-bit quantized model, defaults to False\n engine_cfg: configuration of inference engine\n **kwargs: optional parameters to be passed to the architecture\n\n Returns:\n Recognition predictor\n \"\"\"\n return _predictor(\n arch=arch,\n symmetric_pad=symmetric_pad,\n batch_size=batch_size,\n load_in_8_bit=load_in_8_bit,\n engine_cfg=engine_cfg,\n **kwargs,\n )\n"
},
{
"path": "onnxtr/models/zoo.py",
"content": "# Copyright (C) 2021-2026, Mindee | Felix Dittrich.\n\n# This program is licensed under the Apache License 2.0.\n# See LICENSE or go to for full license details.\n\nfrom typing import Any\n\nfrom .detection.zoo import detection_predictor\nfrom .engine import EngineConfig\nfrom .predictor import OCRPredictor\nfrom .recognition.zoo import recognition_predictor\n\n__all__ = [\"ocr_predictor\"]\n\n\ndef _predictor(\n det_arch: Any,\n reco_arch: Any,\n assume_straight_pages: bool = True,\n preserve_aspect_ratio: bool = True,\n symmetric_pad: bool = True,\n det_bs: int = 2,\n reco_bs: int = 512,\n detect_orientation: bool = False,\n straighten_pages: bool = False,\n detect_language: bool = False,\n load_in_8_bit: bool = False,\n det_engine_cfg: EngineConfig | None = None,\n reco_engine_cfg: EngineConfig | None = None,\n clf_engine_cfg: EngineConfig | None = None,\n **kwargs,\n) -> OCRPredictor:\n # Detection\n det_predictor = detection_predictor(\n det_arch,\n batch_size=det_bs,\n assume_straight_pages=assume_straight_pages,\n preserve_aspect_ratio=preserve_aspect_ratio,\n symmetric_pad=symmetric_pad,\n load_in_8_bit=load_in_8_bit,\n engine_cfg=det_engine_cfg,\n )\n\n # Recognition\n reco_predictor = recognition_predictor(\n reco_arch,\n batch_size=reco_bs,\n load_in_8_bit=load_in_8_bit,\n engine_cfg=reco_engine_cfg,\n )\n\n return OCRPredictor(\n det_predictor,\n reco_predictor,\n assume_straight_pages=assume_straight_pages,\n preserve_aspect_ratio=preserve_aspect_ratio,\n symmetric_pad=symmetric_pad,\n detect_orientation=detect_orientation,\n straighten_pages=straighten_pages,\n detect_language=detect_language,\n clf_engine_cfg=clf_engine_cfg,\n **kwargs,\n )\n\n\ndef ocr_predictor(\n det_arch: Any = \"fast_base\",\n reco_arch: Any = \"crnn_vgg16_bn\",\n assume_straight_pages: bool = True,\n preserve_aspect_ratio: bool = True,\n symmetric_pad: bool = True,\n export_as_straight_boxes: bool = False,\n detect_orientation: bool = False,\n straighten_pages: bool = False,\n detect_language: bool = False,\n load_in_8_bit: bool = False,\n det_engine_cfg: EngineConfig | None = None,\n reco_engine_cfg: EngineConfig | None = None,\n clf_engine_cfg: EngineConfig | None = None,\n **kwargs: Any,\n) -> OCRPredictor:\n \"\"\"End-to-end OCR architecture using one model for localization, and another for text recognition.\n\n >>> import numpy as np\n >>> from onnxtr.models import ocr_predictor\n >>> model = ocr_predictor('db_resnet50', 'crnn_vgg16_bn')\n >>> input_page = (255 * np.random.rand(600, 800, 3)).astype(np.uint8)\n >>> out = model([input_page])\n\n Args:\n det_arch: name of the detection architecture or the model itself to use\n (e.g. 'db_resnet50', 'db_mobilenet_v3_large')\n reco_arch: name of the recognition architecture or the model itself to use\n (e.g. 'crnn_vgg16_bn', 'sar_resnet31')\n assume_straight_pages: if True, speeds up the inference by assuming you only pass straight pages\n without rotated textual elements.\n preserve_aspect_ratio: If True, pad the input document image to preserve the aspect ratio before\n running the detection model on it.\n symmetric_pad: if True, pad the image symmetrically instead of padding at the bottom-right.\n export_as_straight_boxes: when assume_straight_pages is set to False, export final predictions\n (potentially rotated) as straight bounding boxes.\n detect_orientation: if True, the estimated general page orientation will be added to the predictions for each\n page. Doing so will slightly deteriorate the overall latency.\n straighten_pages: if True, estimates the page general orientation\n based on the segmentation map median line orientation.\n Then, rotates page before passing it again to the deep learning detection module.\n Doing so will improve performances for documents with page-uniform rotations.\n detect_language: if True, the language prediction will be added to the predictions for each\n page. Doing so will slightly deteriorate the overall latency.\n load_in_8_bit: whether to load the the 8-bit quantized model, defaults to False\n det_engine_cfg: configuration of the detection engine\n reco_engine_cfg: configuration of the recognition engine\n clf_engine_cfg: configuration of the orientation classification engine\n kwargs: keyword args of `OCRPredictor`\n\n Returns:\n OCR predictor\n \"\"\"\n return _predictor(\n det_arch,\n reco_arch,\n assume_straight_pages=assume_straight_pages,\n preserve_aspect_ratio=preserve_aspect_ratio,\n symmetric_pad=symmetric_pad,\n export_as_straight_boxes=export_as_straight_boxes,\n detect_orientation=detect_orientation,\n straighten_pages=straighten_pages,\n detect_language=detect_language,\n load_in_8_bit=load_in_8_bit,\n det_engine_cfg=det_engine_cfg,\n reco_engine_cfg=reco_engine_cfg,\n clf_engine_cfg=clf_engine_cfg,\n **kwargs,\n )\n"
},
{
"path": "onnxtr/py.typed",
"content": ""
},
{
"path": "onnxtr/transforms/__init__.py",
"content": "from .base import *\n"
},
{
"path": "onnxtr/transforms/base.py",
"content": "# Copyright (C) 2021-2026, Mindee | Felix Dittrich.\n\n# This program is licensed under the Apache License 2.0.\n# See LICENSE or go to for full license details.\n\n\nimport math\n\nimport numpy as np\nfrom PIL import Image, ImageOps\n\n__all__ = [\"Resize\", \"Normalize\"]\n\n\nclass Resize:\n \"\"\"Resize the input image to the given size\n\n Args:\n size: the target size of the image\n interpolation: the interpolation method to use\n preserve_aspect_ratio: whether to preserve the aspect ratio of the image\n symmetric_pad: whether to symmetrically pad the image\n \"\"\"\n\n def __init__(\n self,\n size: int | tuple[int, int],\n interpolation=Image.Resampling.BILINEAR,\n preserve_aspect_ratio: bool = False,\n symmetric_pad: bool = False,\n ) -> None:\n self.size = size if isinstance(size, tuple) else (size, size)\n self.interpolation = interpolation\n self.preserve_aspect_ratio = preserve_aspect_ratio\n self.symmetric_pad = symmetric_pad\n self.output_size = size if isinstance(size, tuple) else (size, size)\n\n if not isinstance(self.size, (tuple, int)):\n raise AssertionError(\"size should be either a tuple or an int\")\n\n def __call__(self, img: np.ndarray) -> np.ndarray:\n if img.dtype != np.uint8:\n img_pil = Image.fromarray((img * 255).clip(0, 255).astype(np.uint8))\n else:\n img_pil = Image.fromarray(img)\n\n sh, sw = self.size\n w, h = img_pil.size\n\n if not self.preserve_aspect_ratio:\n img_resized_pil = img_pil.resize((sw, sh), resample=self.interpolation)\n return np.array(img_resized_pil)\n\n actual_ratio = h / w\n target_ratio = sh / sw\n\n if actual_ratio > target_ratio:\n new_h = sh\n new_w = max(int(sh / actual_ratio), 1)\n else:\n new_w = sw\n new_h = max(int(sw * actual_ratio), 1)\n\n img_resized_pil = img_pil.resize((new_w, new_h), resample=self.interpolation)\n\n delta_w = sw - new_w\n delta_h = sh - new_h\n\n if self.symmetric_pad:\n # Symmetric padding\n pad_left = math.ceil(delta_w / 2)\n pad_right = math.floor(delta_w / 2)\n pad_top = math.ceil(delta_h / 2)\n pad_bottom = math.floor(delta_h / 2)\n else:\n # Asymmetric padding\n pad_left, pad_top = 0, 0\n pad_right, pad_bottom = delta_w, delta_h\n\n img_padded_pil = ImageOps.expand(\n img_resized_pil,\n border=(pad_left, pad_top, pad_right, pad_bottom),\n fill=0,\n )\n\n return np.array(img_padded_pil)\n\n def __repr__(self) -> str:\n interpolate_str = self.interpolation\n _repr = f\"output_size={self.size}, interpolation='{interpolate_str}'\"\n if self.preserve_aspect_ratio:\n _repr += f\", preserve_aspect_ratio={self.preserve_aspect_ratio}, symmetric_pad={self.symmetric_pad}\"\n return f\"{self.__class__.__name__}({_repr})\"\n\n\nclass Normalize:\n \"\"\"Normalize the input image\n\n Args:\n mean: mean values to subtract\n std: standard deviation values to divide\n \"\"\"\n\n def __init__(\n self,\n mean: float | tuple[float, float, float] = (0.485, 0.456, 0.406),\n std: float | tuple[float, float, float] = (0.229, 0.224, 0.225),\n ) -> None:\n self.mean = mean\n self.std = std\n\n if not isinstance(self.mean, (float, tuple, list)):\n raise AssertionError(\"mean should be either a tuple, a list or a float\")\n if not isinstance(self.std, (float, tuple, list)):\n raise AssertionError(\"std should be either a tuple, a list or a float\")\n\n def __call__(\n self,\n img: np.ndarray,\n ) -> np.ndarray:\n # Normalize image\n return (img - np.array(self.mean).astype(img.dtype)) / np.array(self.std).astype(img.dtype)\n\n def __repr__(self) -> str:\n _repr = f\"mean={self.mean}, std={self.std}\"\n return f\"{self.__class__.__name__}({_repr})\"\n"
},
{
"path": "onnxtr/utils/__init__.py",
"content": "from .common_types import *\nfrom .data import *\nfrom .geometry import *\nfrom .vocabs import *\n"
},
{
"path": "onnxtr/utils/common_types.py",
"content": "# Copyright (C) 2021-2026, Mindee | Felix Dittrich.\n\n# This program is licensed under the Apache License 2.0.\n# See LICENSE or go to for full license details.\n\nfrom pathlib import Path\n\n__all__ = [\"Point2D\", \"BoundingBox\", \"Polygon4P\", \"Polygon\", \"Bbox\"]\n\n\nPoint2D = tuple[float, float]\nBoundingBox = tuple[Point2D, Point2D]\nPolygon4P = tuple[Point2D, Point2D, Point2D, Point2D]\nPolygon = list[Point2D]\nAbstractPath = str | Path\nAbstractFile = AbstractPath | bytes\nBbox = tuple[float, float, float, float]\n"
},
{
"path": "onnxtr/utils/data.py",
"content": "# Copyright (C) 2021-2026, Mindee | Felix Dittrich.\n\n# This program is licensed under the Apache License 2.0.\n# See LICENSE or go to for full license details.\n\n# Adapted from https://github.com/pytorch/vision/blob/master/torchvision/datasets/utils.py\n\nimport hashlib\nimport logging\nimport os\nimport re\nimport urllib.error\nimport urllib.request\nfrom pathlib import Path\n\nfrom tqdm.auto import tqdm\n\n__all__ = [\"download_from_url\"]\n\n\n# matches bfd8deac from resnet18-bfd8deac.ckpt\nHASH_REGEX = re.compile(r\"-([a-f0-9]*)\\.\")\nUSER_AGENT = \"felixdittrich92/OnnxTR\"\n\n\ndef _urlretrieve(url: str, filename: Path | str, chunk_size: int = 1024) -> None:\n with open(filename, \"wb\") as fh:\n with urllib.request.urlopen(urllib.request.Request(url, headers={\"User-Agent\": USER_AGENT})) as response:\n with tqdm(total=response.length) as pbar:\n for chunk in iter(lambda: response.read(chunk_size), \"\"):\n if not chunk:\n break\n pbar.update(chunk_size)\n fh.write(chunk)\n\n\ndef _check_integrity(file_path: str | Path, hash_prefix: str) -> bool:\n with open(file_path, \"rb\") as f:\n sha_hash = hashlib.sha256(f.read()).hexdigest()\n\n return sha_hash[: len(hash_prefix)] == hash_prefix\n\n\ndef download_from_url(\n url: str,\n file_name: str | None = None,\n hash_prefix: str | None = None,\n cache_dir: str | None = None,\n cache_subdir: str | None = None,\n) -> Path:\n \"\"\"Download a file using its URL\n\n >>> from onnxtr.models import download_from_url\n >>> download_from_url(\"https://yoursource.com/yourcheckpoint-yourhash.zip\")\n\n Args:\n url: the URL of the file to download\n file_name: optional name of the file once downloaded\n hash_prefix: optional expected SHA256 hash of the file\n cache_dir: cache directory\n cache_subdir: subfolder to use in the cache\n\n Returns:\n the location of the downloaded file\n\n Note:\n You can change cache directory location by using `ONNXTR_CACHE_DIR` environment variable.\n \"\"\"\n if not isinstance(file_name, str):\n file_name = url.rpartition(\"/\")[-1].split(\"&\")[0]\n\n cache_dir = (\n str(os.environ.get(\"ONNXTR_CACHE_DIR\", os.path.join(os.path.expanduser(\"~\"), \".cache\", \"onnxtr\")))\n if cache_dir is None\n else cache_dir\n )\n\n # Check hash in file name\n if hash_prefix is None:\n r = HASH_REGEX.search(file_name)\n hash_prefix = r.group(1) if r else None\n\n folder_path = Path(cache_dir) if cache_subdir is None else Path(cache_dir, cache_subdir)\n file_path = folder_path.joinpath(file_name)\n # Check file existence\n if file_path.is_file() and (hash_prefix is None or _check_integrity(file_path, hash_prefix)):\n logging.info(f\"Using downloaded & verified file: {file_path}\")\n return file_path\n\n try:\n # Create folder hierarchy\n folder_path.mkdir(parents=True, exist_ok=True)\n except OSError:\n error_message = f\"Failed creating cache direcotry at {folder_path}\"\n if os.environ.get(\"ONNXTR_CACHE_DIR\", \"\"):\n error_message += \" using path from 'ONNXTR_CACHE_DIR' environment variable.\"\n else:\n error_message += (\n \". You can change default cache directory using 'ONNXTR_CACHE_DIR' environment variable if needed.\"\n )\n logging.error(error_message)\n raise\n # Download the file\n try:\n print(f\"Downloading {url} to {file_path}\")\n _urlretrieve(url, file_path)\n except (urllib.error.URLError, IOError) as e: # pragma: no cover\n if url[:5] == \"https\":\n url = url.replace(\"https:\", \"http:\")\n print(f\"Failed download. Trying https -> http instead. Downloading {url} to {file_path}\")\n _urlretrieve(url, file_path)\n else:\n raise e\n\n # Remove corrupted files\n if isinstance(hash_prefix, str) and not _check_integrity(file_path, hash_prefix): # pragma: no cover\n # Remove file\n os.remove(file_path)\n raise ValueError(f\"corrupted download, the hash of {url} does not match its expected value\")\n\n return file_path\n"
},
{
"path": "onnxtr/utils/fonts.py",
"content": "# Copyright (C) 2021-2026, Mindee | Felix Dittrich.\n\n# This program is licensed under the Apache License 2.0.\n# See LICENSE or go to for full license details.\n\nimport logging\nimport platform\n\nfrom PIL import ImageFont\n\n__all__ = [\"get_font\"]\n\n\ndef get_font(font_family: str | None = None, font_size: int = 13) -> ImageFont.FreeTypeFont | ImageFont.ImageFont:\n \"\"\"Resolves a compatible ImageFont for the system\n\n Args:\n font_family: the font family to use\n font_size: the size of the font upon rendering\n\n Returns:\n the Pillow font\n \"\"\"\n # Font selection\n if font_family is None:\n try:\n font = ImageFont.truetype(\"FreeMono.ttf\" if platform.system() == \"Linux\" else \"Arial.ttf\", font_size)\n except OSError: # pragma: no cover\n font = ImageFont.load_default() # type: ignore[assignment]\n logging.warning(\n \"unable to load recommended font family. Loading default PIL font,\"\n \"font size issues may be expected.\"\n \"To prevent this, it is recommended to specify the value of 'font_family'.\"\n )\n else: # pragma: no cover\n font = ImageFont.truetype(font_family, font_size)\n\n return font\n"
},
{
"path": "onnxtr/utils/geometry.py",
"content": "# Copyright (C) 2021-2026, Mindee | Felix Dittrich.\n\n# This program is licensed under the Apache License 2.0.\n# See LICENSE or go to for full license details.\n\nfrom copy import deepcopy\nfrom math import ceil\n\nimport cv2\nimport numpy as np\n\nfrom .common_types import BoundingBox, Polygon4P\n\n__all__ = [\n \"bbox_to_polygon\",\n \"polygon_to_bbox\",\n \"order_points\",\n \"resolve_enclosing_bbox\",\n \"resolve_enclosing_rbbox\",\n \"rotate_boxes\",\n \"compute_expanded_shape\",\n \"rotate_image\",\n \"estimate_page_angle\",\n \"convert_to_relative_coords\",\n \"rotate_abs_geoms\",\n \"extract_crops\",\n \"extract_rcrops\",\n \"shape_translate\",\n \"detach_scores\",\n]\n\n\ndef bbox_to_polygon(bbox: BoundingBox) -> Polygon4P:\n \"\"\"Convert a bounding box to a polygon\n\n Args:\n bbox: a bounding box\n\n Returns:\n a polygon\n \"\"\"\n return bbox[0], (bbox[1][0], bbox[0][1]), (bbox[0][0], bbox[1][1]), bbox[1]\n\n\ndef polygon_to_bbox(polygon: Polygon4P) -> BoundingBox:\n \"\"\"Convert a polygon to a bounding box\n\n Args:\n polygon: a polygon\n\n Returns:\n a bounding box\n \"\"\"\n x, y = zip(*polygon)\n return (min(x), min(y)), (max(x), max(y))\n\n\ndef order_points(pts: np.ndarray) -> np.ndarray:\n \"\"\"Order points in the following order: top-left, top-right, bottom-right, bottom-left\n\n Args:\n pts: array of shape (4, 2) or (4,) with the coordinates of the points\n\n Returns:\n ordered points in the following order: top-left, top-right, bottom-right, bottom-left\n \"\"\"\n pts = np.asarray(pts)\n\n # (xmin, ymin, xmax, ymax)\n if pts.shape == (4,):\n xmin, ymin, xmax, ymax = pts\n return np.array(\n [\n [xmin, ymin], # top-left\n [xmax, ymin], # top-right\n [xmax, ymax], # bottom-right\n [xmin, ymax], # bottom-left\n ],\n dtype=pts.dtype,\n )\n\n # (4, 2) quadrangle\n if pts.shape == (4, 2):\n c = pts.mean(axis=0)\n\n # compute angle of each point around centroid\n angles = np.arctan2(pts[:, 1] - c[1], pts[:, 0] - c[0])\n\n # sort by angle (counter-clockwise ordering)\n pts = pts[np.argsort(angles)]\n\n # ensure consistent starting point (top-left)\n start_idx = np.argmin(pts.sum(axis=1))\n pts = np.roll(pts, -start_idx, axis=0)\n\n # ensure order is TL, TR, BR, BL (clockwise)\n def area(poly):\n return 0.5 * np.sum(poly[:, 0] * np.roll(poly[:, 1], -1) - poly[:, 1] * np.roll(poly[:, 0], -1))\n\n if area(pts) < 0:\n pts = np.roll(pts[::-1], 1, axis=0)\n\n return pts.astype(pts.dtype)\n\n raise ValueError(f\"Unsupported shape {pts.shape}, expected (4,) or (4,2)\")\n\n\ndef detach_scores(boxes: list[np.ndarray]) -> tuple[list[np.ndarray], list[np.ndarray]]:\n \"\"\"Detach the objectness scores from box predictions\n\n Args:\n boxes: list of arrays with boxes of shape (N, 5) or (N, 5, 2)\n\n Returns:\n a tuple of two lists: the first one contains the boxes without the objectness scores,\n the second one contains the objectness scores\n \"\"\"\n\n def _detach(boxes: np.ndarray) -> tuple[np.ndarray, np.ndarray]:\n if boxes.ndim == 2:\n return boxes[:, :-1], boxes[:, -1]\n return boxes[:, :-1], boxes[:, -1, -1]\n\n loc_preds, obj_scores = zip(*(_detach(box) for box in boxes))\n return list(loc_preds), list(obj_scores)\n\n\ndef shape_translate(data: np.ndarray, format: str) -> np.ndarray:\n \"\"\"Translate the shape of the input data to the desired format\n\n Args:\n data: input data in shape (B, C, H, W) or (B, H, W, C) or (C, H, W) or (H, W, C)\n format: target format ('BCHW', 'BHWC', 'CHW', or 'HWC')\n\n Returns:\n the reshaped data\n \"\"\"\n # Get the current shape\n current_shape = data.shape\n\n # Check the number of dimensions\n num_dims = len(current_shape)\n\n if num_dims != len(format):\n return data\n\n if format == \"BCHW\" and data.shape[1] in [1, 3]:\n return data\n elif format == \"BHWC\" and data.shape[-1] in [1, 3]:\n return data\n elif format == \"CHW\" and data.shape[0] in [1, 3]:\n return data\n elif format == \"HWC\" and data.shape[-1] in [1, 3]:\n return data\n elif format == \"BCHW\" and data.shape[1] not in [1, 3]:\n return np.moveaxis(data, -1, 1)\n elif format == \"BHWC\" and data.shape[-1] not in [1, 3]:\n return np.moveaxis(data, 1, -1)\n elif format == \"CHW\" and data.shape[0] not in [1, 3]:\n return np.moveaxis(data, -1, 0)\n elif format == \"HWC\" and data.shape[-1] not in [1, 3]:\n return np.moveaxis(data, 0, -1)\n else:\n return data\n\n\ndef resolve_enclosing_bbox(bboxes: list[BoundingBox] | np.ndarray) -> BoundingBox | np.ndarray:\n \"\"\"Compute enclosing bbox either from:\n\n Args:\n bboxes: boxes in one of the following formats:\n\n - an array of boxes: (*, 4), where boxes have this shape:\n (xmin, ymin, xmax, ymax)\n\n - a list of BoundingBox\n\n Returns:\n a (1, 4) array (enclosing boxarray), or a BoundingBox\n \"\"\"\n if isinstance(bboxes, np.ndarray):\n xmin, ymin, xmax, ymax = np.split(bboxes, 4, axis=1)\n return np.array([xmin.min(), ymin.min(), xmax.max(), ymax.max()])\n else:\n x, y = zip(*[point for box in bboxes for point in box])\n return (min(x), min(y)), (max(x), max(y))\n\n\ndef resolve_enclosing_rbbox(rbboxes: list[np.ndarray], intermed_size: int = 1024) -> np.ndarray:\n \"\"\"Compute enclosing rotated bbox either from:\n\n Args:\n rbboxes: boxes in one of the following formats:\n\n - an array of boxes: (*, 4, 2), where boxes have this shape:\n (x1, y1), (x2, y2), (x3, y3), (x4, y4)\n\n - a list of BoundingBox\n intermed_size: size of the intermediate image\n\n Returns:\n a (4, 2) array (enclosing rotated box)\n \"\"\"\n cloud: np.ndarray = np.concatenate(rbboxes, axis=0)\n # Convert to absolute for minAreaRect\n rect = cv2.minAreaRect(cloud.astype(np.float32) * intermed_size)\n return order_points(cv2.boxPoints(rect) / intermed_size)\n\n\ndef rotate_abs_points(points: np.ndarray, angle: float = 0.0) -> np.ndarray:\n \"\"\"Rotate points counter-clockwise.\n\n Args:\n points: array of size (N, 2)\n angle: angle between -90 and +90 degrees\n\n Returns:\n Rotated points\n \"\"\"\n angle_rad = angle * np.pi / 180.0 # compute radian angle for np functions\n rotation_mat = np.array(\n [[np.cos(angle_rad), -np.sin(angle_rad)], [np.sin(angle_rad), np.cos(angle_rad)]], dtype=points.dtype\n )\n return np.matmul(points, rotation_mat.T)\n\n\ndef compute_expanded_shape(img_shape: tuple[int, int], angle: float) -> tuple[int, int]:\n \"\"\"Compute the shape of an expanded rotated image\n\n Args:\n img_shape: the height and width of the image\n angle: angle between -90 and +90 degrees\n\n Returns:\n the height and width of the rotated image\n \"\"\"\n points: np.ndarray = np.array([\n [img_shape[1] / 2, img_shape[0] / 2],\n [-img_shape[1] / 2, img_shape[0] / 2],\n ])\n\n rotated_points = rotate_abs_points(points, angle)\n\n wh_shape = 2 * np.abs(rotated_points).max(axis=0)\n return wh_shape[1], wh_shape[0]\n\n\ndef rotate_abs_geoms(\n geoms: np.ndarray,\n angle: float,\n img_shape: tuple[int, int],\n expand: bool = True,\n) -> np.ndarray:\n \"\"\"Rotate a batch of bounding boxes or polygons by an angle around the\n image center.\n\n Args:\n geoms: (N, 4) or (N, 4, 2) array of ABSOLUTE coordinate boxes\n angle: anti-clockwise rotation angle in degrees\n img_shape: the height and width of the image\n expand: whether the image should be padded to avoid information loss\n\n Returns:\n A batch of rotated polygons (N, 4, 2)\n \"\"\"\n # Switch to polygons\n polys = (\n np.stack([geoms[:, [0, 1]], geoms[:, [2, 1]], geoms[:, [2, 3]], geoms[:, [0, 3]]], axis=1)\n if geoms.ndim == 2\n else geoms\n )\n polys = polys.astype(np.float32)\n\n # Switch to image center as referential\n polys[..., 0] -= img_shape[1] / 2\n polys[..., 1] = img_shape[0] / 2 - polys[..., 1]\n\n # Rotated them around image center\n rotated_polys = rotate_abs_points(polys.reshape(-1, 2), angle).reshape(-1, 4, 2)\n # Switch back to top-left corner as referential\n target_shape = compute_expanded_shape(img_shape, angle) if expand else img_shape\n # Clip coords to fit since there is no expansion\n rotated_polys[..., 0] = (rotated_polys[..., 0] + target_shape[1] / 2).clip(0, target_shape[1])\n rotated_polys[..., 1] = (target_shape[0] / 2 - rotated_polys[..., 1]).clip(0, target_shape[0])\n\n return rotated_polys\n\n\ndef remap_boxes(loc_preds: np.ndarray, orig_shape: tuple[int, int], dest_shape: tuple[int, int]) -> np.ndarray:\n \"\"\"Remaps a batch of rotated locpred (N, 4, 2) expressed for an origin_shape to a destination_shape.\n This does not impact the absolute shape of the boxes, but allow to calculate the new relative RotatedBbox\n coordinates after a resizing of the image.\n\n Args:\n loc_preds: (N, 4, 2) array of RELATIVE loc_preds\n orig_shape: shape of the origin image\n dest_shape: shape of the destination image\n\n Returns:\n A batch of rotated loc_preds (N, 4, 2) expressed in the destination referencial\n \"\"\"\n if len(dest_shape) != 2:\n raise ValueError(f\"Mask length should be 2, was found at: {len(dest_shape)}\")\n if len(orig_shape) != 2:\n raise ValueError(f\"Image_shape length should be 2, was found at: {len(orig_shape)}\")\n orig_height, orig_width = orig_shape\n dest_height, dest_width = dest_shape\n mboxes = loc_preds.copy()\n mboxes[:, :, 0] = ((loc_preds[:, :, 0] * orig_width) + (dest_width - orig_width) / 2) / dest_width\n mboxes[:, :, 1] = ((loc_preds[:, :, 1] * orig_height) + (dest_height - orig_height) / 2) / dest_height\n\n return mboxes\n\n\ndef rotate_boxes(\n loc_preds: np.ndarray,\n angle: float,\n orig_shape: tuple[int, int],\n min_angle: float = 1.0,\n target_shape: tuple[int, int] | None = None,\n) -> np.ndarray:\n \"\"\"Rotate a batch of straight bounding boxes (xmin, ymin, xmax, ymax, c) or rotated bounding boxes\n (4, 2) of an angle, if angle > min_angle, around the center of the page.\n If target_shape is specified, the boxes are remapped to the target shape after the rotation. This\n is done to remove the padding that is created by rotate_page(expand=True)\n\n Args:\n loc_preds: (N, 4) or (N, 4, 2) array of RELATIVE boxes\n angle: angle between -90 and +90 degrees\n orig_shape: shape of the origin image\n min_angle: minimum angle to rotate boxes\n target_shape: shape of the destination image\n\n Returns:\n A batch of rotated boxes (N, 4, 2): or a batch of straight bounding boxes\n \"\"\"\n # Change format of the boxes to rotated boxes\n _boxes = loc_preds.copy()\n if _boxes.ndim == 2:\n _boxes = np.stack(\n [\n _boxes[:, [0, 1]],\n _boxes[:, [2, 1]],\n _boxes[:, [2, 3]],\n _boxes[:, [0, 3]],\n ],\n axis=1,\n )\n # If small angle, return boxes (no rotation)\n if abs(angle) < min_angle or abs(angle) > 90 - min_angle:\n return _boxes\n # Compute rotation matrix\n angle_rad = angle * np.pi / 180.0 # compute radian angle for np functions\n rotation_mat = np.array(\n [[np.cos(angle_rad), -np.sin(angle_rad)], [np.sin(angle_rad), np.cos(angle_rad)]], dtype=_boxes.dtype\n )\n # Rotate absolute points\n points: np.ndarray = np.stack((_boxes[:, :, 0] * orig_shape[1], _boxes[:, :, 1] * orig_shape[0]), axis=-1)\n image_center = (orig_shape[1] / 2, orig_shape[0] / 2)\n rotated_points = image_center + np.matmul(points - image_center, rotation_mat)\n rotated_boxes: np.ndarray = np.stack(\n (rotated_points[:, :, 0] / orig_shape[1], rotated_points[:, :, 1] / orig_shape[0]), axis=-1\n )\n\n # Apply a mask if requested\n if target_shape is not None:\n rotated_boxes = remap_boxes(rotated_boxes, orig_shape=orig_shape, dest_shape=target_shape)\n\n return rotated_boxes\n\n\ndef rotate_image(\n image: np.ndarray,\n angle: float,\n expand: bool = False,\n preserve_origin_shape: bool = False,\n) -> np.ndarray:\n \"\"\"Rotate an image counterclockwise by an given angle.\n\n Args:\n image: numpy tensor to rotate\n angle: rotation angle in degrees, between -90 and +90\n expand: whether the image should be padded before the rotation\n preserve_origin_shape: if expand is set to True, resizes the final output to the original image size\n\n Returns:\n Rotated array, padded by 0 by default.\n \"\"\"\n # Compute the expanded padding\n exp_img: np.ndarray\n if expand:\n exp_shape = compute_expanded_shape(image.shape[:2], angle)\n h_pad, w_pad = (\n int(max(0, ceil(exp_shape[0] - image.shape[0]))),\n int(max(0, ceil(exp_shape[1] - image.shape[1]))),\n )\n exp_img = np.pad(image, ((h_pad // 2, h_pad - h_pad // 2), (w_pad // 2, w_pad - w_pad // 2), (0, 0)))\n else:\n exp_img = image\n\n height, width = exp_img.shape[:2]\n rot_mat = cv2.getRotationMatrix2D((width / 2, height / 2), angle, 1.0)\n rot_img = cv2.warpAffine(exp_img, rot_mat, (width, height))\n if expand:\n # Pad to get the same aspect ratio\n if (image.shape[0] / image.shape[1]) != (rot_img.shape[0] / rot_img.shape[1]):\n # Pad width\n if (rot_img.shape[0] / rot_img.shape[1]) > (image.shape[0] / image.shape[1]):\n h_pad, w_pad = 0, int(rot_img.shape[0] * image.shape[1] / image.shape[0] - rot_img.shape[1])\n # Pad height\n else:\n h_pad, w_pad = int(rot_img.shape[1] * image.shape[0] / image.shape[1] - rot_img.shape[0]), 0\n rot_img = np.pad(rot_img, ((h_pad // 2, h_pad - h_pad // 2), (w_pad // 2, w_pad - w_pad // 2), (0, 0)))\n if preserve_origin_shape:\n # rescale\n rot_img = cv2.resize(rot_img, image.shape[:-1][::-1], interpolation=cv2.INTER_LINEAR)\n\n return rot_img\n\n\ndef remove_image_padding(image: np.ndarray) -> np.ndarray:\n \"\"\"Remove black border padding from an image\n\n Args:\n image: numpy tensor to remove padding from\n\n Returns:\n Image with padding removed\n \"\"\"\n # Find the bounding box of the non-black region\n rows = np.any(image, axis=1)\n cols = np.any(image, axis=0)\n rmin, rmax = np.where(rows)[0][[0, -1]]\n cmin, cmax = np.where(cols)[0][[0, -1]]\n\n return image[rmin : rmax + 1, cmin : cmax + 1]\n\n\ndef estimate_page_angle(polys: np.ndarray) -> float:\n \"\"\"Takes a batch of rotated previously ORIENTED polys (N, 4, 2) (rectified by the classifier) and return the\n estimated angle ccw in degrees\n \"\"\"\n # Compute mean left points and mean right point with respect to the reading direction (oriented polygon)\n xleft = polys[:, 0, 0] + polys[:, 3, 0]\n yleft = polys[:, 0, 1] + polys[:, 3, 1]\n xright = polys[:, 1, 0] + polys[:, 2, 0]\n yright = polys[:, 1, 1] + polys[:, 2, 1]\n with np.errstate(divide=\"raise\", invalid=\"raise\"):\n try:\n return float(\n np.median(np.arctan((yleft - yright) / (xright - xleft)) * 180 / np.pi) # Y axis from top to bottom!\n )\n except FloatingPointError:\n return 0.0\n\n\ndef convert_to_relative_coords(geoms: np.ndarray, img_shape: tuple[int, int]) -> np.ndarray:\n \"\"\"Convert a geometry to relative coordinates\n\n Args:\n geoms: a set of polygons of shape (N, 4, 2) or of straight boxes of shape (N, 4)\n img_shape: the height and width of the image\n\n Returns:\n the updated geometry\n \"\"\"\n # Polygon\n if geoms.ndim == 3 and geoms.shape[1:] == (4, 2):\n polygons: np.ndarray = np.empty(geoms.shape, dtype=np.float32)\n polygons[..., 0] = geoms[..., 0] / img_shape[1]\n polygons[..., 1] = geoms[..., 1] / img_shape[0]\n return polygons.clip(0, 1)\n if geoms.ndim == 2 and geoms.shape[1] == 4:\n boxes: np.ndarray = np.empty(geoms.shape, dtype=np.float32)\n boxes[:, ::2] = geoms[:, ::2] / img_shape[1]\n boxes[:, 1::2] = geoms[:, 1::2] / img_shape[0]\n return boxes.clip(0, 1)\n\n raise ValueError(f\"invalid format for arg `geoms`: {geoms.shape}\")\n\n\ndef extract_crops(img: np.ndarray, boxes: np.ndarray, channels_last: bool = True) -> list[np.ndarray]:\n \"\"\"Created cropped images from list of bounding boxes\n\n Args:\n img: input image\n boxes: bounding boxes of shape (N, 4) where N is the number of boxes, and the relative\n coordinates (xmin, ymin, xmax, ymax)\n channels_last: whether the channel dimensions is the last one instead of the last one\n\n Returns:\n list of cropped images\n \"\"\"\n if boxes.shape[0] == 0:\n return []\n if boxes.shape[1] != 4:\n raise AssertionError(\"boxes are expected to be relative and in order (xmin, ymin, xmax, ymax)\")\n\n # Project relative coordinates\n _boxes = boxes.copy()\n h, w = img.shape[:2] if channels_last else img.shape[-2:]\n if not np.issubdtype(_boxes.dtype, np.integer):\n _boxes[:, [0, 2]] *= w\n _boxes[:, [1, 3]] *= h\n _boxes = _boxes.round().astype(int)\n # Add last index\n _boxes[2:] += 1\n if channels_last:\n return deepcopy([img[box[1] : box[3], box[0] : box[2]] for box in _boxes])\n\n return deepcopy([img[:, box[1] : box[3], box[0] : box[2]] for box in _boxes])\n\n\ndef extract_rcrops(\n img: np.ndarray, polys: np.ndarray, dtype=np.float32, channels_last: bool = True, assume_horizontal: bool = False\n) -> list[np.ndarray]:\n \"\"\"Created cropped images from list of rotated bounding boxes\n\n Args:\n img: input image\n polys: bounding boxes of shape (N, 4, 2)\n dtype: target data type of bounding boxes\n channels_last: whether the channel dimensions is the last one instead of the last one\n assume_horizontal: whether the boxes are assumed to be only horizontally oriented\n\n Returns:\n list of cropped images\n \"\"\"\n if polys.shape[0] == 0:\n return []\n if polys.shape[1:] != (4, 2):\n raise AssertionError(\"polys are expected to be quadrilateral, of shape (N, 4, 2)\")\n\n # Project relative coordinates\n _boxes = polys.copy()\n height, width = img.shape[:2] if channels_last else img.shape[-2:]\n if not np.issubdtype(_boxes.dtype, np.integer):\n _boxes[:, :, 0] *= width\n _boxes[:, :, 1] *= height\n\n src_img = img if channels_last else img.transpose(1, 2, 0)\n\n # Handle only horizontal oriented boxes\n if assume_horizontal:\n crops = []\n\n for box in _boxes:\n # Calculate the centroid of the quadrilateral\n centroid = np.mean(box, axis=0)\n\n # Divide the points into left and right\n left_points = box[box[:, 0] < centroid[0]]\n right_points = box[box[:, 0] >= centroid[0]]\n\n # Sort the left points according to the y-axis\n left_points = left_points[np.argsort(left_points[:, 1])]\n top_left_pt = left_points[0]\n bottom_left_pt = left_points[-1]\n # Sort the right points according to the y-axis\n right_points = right_points[np.argsort(right_points[:, 1])]\n top_right_pt = right_points[0]\n bottom_right_pt = right_points[-1]\n box_points = np.array(\n [top_left_pt, bottom_left_pt, top_right_pt, bottom_right_pt],\n dtype=dtype,\n )\n\n # Get the width and height of the rectangle that will contain the warped quadrilateral\n width_upper = np.linalg.norm(top_right_pt - top_left_pt)\n width_lower = np.linalg.norm(bottom_right_pt - bottom_left_pt)\n height_left = np.linalg.norm(bottom_left_pt - top_left_pt)\n height_right = np.linalg.norm(bottom_right_pt - top_right_pt)\n\n # Get the maximum width and height\n rect_width = max(int(width_upper), int(width_lower))\n rect_height = max(int(height_left), int(height_right))\n\n dst_pts = np.array(\n [\n [0, 0], # top-left\n # bottom-left\n [0, rect_height - 1],\n # top-right\n [rect_width - 1, 0],\n # bottom-right\n [rect_width - 1, rect_height - 1],\n ],\n dtype=dtype,\n )\n\n # Get the perspective transform matrix using the box points\n affine_mat = cv2.getPerspectiveTransform(box_points, dst_pts)\n\n # Perform the perspective warp to get the rectified crop\n crop = cv2.warpPerspective(\n src_img,\n affine_mat,\n (rect_width, rect_height),\n )\n\n # Add the crop to the list of crops\n crops.append(crop)\n\n # Handle any oriented boxes\n else:\n src_pts = _boxes[:, :3].astype(np.float32)\n # Preserve size\n d1 = np.linalg.norm(src_pts[:, 0] - src_pts[:, 1], axis=-1)\n d2 = np.linalg.norm(src_pts[:, 1] - src_pts[:, 2], axis=-1)\n # (N, 3, 2)\n dst_pts = np.zeros((_boxes.shape[0], 3, 2), dtype=dtype)\n dst_pts[:, 1, 0] = dst_pts[:, 2, 0] = d1 - 1\n dst_pts[:, 2, 1] = d2 - 1\n # Use a warp transformation to extract the crop\n crops = [\n cv2.warpAffine(\n src_img,\n # Transformation matrix\n cv2.getAffineTransform(src_pts[idx], dst_pts[idx]),\n (int(d1[idx]), int(d2[idx])),\n )\n for idx in range(_boxes.shape[0])\n ]\n\n return crops\n"
},
{
"path": "onnxtr/utils/multithreading.py",
"content": "# Copyright (C) 2021-2026, Mindee | Felix Dittrich.\n\n# This program is licensed under the Apache License 2.0.\n# See LICENSE or go to for full license details.\n\n\nimport multiprocessing as mp\nimport os\nfrom collections.abc import Callable, Iterable, Iterator\nfrom multiprocessing.pool import ThreadPool\nfrom typing import Any\n\nfrom onnxtr.file_utils import ENV_VARS_TRUE_VALUES\n\n__all__ = [\"multithread_exec\"]\n\n\ndef multithread_exec(func: Callable[[Any], Any], seq: Iterable[Any], threads: int | None = None) -> Iterator[Any]:\n \"\"\"Execute a given function in parallel for each element of a given sequence\n\n >>> from onnxtr.utils.multithreading import multithread_exec\n >>> entries = [1, 4, 8]\n >>> results = multithread_exec(lambda x: x ** 2, entries)\n\n Args:\n func: function to be executed on each element of the iterable\n seq: iterable\n threads: number of workers to be used for multiprocessing\n\n Returns:\n iterator of the function's results using the iterable as inputs\n\n Notes:\n This function uses ThreadPool from multiprocessing package, which uses `/dev/shm` directory for shared memory.\n If you do not have write permissions for this directory (if you run `onnxtr` on AWS Lambda for instance),\n you might want to disable multiprocessing. To achieve that, set 'ONNXTR_MULTIPROCESSING_DISABLE' to 'TRUE'.\n \"\"\"\n threads = threads if isinstance(threads, int) else min(16, mp.cpu_count())\n # Single-thread\n if threads < 2 or os.environ.get(\"ONNXTR_MULTIPROCESSING_DISABLE\", \"\").upper() in ENV_VARS_TRUE_VALUES:\n results = map(func, seq)\n # Multi-threading\n else:\n with ThreadPool(threads) as tp:\n # ThreadPool's map function returns a list, but seq could be of a different type\n # That's why wrapping result in map to return iterator\n results = map(lambda x: x, tp.map(func, seq)) # noqa: C417\n return results\n"
},
{
"path": "onnxtr/utils/reconstitution.py",
"content": "# Copyright (C) 2021-2026, Mindee | Felix Dittrich.\n\n# This program is licensed under the Apache License 2.0.\n# See LICENSE or go to for full license details.\nimport logging\nfrom typing import Any\n\nimport numpy as np\nfrom anyascii import anyascii\nfrom PIL import Image, ImageDraw\n\nfrom .fonts import get_font\n\n__all__ = [\"synthesize_page\"]\n\n\n# Global variable to avoid multiple warnings\nROTATION_WARNING = False\n\n\ndef _warn_rotation(entry: dict[str, Any]) -> None: # pragma: no cover\n global ROTATION_WARNING\n if not ROTATION_WARNING and len(entry[\"geometry\"]) == 4:\n logging.warning(\"Polygons with larger rotations will lead to inaccurate rendering\")\n ROTATION_WARNING = True\n\n\ndef _synthesize(\n response: Image.Image,\n entry: dict[str, Any],\n w: int,\n h: int,\n draw_proba: bool = False,\n font_family: str | None = None,\n smoothing_factor: float = 0.75,\n min_font_size: int = 6,\n max_font_size: int = 50,\n) -> Image.Image:\n if len(entry[\"geometry\"]) == 2:\n (xmin, ymin), (xmax, ymax) = entry[\"geometry\"]\n polygon = [(xmin, ymin), (xmax, ymin), (xmax, ymax), (xmin, ymax)]\n else:\n polygon = entry[\"geometry\"]\n\n # Calculate the bounding box of the word\n x_coords, y_coords = zip(*polygon)\n xmin, ymin, xmax, ymax = (\n int(round(w * min(x_coords))),\n int(round(h * min(y_coords))),\n int(round(w * max(x_coords))),\n int(round(h * max(y_coords))),\n )\n word_width = xmax - xmin\n word_height = ymax - ymin\n\n # If lines are provided instead of words, concatenate the word entries\n if \"words\" in entry:\n word_text = \" \".join(word[\"value\"] for word in entry[\"words\"])\n else:\n word_text = entry[\"value\"]\n # Find the optimal font size\n try:\n font_size = min(word_height, max_font_size)\n font = get_font(font_family, font_size)\n text_width, text_height = font.getbbox(word_text)[2:4]\n\n while (text_width > word_width or text_height > word_height) and font_size > min_font_size:\n font_size = max(int(font_size * smoothing_factor), min_font_size)\n font = get_font(font_family, font_size)\n text_width, text_height = font.getbbox(word_text)[2:4]\n except ValueError: # pragma: no cover\n font = get_font(font_family, min_font_size)\n\n # Create a mask for the word\n mask = Image.new(\"L\", (w, h), 0)\n ImageDraw.Draw(mask).polygon([(int(round(w * x)), int(round(h * y))) for x, y in polygon], fill=255)\n\n # Draw the word text\n d = ImageDraw.Draw(response)\n try:\n try:\n d.text((xmin, ymin), word_text, font=font, fill=(0, 0, 0), anchor=\"lt\")\n except UnicodeEncodeError: # pragma: no cover\n d.text((xmin, ymin), anyascii(word_text), font=font, fill=(0, 0, 0), anchor=\"lt\")\n # Catch generic exceptions to avoid crashing the whole rendering\n except Exception: # pragma: no cover\n logging.warning(f\"Could not render word: {word_text}\")\n\n if draw_proba:\n confidence = (\n entry[\"confidence\"]\n if \"confidence\" in entry\n else sum(w[\"confidence\"] for w in entry[\"words\"]) / len(entry[\"words\"])\n )\n p = int(255 * confidence)\n color = (255 - p, 0, p) # Red to blue gradient based on probability\n d.rectangle([(xmin, ymin), (xmax, ymax)], outline=color, width=2)\n\n prob_font = get_font(font_family, 20)\n prob_text = f\"{confidence:.2f}\"\n prob_text_width, prob_text_height = prob_font.getbbox(prob_text)[2:4]\n\n # Position the probability slightly above the bounding box\n prob_x_offset = (word_width - prob_text_width) // 2\n prob_y_offset = ymin - prob_text_height - 2\n prob_y_offset = max(0, prob_y_offset)\n\n d.text((xmin + prob_x_offset, prob_y_offset), prob_text, font=prob_font, fill=color, anchor=\"lt\")\n\n return response\n\n\ndef synthesize_page(\n page: dict[str, Any],\n draw_proba: bool = False,\n font_family: str | None = None,\n smoothing_factor: float = 0.95,\n min_font_size: int = 8,\n max_font_size: int = 50,\n) -> np.ndarray:\n \"\"\"Draw a the content of the element page (OCR response) on a blank page.\n\n Args:\n page: exported Page object to represent\n draw_proba: if True, draw words in colors to represent confidence. Blue: p=1, red: p=0\n font_family: family of the font\n smoothing_factor: factor to smooth the font size\n min_font_size: minimum font size\n max_font_size: maximum font size\n\n Returns:\n the synthesized page\n \"\"\"\n # Draw template\n h, w = page[\"dimensions\"]\n response = Image.new(\"RGB\", (w, h), color=(255, 255, 255))\n\n for block in page[\"blocks\"]:\n # If lines are provided use these to get better rendering results\n if len(block[\"lines\"]) > 1:\n for line in block[\"lines\"]:\n _warn_rotation(block) # pragma: no cover\n response = _synthesize(\n response=response,\n entry=line,\n w=w,\n h=h,\n draw_proba=draw_proba,\n font_family=font_family,\n smoothing_factor=smoothing_factor,\n min_font_size=min_font_size,\n max_font_size=max_font_size,\n )\n # Otherwise, draw each word\n else:\n for line in block[\"lines\"]:\n _warn_rotation(block) # pragma: no cover\n for word in line[\"words\"]:\n response = _synthesize(\n response=response,\n entry=word,\n w=w,\n h=h,\n draw_proba=draw_proba,\n font_family=font_family,\n smoothing_factor=smoothing_factor,\n min_font_size=min_font_size,\n max_font_size=max_font_size,\n )\n\n return np.array(response, dtype=np.uint8)\n"
},
{
"path": "onnxtr/utils/repr.py",
"content": "# Copyright (C) 2021-2026, Mindee | Felix Dittrich.\n\n# This program is licensed under the Apache License 2.0.\n# See LICENSE or go to for full license details.\n\n# Adapted from https://github.com/pytorch/torch/blob/master/torch/nn/modules/module.py\n\n\n__all__ = [\"NestedObject\"]\n\n\ndef _addindent(s_, num_spaces):\n s = s_.split(\"\\n\")\n # don't do anything for single-line stuff\n if len(s) == 1:\n return s_\n first = s.pop(0)\n s = [(num_spaces * \" \") + line for line in s]\n s = \"\\n\".join(s)\n s = first + \"\\n\" + s\n return s\n\n\nclass NestedObject:\n \"\"\"Base class for all nested objects in onnxtr\"\"\"\n\n _children_names: list[str]\n\n def extra_repr(self) -> str:\n return \"\"\n\n def __repr__(self):\n # We treat the extra repr like the sub-object, one item per line\n extra_lines = []\n extra_repr = self.extra_repr()\n # empty string will be split into list ['']\n if extra_repr:\n extra_lines = extra_repr.split(\"\\n\")\n child_lines = []\n if hasattr(self, \"_children_names\"):\n for key in self._children_names:\n child = getattr(self, key)\n if isinstance(child, list) and len(child) > 0:\n child_str = \",\\n\".join([repr(subchild) for subchild in child])\n if len(child) > 1:\n child_str = _addindent(f\"\\n{child_str},\", 2) + \"\\n\"\n child_str = f\"[{child_str}]\"\n else:\n child_str = repr(child)\n child_str = _addindent(child_str, 2)\n child_lines.append(\"(\" + key + \"): \" + child_str)\n lines = extra_lines + child_lines\n\n main_str = self.__class__.__name__ + \"(\"\n if lines:\n # simple one-liner info, which most builtin Modules will use\n if len(extra_lines) == 1 and not child_lines:\n main_str += extra_lines[0]\n else:\n main_str += \"\\n \" + \"\\n \".join(lines) + \"\\n\"\n\n main_str += \")\"\n return main_str\n"
},
{
"path": "onnxtr/utils/visualization.py",
"content": "# Copyright (C) 2021-2026, Mindee | Felix Dittrich.\n\n# This program is licensed under the Apache License 2.0.\n# See LICENSE or go to for full license details.\n\nfrom copy import deepcopy\nfrom typing import Any\n\nimport cv2\nimport matplotlib.patches as patches\nimport matplotlib.pyplot as plt\nimport numpy as np\nfrom matplotlib.figure import Figure\n\nfrom .common_types import BoundingBox, Polygon4P\n\n__all__ = [\"visualize_page\", \"draw_boxes\"]\n\n\ndef rect_patch(\n geometry: BoundingBox,\n page_dimensions: tuple[int, int],\n label: str | None = None,\n color: tuple[float, float, float] = (0, 0, 0),\n alpha: float = 0.3,\n linewidth: int = 2,\n fill: bool = True,\n preserve_aspect_ratio: bool = False,\n) -> patches.Rectangle:\n \"\"\"Create a matplotlib rectangular patch for the element\n\n Args:\n geometry: bounding box of the element\n page_dimensions: dimensions of the Page in format (height, width)\n label: label to display when hovered\n color: color to draw box\n alpha: opacity parameter to fill the boxes, 0 = transparent\n linewidth: line width\n fill: whether the patch should be filled\n preserve_aspect_ratio: pass True if you passed True to the predictor\n\n Returns:\n a rectangular Patch\n \"\"\"\n if len(geometry) != 2 or any(not isinstance(elt, tuple) or len(elt) != 2 for elt in geometry):\n raise ValueError(\"invalid geometry format\")\n\n # Unpack\n height, width = page_dimensions\n (xmin, ymin), (xmax, ymax) = geometry\n # Switch to absolute coords\n if preserve_aspect_ratio:\n width = height = max(height, width)\n xmin, w = xmin * width, (xmax - xmin) * width\n ymin, h = ymin * height, (ymax - ymin) * height\n\n return patches.Rectangle(\n (xmin, ymin),\n w,\n h,\n fill=fill,\n linewidth=linewidth,\n edgecolor=(*color, alpha),\n facecolor=(*color, alpha),\n label=label,\n )\n\n\ndef polygon_patch(\n geometry: np.ndarray,\n page_dimensions: tuple[int, int],\n label: str | None = None,\n color: tuple[float, float, float] = (0, 0, 0),\n alpha: float = 0.3,\n linewidth: int = 2,\n fill: bool = True,\n preserve_aspect_ratio: bool = False,\n) -> patches.Polygon:\n \"\"\"Create a matplotlib polygon patch for the element\n\n Args:\n geometry: bounding box of the element\n page_dimensions: dimensions of the Page in format (height, width)\n label: label to display when hovered\n color: color to draw box\n alpha: opacity parameter to fill the boxes, 0 = transparent\n linewidth: line width\n fill: whether the patch should be filled\n preserve_aspect_ratio: pass True if you passed True to the predictor\n\n Returns:\n a polygon Patch\n \"\"\"\n if not geometry.shape == (4, 2):\n raise ValueError(\"invalid geometry format\")\n\n # Unpack\n height, width = page_dimensions\n geometry[:, 0] = geometry[:, 0] * (max(width, height) if preserve_aspect_ratio else width)\n geometry[:, 1] = geometry[:, 1] * (max(width, height) if preserve_aspect_ratio else height)\n\n return patches.Polygon(\n geometry,\n fill=fill,\n linewidth=linewidth,\n edgecolor=(*color, alpha),\n facecolor=(*color, alpha),\n label=label,\n )\n\n\ndef create_obj_patch(\n geometry: BoundingBox | Polygon4P | np.ndarray,\n page_dimensions: tuple[int, int],\n **kwargs: Any,\n) -> patches.Patch:\n \"\"\"Create a matplotlib patch for the element\n\n Args:\n geometry: bounding box (straight or rotated) of the element\n page_dimensions: dimensions of the page in format (height, width)\n **kwargs: keyword arguments for the patch\n\n Returns:\n a matplotlib Patch\n \"\"\"\n if isinstance(geometry, tuple):\n if len(geometry) == 2: # straight word BB (2 pts)\n return rect_patch(geometry, page_dimensions, **kwargs)\n elif len(geometry) == 4: # rotated word BB (4 pts)\n return polygon_patch(np.asarray(geometry), page_dimensions, **kwargs)\n elif isinstance(geometry, np.ndarray) and geometry.shape == (4, 2): # rotated line\n return polygon_patch(geometry, page_dimensions, **kwargs)\n raise ValueError(\"invalid geometry format\")\n\n\ndef visualize_page(\n page: dict[str, Any],\n image: np.ndarray,\n words_only: bool = True,\n display_artefacts: bool = True,\n scale: float = 10,\n interactive: bool = True,\n add_labels: bool = True,\n **kwargs: Any,\n) -> Figure:\n \"\"\"Visualize a full page with predicted blocks, lines and words\n\n >>> import numpy as np\n >>> import matplotlib.pyplot as plt\n >>> from onnxtr.utils.visualization import visualize_page\n >>> from onnxtr.models import ocr_db_crnn\n >>> model = ocr_db_crnn()\n >>> input_page = (255 * np.random.rand(600, 800, 3)).astype(np.uint8)\n >>> out = model([[input_page]])\n >>> visualize_page(out[0].pages[0].export(), input_page)\n >>> plt.show()\n\n Args:\n page: the exported Page of a Document\n image: np array of the page, needs to have the same shape than page['dimensions']\n words_only: whether only words should be displayed\n display_artefacts: whether artefacts should be displayed\n scale: figsize of the largest windows side\n interactive: whether the plot should be interactive\n add_labels: for static plot, adds text labels on top of bounding box\n **kwargs: keyword arguments for the polygon patch\n\n Returns:\n the matplotlib figure\n \"\"\"\n # Get proper scale and aspect ratio\n h, w = image.shape[:2]\n size = (scale * w / h, scale) if h > w else (scale, h / w * scale)\n fig, ax = plt.subplots(figsize=size)\n # Display the image\n ax.imshow(image)\n # hide both axis\n ax.axis(\"off\")\n\n if interactive:\n artists: list[patches.Patch] = [] # instantiate an empty list of patches (to be drawn on the page)\n\n for block in page[\"blocks\"]:\n if not words_only:\n rect = create_obj_patch(\n block[\"geometry\"], page[\"dimensions\"], label=\"block\", color=(0, 1, 0), linewidth=1, **kwargs\n )\n # add patch on figure\n ax.add_patch(rect)\n if interactive:\n # add patch to cursor's artists\n artists.append(rect)\n\n for line in block[\"lines\"]:\n if not words_only:\n rect = create_obj_patch(\n line[\"geometry\"], page[\"dimensions\"], label=\"line\", color=(1, 0, 0), linewidth=1, **kwargs\n )\n ax.add_patch(rect)\n if interactive:\n artists.append(rect)\n\n for word in line[\"words\"]:\n rect = create_obj_patch(\n word[\"geometry\"],\n page[\"dimensions\"],\n label=f\"{word['value']} (confidence: {word['confidence']:.2%})\",\n color=(0, 0, 1),\n **kwargs,\n )\n ax.add_patch(rect)\n if interactive:\n artists.append(rect)\n elif add_labels:\n if len(word[\"geometry\"]) == 5:\n text_loc = (\n int(page[\"dimensions\"][1] * (word[\"geometry\"][0] - word[\"geometry\"][2] / 2)),\n int(page[\"dimensions\"][0] * (word[\"geometry\"][1] - word[\"geometry\"][3] / 2)),\n )\n else:\n text_loc = (\n int(page[\"dimensions\"][1] * word[\"geometry\"][0][0]),\n int(page[\"dimensions\"][0] * word[\"geometry\"][0][1]),\n )\n\n if len(word[\"geometry\"]) == 2:\n # We draw only if boxes are in straight format\n ax.text(\n *text_loc,\n word[\"value\"],\n size=10,\n alpha=0.5,\n color=(0, 0, 1),\n )\n\n if display_artefacts:\n for artefact in block[\"artefacts\"]:\n rect = create_obj_patch(\n artefact[\"geometry\"],\n page[\"dimensions\"],\n label=\"artefact\",\n color=(0.5, 0.5, 0.5),\n linewidth=1,\n **kwargs,\n )\n ax.add_patch(rect)\n if interactive:\n artists.append(rect)\n\n if interactive:\n import mplcursors\n\n # Create mlp Cursor to hover patches in artists\n mplcursors.Cursor(artists, hover=2).connect(\"add\", lambda sel: sel.annotation.set_text(sel.artist.get_label()))\n fig.tight_layout(pad=0.0)\n\n return fig\n\n\ndef draw_boxes(boxes: np.ndarray, image: np.ndarray, color: tuple[int, int, int] | None = None, **kwargs) -> None:\n \"\"\"Draw an array of relative straight boxes on an image\n\n Args:\n boxes: array of relative boxes, of shape (*, 4)\n image: np array, float32 or uint8\n color: color to use for bounding box edges\n **kwargs: keyword arguments from `matplotlib.pyplot.plot`\n \"\"\"\n h, w = image.shape[:2]\n # Convert boxes to absolute coords\n _boxes = deepcopy(boxes)\n _boxes[:, [0, 2]] *= w\n _boxes[:, [1, 3]] *= h\n _boxes = _boxes.astype(np.int32)\n for box in _boxes.tolist():\n xmin, ymin, xmax, ymax = box\n image = cv2.rectangle(\n image,\n (xmin, ymin),\n (xmax, ymax),\n color=color if isinstance(color, tuple) else (0, 0, 255),\n thickness=2,\n )\n plt.imshow(image)\n plt.plot(**kwargs)\n"
},
{
"path": "onnxtr/utils/vocabs.py",
"content": "# Copyright (C) 2021-2026, Mindee | Felix Dittrich.\n\n# This program is licensed under the Apache License 2.0.\n# See LICENSE or go to for full license details.\n\nimport re\nimport string\n\n__all__ = [\"VOCABS\"]\n\n_BASE_VOCABS = {\n # Latin\n \"digits\": string.digits,\n \"ascii_letters\": string.ascii_letters,\n \"punctuation\": string.punctuation,\n \"currency\": \"£€¥¢฿\",\n # Cyrillic\n \"generic_cyrillic_letters\": \"абвгдежзийклмнопрстуфхцчшщьюяАБВГДЕЖЗИЙКЛМНОПРСТУФХЦЧШЩЬЮЯ\",\n \"russian_cyrillic_letters\": \"ёыэЁЫЭ\",\n \"russian_signs\": \"ъЪ\",\n # Greek\n \"ancient_greek\": \"αβγδεζηθικλμνξοπρστςυφχψωΑΒΓΔΕΖΗΘΙΚΛΜΝΞΟΠΡΣΤΥΦΧΨΩ\",\n # Arabic & Persian\n \"arabic_diacritics\": \"\".join([\"ً\", \"ٌ\", \"ٍ\", \"َ\", \"ُ\", \"ِ\", \"ّ\", \"ْ\", \"ٕ\", \"ٓ\", \"ٔ\", \"ٚ\"]),\n \"arabic_digits\": \"٠١٢٣٤٥٦٧٨٩\",\n \"arabic_letters\": \"ءآأؤإئابةتثجحخدذرزسشصضطظعغـفقكلمنهوىيٱ\",\n \"arabic_punctuation\": \"؟؛«»—،\",\n \"persian_letters\": \"پچژڢڤگکی\",\n # Bengali\n \"bengali_consonants\": \"কখগঘঙচছজঝঞটঠডঢণতথদধনপফবভমযরলশষসহড়ঢ়য়ৰৱৼ\",\n \"bengali_vowels\": \"অআইঈউঊঋঌএঐওঔৠৡ\",\n \"bengali_digits\": \"০১২৩৪৫৬৭৮৯\",\n \"bengali_matras\": \"\".join([\"া\", \"ি\", \"ী\", \"ু\", \"ূ\", \"ৃ\", \"ে\", \"ৈ\", \"ো\", \"ৌ\", \"ৗ\"]),\n \"bengali_virama\": \"্\",\n \"bengali_punctuation\": \"ঽৎ৽৺৻\",\n \"bengali_signs\": \"\".join([\"ঁ\", \"ং\", \"ঃ\", \"়\"]),\n # Gujarati\n \"gujarati_consonants\": \"કખગઘઙચછજઝઞટઠડઢણતથદધનપફબભમયરલળવશષસહ\",\n \"gujarati_vowels\": \"અઆઇઈઉઊઋઌઍએઐઑઓઔ\",\n \"gujarati_digits\": \"૦૧૨૩૪૫૬૭૮૯\",\n \"gujarati_matras\": \"\".join([\n \"ઁ\",\n \"ં\",\n \"ઃ\",\n \"઼\",\n \"ા\",\n \"િ\",\n \"ી\",\n \"ુ\",\n \"ૂ\",\n \"ૃ\",\n \"ૄ\",\n \"ૅ\",\n \"ે\",\n \"ૈ\",\n \"ૉ\",\n \"ો\",\n \"ૌ\",\n \"ૢ\",\n \"ૣ\",\n \"ૺ\",\n \"ૻ\",\n \"ૼ\",\n \"૽\",\n \"૾\",\n \"૿\",\n ]),\n \"gujarati_virama\": \"્\",\n \"gujarati_punctuation\": \"ઽ॥\",\n \"gujarati_signs\": \"ૐ૰\",\n # Devanagari\n \"devanagari_consonants\": \"कखगघङचछजझञटठडढणतथदधनपफबभमयरलवशषसहऴऩळक़ख़ग़ज़ड़ढ़फ़य़ऱॺॻॼॽॾ\",\n \"devanagari_vowels\": \"अआइईउऊऋऌऍऎएऐऑऒओऔॠॡॲऄॵॶॳॴॷॸॹ\",\n \"devanagari_digits\": \"०१२३४५६७८९\",\n \"devanagari_matras\": \"\".join([\n \"़\",\n \"ं\",\n \"ँ\",\n \"ः\",\n \"॑\",\n \"॒\",\n \"ा\",\n \"ि\",\n \"ी\",\n \"ु\",\n \"ू\",\n \"ृ\",\n \"ॄ\",\n \"ॅ\",\n \"ॆ\",\n \"े\",\n \"ै\",\n \"ॉ\",\n \"ॊ\",\n \"ो\",\n \"ौ\",\n \"ॢ\",\n \"ॣ\",\n \"ॏ\",\n \"ॎ\",\n ]),\n \"devanagari_virama\": \"्\",\n \"devanagari_punctuation\": \"।॥॰ऽꣲ\",\n \"devanagari_signs\": \"ॐ\",\n # Punjabi (Gurmukhi script)\n \"punjabi_consonants\": \"ਕਖਗਘਙਚਛਜਝਞਟਠਡਢਣਤਥਦਧਨਪਫਬਭਮਯਰਲਵਸ਼ਸਹਖ਼ਗ਼ਜ਼ਫ਼ੜਲ਼\",\n \"punjabi_vowels\": \"ਅਆਇਈਉਊਏਐਓਔੲੳ\",\n \"punjabi_digits\": \"੦੧੨੩੪੫੬੭੮੯\",\n \"punjabi_matras\": \"\".join([\"ਂ\", \"਼\", \"ਾ\", \"ਿ\", \"ੀ\", \"ੁ\", \"ੂ\", \"ੇ\", \"ੈ\", \"ੋ\", \"ੌ\", \"ੑ\", \"ੰ\", \"ੱ\", \"ੵ\"]),\n \"punjabi_virama\": \"੍\",\n \"punjabi_punctuation\": \"।॥\",\n \"punjabi_signs\": \"ੴ\",\n # Tamil\n \"tamil_consonants\": \"கஙசஞடணதநபமயரலவழளறன\",\n \"tamil_vowels\": \"அஆஇஈஉஊஎஏஐஒஓஔ\",\n \"tamil_digits\": \"௦௧௨௩௪௫௬௭௮௯\",\n \"tamil_matras\": \"\".join([\"ா\", \"ி\", \"ீ\", \"ு\", \"ூ\", \"ெ\", \"ே\", \"ை\", \"ொ\", \"ோ\", \"ௌ\"]),\n \"tamil_virama\": \"்\",\n \"tamil_punctuation\": \"௰௱௲\",\n \"tamil_signs\": \"ஃௐ\",\n \"tamil_fractions\": \"௳௴௵௶௷௸௹௺\",\n # Telugu\n \"telugu_consonants\": \"కఖగఘఙచఛజఝఞటఠడఢణతథదధనపఫబభమయరఱలళవశషసహఴ\",\n \"telugu_digits\": \"౦౧౨౩౪౫౬౭౮౯\" + \"౸౹౺౻\", # Telugu digits and fractional digits\n \"telugu_vowels\": \"అఆఇఈఉఊఋఌఎఏఐఒఓఔౠౡ\",\n \"telugu_matras\": \"\".join([\"ా\", \"ి\", \"ీ\", \"ు\", \"ూ\", \"ృ\", \"ౄ\", \"ె\", \"ే\", \"ై\", \"ొ\", \"ో\", \"ౌ\", \"ౢ\", \"ౣ\"]),\n \"telugu_virama\": \"్\",\n \"telugu_punctuation\": \"ఽ\",\n \"telugu_signs\": \"\".join([\"ఁ\", \"ం\", \"ః\"]),\n # Kannada\n \"kannada_consonants\": \"ಕಖಗಘಙಚಛಜಝಞಟಠಡಢಣತಥದಧನಪಫಬಭಮಯರಲವಶಷಸಹಳ\",\n \"kannada_vowels\": \"ಅಆಇಈಉಊಋॠಌೡಎಏಐಒಓಔ\",\n \"kannada_digits\": \"೦೧೨೩೪೫೬೭೮೯\",\n \"kannada_matras\": \"\".join([\"ಾ\", \"ಿ\", \"ೀ\", \"ು\", \"ೂ\", \"ೃ\", \"ೄ\", \"ೆ\", \"ೇ\", \"ೈ\", \"ೊ\", \"ೋ\", \"ೌ\"]),\n \"kannada_virama\": \"್\",\n \"kannada_punctuation\": \"।॥ೱೲ\",\n \"kannada_signs\": \"\".join([\"ಂ\", \"ಃ\", \"ಁ\"]),\n # Sinhala\n \"sinhala_consonants\": \"කඛගඝඞචඡජඣඤටඨඩඪණතථදධනපඵබභමයරලවශෂසහළෆ\",\n \"sinhala_vowels\": \"අආඇඈඉඊඋඌඍඎඏඐඑඒඓඔඕඖ\",\n \"sinhala_digits\": \"෦෧෨෩෪෫෬෭෮෯\",\n \"sinhala_matras\": \"\".join([\"ා\", \"ැ\", \"ෑ\", \"ි\", \"ී\", \"ු\", \"ූ\", \"ෙ\", \"ේ\", \"ෛ\", \"ො\", \"ෝ\", \"ෞ\"]),\n \"sinhala_virama\": \"්\",\n \"sinhala_punctuation\": \"෴\",\n \"sinhala_signs\": \"\".join([\"ං\", \"ඃ\"]),\n # Malayalam\n \"malayalam_consonants\": \"കഖഗഘങചഛജഝഞടഠഡഢണതഥദധനപഫബഭമയരറലളഴവശഷസഹ\",\n \"malayalam_vowels\": \"അആഇഈഉഊഋൠഌൡഎഏഐഒഓഔ\",\n \"malayalam_digits\": \"൦൧൨൩൪൫൬൭൮൯\",\n \"malayalam_matras\": \"\".join([\"ാ\", \"ി\", \"ീ\", \"ു\", \"ൂ\", \"ൃ\", \"ൄ\", \"ൢ\", \"ൣ\", \"െ\", \"േ\", \"ൈ\", \"ൊ\", \"ോ\", \"ൌ\"]),\n \"malayalam_virama\": \"്\",\n \"malayalam_signs\": \"\".join([\"ഃ\", \"൹\", \"ഽ\", \"൏\", \"ം\"]),\n # Odia (Oriya)\n \"odia_consonants\": \"କଖଗଘଙଚଛଜଝଞଟଠଡଢଣତଥଦଧନପଫବଭମଯରଲଳଵଶଷସହୟୱଡ଼ଢ଼\",\n \"odia_vowels\": \"ଅଆଇଈଉଊଋଌଏଐଓଔୡୠ\",\n \"odia_digits\": \"୦୧୨୩୪୫୬୭୮୯\" + \"୲୳୴୵୶୷\", # Odia digits and fractional digits\n \"odia_matras\": \"\".join([\"ା\", \"ି\", \"ୀ\", \"ୁ\", \"ୂ\", \"ୃ\", \"ୄ\", \"େ\", \"ୈ\", \"ୋ\", \"ୌ\", \"ୢ\", \"ୣ\"]),\n \"odia_virama\": \"୍\",\n \"odia_punctuation\": \"ଽ\",\n \"odia_signs\": \"\".join([\"ଂ\", \"ଃ\", \"ଁ\", \"଼\", \"୰\"]),\n # Khmer\n \"khmer_consonants\": \"កខគឃងចឆជឈញដឋឌឍណតថទធនបផពភមយរលវឝឞសហឡអ\",\n \"khmer_vowels\": \"ឣឤឥឦឧឨឩឪឫឬឭឮឯឰឱឲឳ\",\n \"khmer_digits\": \"០១២៣៤៥៦៧៨៩\",\n \"khmer_matras\": \"\".join([\"ា\", \"ិ\", \"ី\", \"ឹ\", \"ឺ\", \"ុ\", \"ូ\", \"ួ\", \"ើ\", \"ឿ\", \"ៀ\", \"េ\", \"ែ\", \"ៃ\", \"ោ\", \"ៅ\"]),\n \"khmer_diacritics\": \"\".join([\"ំ\", \"ះ\", \"ៈ\", \"៉\", \"៊\", \"់\", \"៌\", \"៍\", \"៎\", \"៏\", \"័\", \"៑\", \"៓\", \"៝\"]),\n \"khmer_virama\": \"្\",\n \"khmer_punctuation\": \"។៕៖៘៙៚ៗៜ\",\n # Burmese\n \"burmese_consonants\": \"ကခဂဃငစဆဇဈဉညဋဌဍဎဏတထဒဓနပဖဗဘမယရလဝသဟဠအၐၑၒၓၔၕၚၛၜၝၡၥၦၮၯၰၵၶၷၸၹၺၻၼၽၾၿႀႁႎ\",\n \"burmese_vowels\": \"ဣဤဥဦဧဩဪဿ\",\n \"burmese_digits\": \"၀၁၂၃၄၅၆၇၈၉\" + \"႐႑႒႓႔႕႖႗႘႙\", # Burmese digits and Shan digits\n \"burmese_diacritics\": \"\".join([\"့\", \"း\", \"ံ\", \"ါ\", \"ာ\", \"ိ\", \"ီ\", \"ု\", \"ူ\", \"ေ\", \"ဲ\", \"ဳ\", \"ဴ\", \"ဵ\", \"ျြွှ\"]), # းံါာိီုူေဲံ့းှျြွှ\n # ္ (virama) and ် (final consonant) - the first is used to stack consonants, the second is used for final consonants\n \"burmese_virama\": \"\".join([\n \"္\",\n \"်\",\n ]),\n \"burmese_punctuation\": \"၊။၌၍၎၏\" + \"ၤ\" + \"ၗ\", # Includes ၗ and ၤ\n # Javanese\n \"javanese_consonants\": \"ꦏꦐꦑꦒꦓꦔꦕꦖꦗꦘꦙꦚꦛꦜꦝꦞꦟꦠꦡꦢꦣꦤꦥꦦꦧꦨꦩꦪꦫꦬꦭꦮꦯꦰꦱꦲ\",\n \"javanese_vowels\": \"ꦄꦅꦆꦇꦈꦉꦊꦋꦌꦍꦎ\" + \"ꦴꦵꦶꦷꦸꦹꦺꦻꦼ\", # sec: Dependent vowels ꦴꦵꦶꦷꦸꦹꦺꦻꦼ\n \"javanese_digits\": \"꧐꧑꧒꧓꧔꧕꧖꧗꧘꧙\",\n \"javanese_diacritics\": \"\".join([\"ꦀ\", \"ꦁ\", \"ꦂ\", \"ꦃ\", \"꦳\", \"ꦽ\", \"ꦾ\", \"ꦿ\"]), # ꦀꦁꦂꦃ꦳ꦽꦾꦿ\n \"javanese_virama\": \"꧀\",\n \"javanese_punctuation\": \"\".join([\"꧈\", \"꧉\", \"꧊\", \"꧋\", \"꧌\", \"꧍\", \"ꧏ\"]),\n # Sudanese\n \"sudanese_consonants\": \"ᮊᮋᮌᮍᮎᮏᮐᮑᮒᮓᮔᮕᮖᮗᮘᮙᮚᮛᮜᮝᮞᮟᮠᮮᮯᮺᮻᮼᮽᮾᮿ\",\n \"sudanese_vowels\": \"ᮃᮄᮅᮆᮇᮈᮉ\",\n \"sudanese_digits\": \"᮰᮱᮲᮳᮴᮵᮶᮷᮸᮹\",\n \"sudanese_diacritics\": \"\".join([\"ᮀ\", \"ᮁ\", \"ᮂ\", \"ᮡ\", \"ᮢ\", \"ᮣ\", \"ᮤ\", \"ᮥ\", \"ᮦ\", \"ᮧ\", \"ᮨ\", \"ᮩ\", \"᮪\", \"᮫\", \"ᮬ\", \"ᮭ\"]), # \"ᮀᮁᮂᮡᮢᮣᮤᮥᮦᮧᮨᮩ᮪᮫ᮬᮭ\"\n # Hebrew\n \"hebrew_cantillations\": \"\".join([\n \"֑\",\n \"֒\",\n \"֓\",\n \"֔\",\n \"֕\",\n \"֖\",\n \"֗\",\n \"֘\",\n \"֙\",\n \"֚\",\n \"֛\",\n \"֜\",\n \"֝\",\n \"֞\",\n \"֟\",\n \"֠\",\n \"֡\",\n \"֢\",\n \"֣\",\n \"֤\",\n \"֥\",\n \"֦\",\n \"֧\",\n \"֨\",\n \"֩\",\n \"֪\",\n \"֫\",\n \"֬\",\n \"֭\",\n \"֮\",\n \"֯\",\n ]),\n \"hebrew_consonants\": \"אבגדהוזחטיךכלםמןנסעףפץצקרשת\",\n \"hebrew_specials\": \"ׯװױײיִﬞײַﬠﬡﬢﬣﬤﬥﬦﬧﬨ﬩שׁשׂשּׁשּׂאַאָאּבּגּדּהּוּזּטּיּךּכּלּמּנּסּףּפּצּקּרּשּתּוֹבֿכֿפֿﭏ\",\n \"hebrew_punctuation\": \"\".join([\"ֽ\", \"־\", \"ֿ\", \"׀\", \"ׁ\", \"ׂ\", \"׃\", \"ׄ\", \"ׅ\", \"׆\", \"׳\", \"״\"]),\n \"hebrew_vowels\": \"\".join([\"ְ\", \"ֱ\", \"ֲ\", \"ֳ\", \"ִ\", \"ֵ\", \"ֶ\", \"ַ\", \"ָ\", \"ֹ\", \"ֺ\", \"ֻ\", \"ׇ\"]),\n}\n\n\nVOCABS: dict[str, str] = {}\n\nfor key, value in _BASE_VOCABS.items():\n VOCABS[key] = value\n\n# Latin & latin-dependent alphabets\nVOCABS[\"latin\"] = _BASE_VOCABS[\"digits\"] + _BASE_VOCABS[\"ascii_letters\"] + _BASE_VOCABS[\"punctuation\"]\nVOCABS[\"english\"] = VOCABS[\"latin\"] + \"°\" + _BASE_VOCABS[\"currency\"]\n\nVOCABS[\"albanian\"] = VOCABS[\"english\"] + \"çëÇË\"\n\nVOCABS[\"afrikaans\"] = VOCABS[\"english\"] + \"èëïîôûêÈËÏÎÔÛÊ\"\n\nVOCABS[\"azerbaijani\"] = re.sub(r\"[Ww]\", \"\", VOCABS[\"english\"]) + \"çəğöşüÇƏĞÖŞÜ\" + \"₼\"\n\nVOCABS[\"basque\"] = VOCABS[\"english\"] + \"ñçÑÇ\"\n\nVOCABS[\"bosnian\"] = re.sub(r\"[QqWwXxYy]\", \"\", VOCABS[\"english\"]) + \"čćđšžČĆĐŠŽ\"\n\nVOCABS[\"catalan\"] = VOCABS[\"english\"] + \"àèéíïòóúüçÀÈÉÍÏÒÓÚÜÇ\"\n\nVOCABS[\"croatian\"] = VOCABS[\"english\"] + \"ČčĆćĐ𩹮ž\"\n\nVOCABS[\"czech\"] = VOCABS[\"english\"] + \"áčďéěíňóřšťúůýžÁČĎÉĚÍŇÓŘŠŤÚŮÝŽ\"\n\nVOCABS[\"danish\"] = VOCABS[\"english\"] + \"æøåÆØÅ\"\n\nVOCABS[\"dutch\"] = VOCABS[\"english\"] + \"áéíóúüñÁÉÍÓÚÜÑ\"\n\nVOCABS[\"estonian\"] = VOCABS[\"english\"] + \"šžõäöüŠŽÕÄÖÜ\"\n\nVOCABS[\"esperanto\"] = re.sub(r\"[QqWwXxYy]\", \"\", VOCABS[\"english\"]) + \"ĉĝĥĵŝŭĈĜĤĴŜŬ\" + \"₷\"\n\nVOCABS[\"french\"] = VOCABS[\"english\"] + \"àâéèêëîïôùûüçÀÂÉÈÊËÎÏÔÙÛÜÇ\"\n\nVOCABS[\"finnish\"] = VOCABS[\"english\"] + \"äöÄÖ\"\n\nVOCABS[\"frisian\"] = re.sub(r\"[QqXx]\", \"\", VOCABS[\"english\"]) + \"âêôûúÂÊÔÛÚ\" + \"ƒƑ\"\n\nVOCABS[\"galician\"] = re.sub(r\"[JjKkWw]\", \"\", VOCABS[\"english\"]) + \"ñÑçÇ\"\n\nVOCABS[\"german\"] = VOCABS[\"english\"] + \"äöüßÄÖÜẞ\"\n\nVOCABS[\"hausa\"] = re.sub(r\"[PpQqVvXx]\", \"\", VOCABS[\"english\"]) + \"ɓɗƙƴƁƊƘƳ\" + \"₦\"\n\nVOCABS[\"hungarian\"] = VOCABS[\"english\"] + \"áéíóöúüÁÉÍÓÖÚÜ\"\n\nVOCABS[\"icelandic\"] = re.sub(r\"[CcQqWw]\", \"\", VOCABS[\"english\"]) + \"ðáéíóúýþæöÐÁÉÍÓÚÝÞÆÖ\"\n\nVOCABS[\"indonesian\"] = VOCABS[\"english\"]\n\nVOCABS[\"irish\"] = VOCABS[\"english\"] + \"áéíóúÁÉÍÓÚ\"\n\nVOCABS[\"italian\"] = VOCABS[\"english\"] + \"àèéìíîòóùúÀÈÉÌÍÎÒÓÙÚ\"\n\nVOCABS[\"latvian\"] = re.sub(r\"[QqWwXx]\", \"\", VOCABS[\"english\"]) + \"āčēģīķļņšūžĀČĒĢĪĶĻŅŠŪŽ\"\n\nVOCABS[\"lithuanian\"] = re.sub(r\"[QqWwXx]\", \"\", VOCABS[\"english\"]) + \"ąčęėįšųūžĄČĘĖĮŠŲŪŽ\"\n\nVOCABS[\"luxembourgish\"] = VOCABS[\"english\"] + \"äöüéëÄÖÜÉË\"\n\nVOCABS[\"malagasy\"] = re.sub(r\"[CcQqUuWwXx]\", \"\", VOCABS[\"english\"]) + \"ôñÔÑ\"\n\nVOCABS[\"malay\"] = VOCABS[\"english\"]\n\nVOCABS[\"maltese\"] = re.sub(r\"[CcYy]\", \"\", VOCABS[\"english\"]) + \"ċġħżĊĠĦŻ\"\n\nVOCABS[\"maori\"] = re.sub(r\"[BbCcDdFfJjLlOoQqSsVvXxYyZz]\", \"\", VOCABS[\"english\"]) + \"āēīōūĀĒĪŌŪ\"\n\nVOCABS[\"montenegrin\"] = re.sub(r\"[QqWwXxYy]\", \"\", VOCABS[\"english\"]) + \"čćšžźČĆŠŚŽŹ\"\n\nVOCABS[\"norwegian\"] = VOCABS[\"english\"] + \"æøåÆØÅ\"\n\nVOCABS[\"polish\"] = VOCABS[\"english\"] + \"ąćęłńóśźżĄĆĘŁŃÓŚŹŻ\"\n\nVOCABS[\"portuguese\"] = VOCABS[\"english\"] + \"áàâãéêíïóôõúüçÁÀÂÃÉÊÍÏÓÔÕÚÜÇ\"\n\nVOCABS[\"quechua\"] = re.sub(r\"[BbDdFfGgJjVvXxZz]\", \"\", VOCABS[\"english\"]) + \"ñÑĉĈçÇ\"\n\nVOCABS[\"romanian\"] = VOCABS[\"english\"] + \"ăâîșțĂÂÎȘȚ\"\n\nVOCABS[\"scottish_gaelic\"] = re.sub(r\"[JjKkQqVvWwXxYyZz]\", \"\", VOCABS[\"english\"]) + \"àèìòùÀÈÌÒÙ\"\n\nVOCABS[\"serbian_latin\"] = VOCABS[\"english\"] + \"čćđžšČĆĐŽŠ\"\n\nVOCABS[\"slovak\"] = VOCABS[\"english\"] + \"ôäčďľňšťžáéíĺóŕúýÔÄČĎĽŇŠŤŽÁÉÍĹÓŔÚÝ\"\n\nVOCABS[\"slovene\"] = re.sub(r\"[QqWwXxYy]\", \"\", VOCABS[\"english\"]) + \"čćđšžČĆĐŠŽ\"\n\nVOCABS[\"somali\"] = re.sub(r\"[PpVvZz]\", \"\", VOCABS[\"english\"])\n\nVOCABS[\"spanish\"] = VOCABS[\"english\"] + \"áéíóúüñÁÉÍÓÚÜÑ\" + \"¡¿\"\n\nVOCABS[\"swahili\"] = re.sub(r\"[QqXx]\", \"\", VOCABS[\"english\"])\n\nVOCABS[\"swedish\"] = VOCABS[\"english\"] + \"åäöÅÄÖ\"\n\nVOCABS[\"tagalog\"] = re.sub(r\"[CcQqWwXx]\", \"\", VOCABS[\"english\"]) + \"ñÑ\" + \"₱\"\n\nVOCABS[\"turkish\"] = re.sub(r\"[QqWwXx]\", \"\", VOCABS[\"english\"]) + \"çğıöşüâîûÇĞİÖŞÜÂÎÛ\" + \"₺\"\n\nVOCABS[\"uzbek_latin\"] = re.sub(r\"[Ww]\", \"\", VOCABS[\"english\"]) + \"çğɉñöşÇĞɈÑÖŞ\"\n\nVOCABS[\"vietnamese\"] = (\n VOCABS[\"english\"]\n + \"áàảạãăắằẳẵặâấầẩẫậđéèẻẽẹêếềểễệóòỏõọôốồổộỗơớờởợỡúùủũụưứừửữựíìỉĩịýỳỷỹỵ\"\n + \"ÁÀẢẠÃĂẮẰẲẴẶÂẤẦẨẪẬĐÉÈẺẼẸÊẾỀỂỄỆÓÒỎÕỌÔỐỒỔỘỖƠỚỜỞỢỠÚÙỦŨỤƯỨỪỬỮỰÍÌỈĨỊÝỲỶỸỴ\"\n + \"₫\" # currency\n)\n\nVOCABS[\"welsh\"] = re.sub(r\"[KkQqVvXxZz]\", \"\", VOCABS[\"english\"]) + \"âêîôŵŷÂÊÎÔŴŶ\"\n\nVOCABS[\"yoruba\"] = re.sub(r\"[CcQqVvXxZz]\", \"\", VOCABS[\"english\"]) + \"ẹọṣẸỌṢ\" + \"₦\"\n\nVOCABS[\"zulu\"] = VOCABS[\"english\"]\n\n# Non-latin alphabets.\n\n# Cyrillic\nVOCABS[\"russian\"] = (\n _BASE_VOCABS[\"generic_cyrillic_letters\"]\n + _BASE_VOCABS[\"russian_cyrillic_letters\"]\n + _BASE_VOCABS[\"russian_signs\"]\n + _BASE_VOCABS[\"digits\"]\n + _BASE_VOCABS[\"punctuation\"]\n + _BASE_VOCABS[\"currency\"]\n + \"₽\"\n)\n\nVOCABS[\"belarusian\"] = (\n _BASE_VOCABS[\"generic_cyrillic_letters\"]\n + _BASE_VOCABS[\"russian_cyrillic_letters\"]\n + _BASE_VOCABS[\"digits\"]\n + _BASE_VOCABS[\"punctuation\"]\n + _BASE_VOCABS[\"currency\"]\n + \"ўiЎI\"\n + \"₽\"\n)\n\nVOCABS[\"ukrainian\"] = (\n _BASE_VOCABS[\"generic_cyrillic_letters\"]\n + _BASE_VOCABS[\"digits\"]\n + _BASE_VOCABS[\"punctuation\"]\n + _BASE_VOCABS[\"currency\"]\n + \"ґіїєҐІЇЄ\"\n + \"₴\"\n)\n\nVOCABS[\"tatar\"] = VOCABS[\"russian\"] + \"ӘәҖҗҢңӨөҮү\"\n\nVOCABS[\"tajik\"] = VOCABS[\"russian\"].replace(\"₽\", \"\") + \"ҒғҚқҲҳҶҷӢӣӮӯ\"\n\nVOCABS[\"kazakh\"] = VOCABS[\"russian\"].replace(\"₽\", \"\") + \"ӘәҒғҚқҢңӨөҰұҮүҺһІі\" + \"₸\"\n\nVOCABS[\"kyrgyz\"] = VOCABS[\"russian\"].replace(\"₽\", \"\") + \"ҢңӨөҮү\"\n\nVOCABS[\"bulgarian\"] = (\n _BASE_VOCABS[\"generic_cyrillic_letters\"]\n + _BASE_VOCABS[\"russian_signs\"]\n + _BASE_VOCABS[\"digits\"]\n + _BASE_VOCABS[\"punctuation\"]\n + _BASE_VOCABS[\"currency\"]\n)\n\nVOCABS[\"macedonian\"] = (\n _BASE_VOCABS[\"generic_cyrillic_letters\"]\n + _BASE_VOCABS[\"digits\"]\n + _BASE_VOCABS[\"punctuation\"]\n + _BASE_VOCABS[\"currency\"]\n + \"ЃѓЅѕЈјЉљЊњЌќЏџ\"\n)\n\nVOCABS[\"mongolian\"] = (\n _BASE_VOCABS[\"generic_cyrillic_letters\"]\n + _BASE_VOCABS[\"russian_cyrillic_letters\"]\n + _BASE_VOCABS[\"russian_signs\"]\n + _BASE_VOCABS[\"digits\"]\n + _BASE_VOCABS[\"punctuation\"]\n + _BASE_VOCABS[\"currency\"]\n + \"ӨөҮү\"\n + \"᠐᠑᠒᠓᠔᠕᠖᠗᠘᠙\" # Mongolian digits\n + \"₮\"\n)\n\nVOCABS[\"yakut\"] = (\n _BASE_VOCABS[\"generic_cyrillic_letters\"]\n + _BASE_VOCABS[\"russian_cyrillic_letters\"]\n + _BASE_VOCABS[\"russian_signs\"]\n + _BASE_VOCABS[\"digits\"]\n + _BASE_VOCABS[\"punctuation\"]\n + _BASE_VOCABS[\"currency\"]\n + \"ҔҕҤҥӨөҺһҮү\"\n + \"₽\"\n)\n\nVOCABS[\"serbian_cyrillic\"] = (\n \"абвгдежзиклмнопрстуфхцчшАБВГДЕЖЗИКЛМНОПРСТУФХЦЧШ\" # limited cyrillic\n + \"JjЂђЉљЊњЋћЏџ\" # Serbian specials\n + _BASE_VOCABS[\"digits\"]\n + _BASE_VOCABS[\"punctuation\"]\n + _BASE_VOCABS[\"currency\"]\n)\n\nVOCABS[\"uzbek_cyrillic\"] = (\n _BASE_VOCABS[\"generic_cyrillic_letters\"]\n + _BASE_VOCABS[\"russian_cyrillic_letters\"]\n + _BASE_VOCABS[\"russian_signs\"]\n + _BASE_VOCABS[\"digits\"]\n + _BASE_VOCABS[\"punctuation\"]\n + _BASE_VOCABS[\"currency\"]\n + \"ЎўҚқҒғҲҳ\"\n)\n\nVOCABS[\"ukrainian\"] = (\n _BASE_VOCABS[\"generic_cyrillic_letters\"]\n + _BASE_VOCABS[\"digits\"]\n + _BASE_VOCABS[\"punctuation\"]\n + _BASE_VOCABS[\"currency\"]\n + \"ґіїєҐІЇЄ₴\"\n)\n\n# Greek\nVOCABS[\"greek\"] = (\n _BASE_VOCABS[\"punctuation\"] + _BASE_VOCABS[\"ancient_greek\"] + _BASE_VOCABS[\"currency\"] + \"άέήίϊΐόύϋΰώΆΈΉΊΪΌΎΫΏ\"\n)\nVOCABS[\"greek_extended\"] = (\n VOCABS[\"greek\"]\n + \"ͶͷϜϝἀἁἂἃἄἅἆἇἈἉἊἋἌἍἎἏἐἑἒἓἔἕἘἙἚἛἜἝἠἡἢἣἤἥἦἧἨἩἪἫἬἭἮἯἰἱἲἳἴἵἶἷἸἹἺἻἼἽἾἿ\"\n + \"ὀὁὂὃὄὅὈὉὊὋὌὍὐὑὒὓὔὕὖὗὙὛὝὟὠὡὢὣὤὥὦὧὨὩὪὫὬὭὮὯὰὲὴὶὸὺὼᾀᾁᾂᾃᾄᾅᾆᾇᾈᾉᾊᾋᾌᾍᾎᾏᾐ\"\n + \"ᾑᾒᾓᾔᾕᾖᾗᾘᾙᾚᾛᾜᾝᾞᾟᾠᾡᾢᾣᾤᾥᾦᾧᾨᾩᾪᾫᾬᾭᾮᾯᾲᾳᾴᾶᾷᾺᾼῂῃῄῆῇῈῊῌῒΐῖῗῚῢΰῤῥῦῧῪῬῲῳῴῶῷῸῺῼ\"\n)\n\n# Hebrew\nVOCABS[\"hebrew\"] = (\n _BASE_VOCABS[\"digits\"]\n + _BASE_VOCABS[\"punctuation\"]\n + _BASE_VOCABS[\"hebrew_consonants\"]\n + _BASE_VOCABS[\"hebrew_vowels\"]\n + _BASE_VOCABS[\"hebrew_punctuation\"]\n + _BASE_VOCABS[\"hebrew_cantillations\"]\n + _BASE_VOCABS[\"hebrew_specials\"]\n + \"₪\"\n)\n\n# Arabic\nVOCABS[\"arabic\"] = (\n _BASE_VOCABS[\"digits\"]\n + _BASE_VOCABS[\"arabic_digits\"]\n + _BASE_VOCABS[\"arabic_letters\"]\n + _BASE_VOCABS[\"persian_letters\"]\n + _BASE_VOCABS[\"arabic_diacritics\"]\n + _BASE_VOCABS[\"arabic_punctuation\"]\n + _BASE_VOCABS[\"punctuation\"]\n)\n\nVOCABS[\"persian\"] = VOCABS[\"arabic\"]\n\nVOCABS[\"urdu\"] = VOCABS[\"persian\"] + \"ٹڈڑںھےہۃ\"\n\nVOCABS[\"pashto\"] = VOCABS[\"persian\"] + \"ټډړږښځڅڼېۍ\"\n\nVOCABS[\"kurdish\"] = VOCABS[\"persian\"] + \"ڵڕۆێە\"\n\nVOCABS[\"uyghur\"] = VOCABS[\"persian\"] + \"ەېۆۇۈڭھ\"\n\nVOCABS[\"sindhi\"] = VOCABS[\"persian\"] + \"ڀٿٺٽڦڄڃڇڏڌڊڍڙڳڱڻھ\"\n\n# Indic scripts\n# Rules:\n# Any consonant can be \"combined\" with any matra\n# The virama is used to create consonant clusters - so C + Virama + C = CC\n\n# Devanagari based\nVOCABS[\"devanagari\"] = (\n _BASE_VOCABS[\"devanagari_consonants\"]\n + _BASE_VOCABS[\"devanagari_vowels\"]\n + _BASE_VOCABS[\"devanagari_digits\"]\n + _BASE_VOCABS[\"devanagari_matras\"]\n + _BASE_VOCABS[\"devanagari_virama\"]\n + _BASE_VOCABS[\"devanagari_punctuation\"]\n + _BASE_VOCABS[\"punctuation\"] # western punctuation used in Devanagari\n + \"₹\" # currency\n)\n\nVOCABS[\"hindi\"] = VOCABS[\"devanagari\"]\n\nVOCABS[\"sanskrit\"] = VOCABS[\"devanagari\"]\n\nVOCABS[\"marathi\"] = VOCABS[\"devanagari\"]\n\nVOCABS[\"nepali\"] = VOCABS[\"devanagari\"]\n\n# Gujarati\nVOCABS[\"gujarati\"] = (\n _BASE_VOCABS[\"gujarati_consonants\"]\n + _BASE_VOCABS[\"gujarati_vowels\"]\n + _BASE_VOCABS[\"gujarati_digits\"]\n + _BASE_VOCABS[\"gujarati_matras\"]\n + _BASE_VOCABS[\"gujarati_virama\"]\n + _BASE_VOCABS[\"gujarati_punctuation\"]\n + _BASE_VOCABS[\"punctuation\"] # western punctuation used in Gujarati\n + _BASE_VOCABS[\"gujarati_signs\"]\n + \"૱\" # currency\n)\n\n# Bengali\nVOCABS[\"bengali\"] = (\n _BASE_VOCABS[\"bengali_consonants\"]\n + _BASE_VOCABS[\"bengali_vowels\"]\n + _BASE_VOCABS[\"bengali_digits\"]\n + _BASE_VOCABS[\"bengali_matras\"]\n + _BASE_VOCABS[\"bengali_virama\"]\n + _BASE_VOCABS[\"bengali_punctuation\"]\n + _BASE_VOCABS[\"punctuation\"] # western punctuation used in Bengali\n + _BASE_VOCABS[\"bengali_signs\"]\n + \"৳\" # currency\n)\n\n# Brahmic scripts\nVOCABS[\"tamil\"] = (\n _BASE_VOCABS[\"tamil_consonants\"]\n + _BASE_VOCABS[\"tamil_vowels\"]\n + _BASE_VOCABS[\"tamil_digits\"]\n + _BASE_VOCABS[\"tamil_matras\"]\n + _BASE_VOCABS[\"tamil_virama\"]\n + _BASE_VOCABS[\"tamil_punctuation\"]\n + _BASE_VOCABS[\"punctuation\"] # western punctuation used in Tamil\n + _BASE_VOCABS[\"tamil_fractions\"] # This is a Tamil-specific addition\n + _BASE_VOCABS[\"tamil_signs\"]\n + \"₹\" # currency\n)\n\nVOCABS[\"telugu\"] = (\n _BASE_VOCABS[\"telugu_consonants\"]\n + _BASE_VOCABS[\"telugu_vowels\"]\n + _BASE_VOCABS[\"telugu_digits\"]\n + _BASE_VOCABS[\"telugu_matras\"]\n + _BASE_VOCABS[\"telugu_virama\"]\n + _BASE_VOCABS[\"telugu_punctuation\"]\n + _BASE_VOCABS[\"punctuation\"] # western punctuation used in Telugu\n + _BASE_VOCABS[\"telugu_signs\"]\n + \"₹\" # currency\n)\n\nVOCABS[\"kannada\"] = (\n _BASE_VOCABS[\"kannada_consonants\"]\n + _BASE_VOCABS[\"kannada_vowels\"]\n + _BASE_VOCABS[\"kannada_digits\"]\n + _BASE_VOCABS[\"kannada_matras\"]\n + _BASE_VOCABS[\"kannada_virama\"]\n + _BASE_VOCABS[\"kannada_punctuation\"]\n + _BASE_VOCABS[\"punctuation\"] # western punctuation used in Kannada\n + _BASE_VOCABS[\"kannada_signs\"]\n + \"₹\" # currency\n)\n\nVOCABS[\"sinhala\"] = (\n _BASE_VOCABS[\"sinhala_consonants\"]\n + _BASE_VOCABS[\"sinhala_vowels\"]\n + _BASE_VOCABS[\"sinhala_digits\"]\n + _BASE_VOCABS[\"sinhala_matras\"]\n + _BASE_VOCABS[\"sinhala_virama\"]\n + _BASE_VOCABS[\"sinhala_punctuation\"]\n + _BASE_VOCABS[\"punctuation\"] # western punctuation used in Sinhala\n + _BASE_VOCABS[\"sinhala_signs\"]\n + \"₹\" # currency\n)\n\nVOCABS[\"malayalam\"] = (\n _BASE_VOCABS[\"malayalam_consonants\"]\n + _BASE_VOCABS[\"malayalam_vowels\"]\n + _BASE_VOCABS[\"malayalam_digits\"]\n + _BASE_VOCABS[\"malayalam_matras\"]\n + _BASE_VOCABS[\"malayalam_virama\"]\n + _BASE_VOCABS[\"punctuation\"] # western punctuation used in Malayalam\n + _BASE_VOCABS[\"malayalam_signs\"]\n + \"₹\" # currency\n)\n\nVOCABS[\"punjabi\"] = (\n _BASE_VOCABS[\"punjabi_consonants\"]\n + _BASE_VOCABS[\"punjabi_vowels\"]\n + _BASE_VOCABS[\"punjabi_digits\"]\n + _BASE_VOCABS[\"punjabi_matras\"]\n + _BASE_VOCABS[\"punjabi_virama\"]\n + _BASE_VOCABS[\"punjabi_punctuation\"]\n + _BASE_VOCABS[\"punctuation\"] # western punctuation used in Punjabi\n + _BASE_VOCABS[\"punjabi_signs\"]\n + \"₹\" # currency\n)\n\n\nVOCABS[\"odia\"] = (\n _BASE_VOCABS[\"odia_consonants\"]\n + _BASE_VOCABS[\"odia_vowels\"]\n + _BASE_VOCABS[\"odia_digits\"]\n + _BASE_VOCABS[\"odia_matras\"]\n + _BASE_VOCABS[\"odia_virama\"]\n + _BASE_VOCABS[\"odia_punctuation\"]\n + _BASE_VOCABS[\"punctuation\"] # western punctuation used in Odia\n + _BASE_VOCABS[\"odia_signs\"]\n + \"₹\" # currency\n)\n\nVOCABS[\"khmer\"] = (\n _BASE_VOCABS[\"khmer_consonants\"]\n + _BASE_VOCABS[\"khmer_vowels\"]\n + _BASE_VOCABS[\"khmer_digits\"]\n + _BASE_VOCABS[\"khmer_matras\"]\n + _BASE_VOCABS[\"khmer_virama\"]\n + _BASE_VOCABS[\"khmer_diacritics\"] # This is a Khmer-specific addition\n + _BASE_VOCABS[\"khmer_punctuation\"]\n + _BASE_VOCABS[\"punctuation\"] # western punctuation used in Khmer\n + \"៛\" # Cambodian currency\n)\n\n# Armenian\nVOCABS[\"armenian\"] = (\n \"ԱԲԳԴԵԶԷԸԹԺԻԼԽԾԿՀՁՂՃՄՅՆՇՈՉՊՋՌՍՎՏՐՑՒՓՔՕՖՙՠաբգդեզէըթժիլխծկհձղճմյնշոչպջռսվտրցւփքօֆևֈ\"\n + _BASE_VOCABS[\"digits\"]\n + _BASE_VOCABS[\"punctuation\"]\n + \"՚՛՜՝՞՟։֊\"\n + \"֏\"\n)\n\n# Sudanese\nVOCABS[\"sudanese\"] = (\n _BASE_VOCABS[\"digits\"]\n + _BASE_VOCABS[\"sudanese_digits\"]\n + _BASE_VOCABS[\"sudanese_consonants\"]\n + _BASE_VOCABS[\"sudanese_vowels\"]\n + _BASE_VOCABS[\"sudanese_diacritics\"]\n + _BASE_VOCABS[\"punctuation\"]\n)\n\n# Thai\n# Rules:\n# Diacritics are used to modify the consonants and vowels\nVOCABS[\"thai\"] = (\n _BASE_VOCABS[\"digits\"]\n + \"๐๑๒๓๔๕๖๗๘๙\"\n + _BASE_VOCABS[\"punctuation\"]\n + \"๏๚๛ๆฯ\"\n + \"กขฃคฅฆงจฉชซฌญฎฏฐฑฒณดตถทธนบปผฝพฟภมยรฤลฦวศษสหฬอฮ\" # Thai consonants\n + \"ะาำเแโใไๅ\" # Thai vowels\n + \" ัิีึืฺุู็่้๊๋์ํ๎\".replace(\" \", \"\")\n + \"฿\"\n)\n\nVOCABS[\"lao\"] = (\n _BASE_VOCABS[\"digits\"]\n + \"໐໑໒໓໔໕໖໗໘໙\"\n + _BASE_VOCABS[\"punctuation\"]\n + \"ໆໞໟຯ\"\n + \"ກຂຄຆງຈຉຊຌຍຎຏຐຑຒຓດຕຖທຘນບປຜຝພຟຠມຢຣລວຨຩສຫຬອຮ\" # Lao consonants\n + \"ະາຳຽເແໂໃໄ\" # Lao vowels\n + \"ໜໝ\" # Lao ligature\n + \"\".join([\"ັ\", \"ິ\", \"ີ\", \"ຶ\", \"ື\", \"ຸ\", \"ູ\", \"຺\", \"ົ\", \"ຼ\", \"່\", \"້\", \"໊\", \"໋\", \"໌\", \"ໍ\"])\n)\n\n# Burmese & Javanese\n\n# Rules:\n# - A syllable usually starts with a base consonant.\n# - Diacritics (sandhangan), which represent vowels and consonant modifications, are attached to the base consonant:\n# - Vowel signs (ꦴꦵꦶꦷꦸꦹꦺꦻꦼ) follow the consonant and determine the syllable's vowel sound.\n# - Medial signs like ꦿ (ra), ꦾ (ya), and ꦽ (vocalic r) modify the consonant cluster.\n# - The virama (꧀, called *pangkon*) suppresses the inherent vowel,\n# creating consonant clusters.\n# - Special signs like ꦀ (cecak), ꦁ (layar), ꦂ (cakra), and ꦃ (wignyan)\n# can appear before or after syllables to represent nasal or glottal finals.\n# - Independent vowels (ꦄꦅꦆꦇꦈꦉꦊꦋꦌꦍꦎ) can occur without a base consonant, especially at word/sentence starts.\n# - Use Unicode NFC normalization to ensure composed syllables render correctly.\n\nVOCABS[\"burmese\"] = (\n _BASE_VOCABS[\"digits\"]\n + _BASE_VOCABS[\"burmese_digits\"]\n + _BASE_VOCABS[\"burmese_consonants\"]\n + _BASE_VOCABS[\"burmese_vowels\"]\n + _BASE_VOCABS[\"burmese_diacritics\"]\n + _BASE_VOCABS[\"burmese_virama\"]\n + _BASE_VOCABS[\"burmese_punctuation\"]\n)\n\nVOCABS[\"javanese\"] = (\n _BASE_VOCABS[\"digits\"]\n + _BASE_VOCABS[\"javanese_digits\"]\n + _BASE_VOCABS[\"javanese_consonants\"]\n + _BASE_VOCABS[\"javanese_vowels\"]\n + _BASE_VOCABS[\"javanese_diacritics\"]\n + _BASE_VOCABS[\"javanese_virama\"]\n + _BASE_VOCABS[\"javanese_punctuation\"]\n + _BASE_VOCABS[\"punctuation\"] # western punctuation used in Javanese\n)\n\n# Georgian (Mkhedruli - modern)\nVOCABS[\"georgian\"] = (\n _BASE_VOCABS[\"digits\"]\n + \"ႠႡႢႣႤႥႦႧႨႩႪႫႬႭႮႯႰႱႲႳႴႵႶႷႸႹႺႻႼႽႾႿჀჁჂჃჄჅჇჍაბგდევზთიკლმნოპჟრსტუფქღყშჩცძწჭხჯჰჱჲჳჴჵჶჷჸჹჺჼჽჾჿ\"\n + _BASE_VOCABS[\"punctuation\"]\n + \"჻\"\n + \"₾\" # currency\n)\n\n# Ethiopic\nVOCABS[\"ethiopic\"] = (\n \"ሀሁሂሃሄህሆሇለሉሊላሌልሎሏሐሑሒሓሔሕሖሗመሙሚማሜምሞሟሠሡሢሣሤሥሦሧረሩሪራሬርሮሯሰሱሲሳሴስሶሷሸሹሺሻሼሽሾሿቀቁቂቃቄቅቆቇቈቊቋ\"\n + \"ቌቍቐቑቒቓቔቕቖቘቚቛቜቝበቡቢባቤብቦቧቨቩቪቫቬቭቮቯተቱቲታቴትቶቷቸቹቺቻቼችቾቿኀኁኂኃኄኅኆኇኈኊኋኌኍነኑኒናኔንኖኗኘኙኚኛኜኝኞኟአኡኢኣኤእኦኧ\"\n + \"ከኩኪካኬክኮኯኰኲኳኴኵኸኹኺኻኼኽኾዀዂዃዄዅወዉዊዋዌውዎዏዐዑዒዓዔዕዖዘዙዚዛዜዝዞዟዠዡዢዣዤዥዦዧየዩዪያዬይዮዯደዱዲዳዴድዶዷዸዹዺ\"\n + \"ዻዼዽዾዿጀጁጂጃጄጅጆጇገጉጊጋጌግጎጏጐጒጓጔጕጘጙጚጛጜጝጞጟጠጡጢጣጤጥጦጧጨጩጪጫጬጭጮጯጰጱጲጳጴጵጶጷጸጹጺጻጼጽጾጿፀፁፂፃፄፅፆ\"\n + \"ፇፈፉፊፋፌፍፎፏፐፑፒፓፔፕፖፗፘፙፚᎀᎁᎂᎃᎄᎅᎆᎇᎈᎉᎊᎋᎌᎍᎎᎏ\"\n + \"፩፪፫፬፭፮፯፰፱፲፳፴፵፶፷፸፹፺፻፼\" # digits\n)\n\n# East Asian\nVOCABS[\"japanese\"] = (\n _BASE_VOCABS[\"digits\"]\n + \"ぁあぃいぅうぇえぉおかがきぎくぐけげこごさざしじすずせぜそぞただちぢっつづ\"\n + \"てでとどなにぬねのはばぱひびぴふぶぷへべぺほぼぽまみむめ\"\n + \"もゃやゅゆょよらりるれろゎわゐゑをんゔゕゖゝゞゟ\" # Hiragana\n + \"ァアィイゥウェエォオカガキギクグケゲコゴサザシジスズセゼソゾタダ\"\n + \"チヂッツヅテデトドナニヌネノハバパヒビピフブプヘベペホボポマミムメ\"\n + \"モャヤュユョヨラリルレロヮワヰヱヲンヴヵヶヷヸヹヺーヽヾヿ\" # Katakana\n # Kanji jōyō (incl. numerals)\n + \"亜哀挨愛曖悪握圧扱宛嵐安案暗以衣位囲医依委威為畏胃尉異移萎偉椅彙意違維慰遺緯域育一壱逸茨芋引印因咽姻員院淫陰飲隠韻右宇羽雨唄鬱畝浦運雲\" # noqa: E501\n + \"永泳英映栄営詠影鋭衛易疫益液駅悦越謁閲円延沿炎怨宴媛援園煙猿遠鉛塩演縁艶汚王凹央応往押旺欧殴桜翁奥横岡屋億憶臆虞乙俺卸音恩温穏下化火加\" # noqa: E501\n + \"可仮何花佳価果河苛科架夏家荷華菓貨渦過嫁暇禍靴寡歌箇稼課蚊牙瓦我画芽賀雅餓介回灰会快戒改怪拐悔海界皆械絵開階塊楷解潰壊懐諧貝外劾害崖涯\" # noqa: E501\n + \"街慨蓋該概骸垣柿各角拡革格核殻郭覚較隔閣確獲嚇穫学岳楽額顎掛潟括活喝渇割葛滑褐轄且株釜鎌刈干刊甘汗缶完肝官冠巻看陥乾勘患貫寒喚堪換敢棺\" # noqa: E501\n + \"款間閑勧寛幹感漢慣管関歓監緩憾還館環簡観韓艦鑑丸含岸岩玩眼頑顔願企伎危机気岐希忌汽奇祈季紀軌既記起飢鬼帰基寄規亀喜幾揮期棋貴棄毀旗器畿\" # noqa: E501\n + \"輝機騎技宜偽欺義疑儀戯擬犠議菊吉喫詰却客脚逆虐九久及弓丘旧休吸朽臼求究泣急級糾宮救球給嗅窮牛去巨居拒拠挙虚許距魚御漁凶共叫狂京享供協況\" # noqa: E501\n + \"峡挟狭恐恭胸脅強教郷境橋矯鏡競響驚仰暁業凝曲局極玉巾斤均近金菌勤琴筋僅禁緊錦謹襟吟銀区句苦駆具惧愚空偶遇隅串屈掘窟熊繰君訓勲薫軍郡群兄\" # noqa: E501\n + \"刑形系径茎係型契計恵啓掲渓経蛍敬景軽傾携継詣慶憬稽憩警鶏芸迎鯨隙劇撃激桁欠穴血決結傑潔月犬件見券肩建研県倹兼剣拳軒健険圏堅検嫌献絹遣権\" # noqa: E501\n + \"憲賢謙鍵繭顕験懸元幻玄言弦限原現舷減源厳己戸古呼固股虎孤弧故枯個庫湖雇誇鼓錮顧五互午呉後娯悟碁語誤護口工公勾孔功巧広甲交光向后好江考行\" # noqa: E501\n + \"坑孝抗攻更効幸拘肯侯厚恒洪皇紅荒郊香候校耕航貢降高康控梗黄喉慌港硬絞項溝鉱構綱酵稿興衡鋼講購乞号合拷剛傲豪克告谷刻国黒穀酷獄骨駒込頃今\" # noqa: E501\n + \"困昆恨根婚混痕紺魂墾懇左佐沙査砂唆差詐鎖座挫才再災妻采砕宰栽彩採済祭斎細菜最裁債催塞歳載際埼在材剤財罪崎作削昨柵索策酢搾錯咲冊札刷刹拶\" # noqa: E501\n + \"殺察撮擦雑皿三山参桟蚕惨産傘散算酸賛残斬暫士子支止氏仕史司四市矢旨死糸至伺志私使刺始姉枝祉肢姿思指施師恣紙脂視紫詞歯嗣試詩資飼誌雌摯賜\" # noqa: E501\n + \"諮示字寺次耳自似児事侍治持時滋慈辞磁餌璽鹿式識軸七叱失室疾執湿嫉漆質実芝写社車舎者射捨赦斜煮遮謝邪蛇尺借酌釈爵若弱寂手主守朱取狩首殊珠\" # noqa: E501\n + \"酒腫種趣寿受呪授需儒樹収囚州舟秀周宗拾秋臭修袖終羞習週就衆集愁酬醜蹴襲十汁充住柔重従渋銃獣縦叔祝宿淑粛縮塾熟出述術俊春瞬旬巡盾准殉純循\" # noqa: E501\n + \"順準潤遵処初所書庶暑署緒諸女如助序叙徐除小升少召匠床抄肖尚招承昇松沼昭宵将消症祥称笑唱商渉章紹訟勝掌晶焼焦硝粧詔証象傷奨照詳彰障憧衝賞\" # noqa: E501\n + \"償礁鐘上丈冗条状乗城浄剰常情場畳蒸縄壌嬢錠譲醸色拭食植殖飾触嘱織職辱尻心申伸臣芯身辛侵信津神唇娠振浸真針深紳進森診寝慎新審震薪親人刃仁\" # noqa: E501\n + \"尽迅甚陣尋腎須図水吹垂炊帥粋衰推酔遂睡穂随髄枢崇数据杉裾寸瀬是井世正生成西声制姓征性青斉政星牲省凄逝清盛婿晴勢聖誠精製誓静請整醒税夕斥\" # noqa: E501\n + \"石赤昔析席脊隻惜戚責跡積績籍切折拙窃接設雪摂節説舌絶千川仙占先宣専泉浅洗染扇栓旋船戦煎羨腺詮践箋銭潜線遷選薦繊鮮全前善然禅漸膳繕狙阻祖\" # noqa: E501\n + \"租素措粗組疎訴塑遡礎双壮早争走奏相荘草送倉捜挿桑巣掃曹曽爽窓創喪痩葬装僧想層総遭槽踪操燥霜騒藻造像増憎蔵贈臓即束足促則息捉速側測俗族属\" # noqa: E501\n + \"賊続卒率存村孫尊損遜他多汰打妥唾堕惰駄太対体耐待怠胎退帯泰堆袋逮替貸隊滞態戴大代台第題滝宅択沢卓拓託濯諾濁但達脱奪棚誰丹旦担単炭胆探淡\" # noqa: E501\n + \"短嘆端綻誕鍛団男段断弾暖談壇地池知値恥致遅痴稚置緻竹畜逐蓄築秩窒茶着嫡中仲虫沖宙忠抽注昼柱衷酎鋳駐著貯丁弔庁兆町長挑帳張彫眺釣頂鳥朝貼\" # noqa: E501\n + \"超腸跳徴嘲潮澄調聴懲直勅捗沈珍朕陳賃鎮追椎墜通痛塚漬坪爪鶴低呈廷弟定底抵邸亭貞帝訂庭逓停偵堤提程艇締諦泥的笛摘滴適敵溺迭哲鉄徹撤天典店\" # noqa: E501\n + \"点展添転塡田伝殿電斗吐妬徒途都渡塗賭土奴努度怒刀冬灯当投豆東到逃倒凍唐島桃討透党悼盗陶塔搭棟湯痘登答等筒統稲踏糖頭謄藤闘騰同洞胴動堂童\" # noqa: E501\n + \"道働銅導瞳峠匿特得督徳篤毒独読栃凸突届屯豚頓貪鈍曇丼那奈内梨謎鍋南軟難二尼弐匂肉虹日入乳尿任妊忍認寧熱年念捻粘燃悩納能脳農濃把波派破覇\" # noqa: E501\n + \"馬婆罵拝杯背肺俳配排敗廃輩売倍梅培陪媒買賠白伯拍泊迫剝舶博薄麦漠縛爆箱箸畑肌八鉢発髪伐抜罰閥反半氾犯帆汎伴判坂阪板版班畔般販斑飯搬煩頒\" # noqa: E501\n + \"範繁藩晩番蛮盤比皮妃否批彼披肥非卑飛疲秘被悲扉費碑罷避尾眉美備微鼻膝肘匹必泌筆姫百氷表俵票評漂標苗秒病描猫品浜貧賓頻敏瓶不夫父付布扶府\" # noqa: E501\n + \"怖阜附訃負赴浮婦符富普腐敷膚賦譜侮武部舞封風伏服副幅復福腹複覆払沸仏物粉紛雰噴墳憤奮分文聞丙平兵併並柄陛閉塀幣弊蔽餅米壁璧癖別蔑片辺返\" # noqa: E501\n + \"変偏遍編弁便勉歩保哺捕補舗母募墓慕暮簿方包芳邦奉宝抱放法泡胞俸倣峰砲崩訪報蜂豊飽褒縫亡乏忙坊妨忘防房肪某冒剖紡望傍帽棒貿貌暴膨謀頰北木\" # noqa: E501\n + \"朴牧睦僕墨撲没勃堀本奔翻凡盆麻摩磨魔毎妹枚昧埋幕膜枕又末抹万満慢漫未味魅岬密蜜脈妙民眠矛務無夢霧娘名命明迷冥盟銘鳴滅免面綿麺茂模毛妄盲\" # noqa: E501\n + \"耗猛網目黙門紋問冶夜野弥厄役約訳薬躍闇由油喩愉諭輸癒唯友有勇幽悠郵湧猶裕遊雄誘憂融優与予余誉預幼用羊妖洋要容庸揚揺葉陽溶腰様瘍踊窯養擁\" # noqa: E501\n + \"謡曜抑沃浴欲翌翼拉裸羅来雷頼絡落酪辣乱卵覧濫藍欄吏利里理痢裏履璃離陸立律慄略柳流留竜粒隆硫侶旅虜慮了両良料涼猟陵量僚領寮療瞭糧力緑林厘\" # noqa: E501\n + \"倫輪隣臨瑠涙累塁類令礼冷励戻例鈴零霊隷齢麗暦歴列劣烈裂恋連廉練錬呂炉賂路露老労弄郎朗浪廊楼漏籠六録麓論和話賄脇惑枠湾腕\" # noqa: E501\n + _BASE_VOCABS[\"punctuation\"]\n + \"。・〜°—、「」『』【】゛》《〉〈\"\n + _BASE_VOCABS[\"currency\"]\n)\n\nVOCABS[\"korean\"] = (\n _BASE_VOCABS[\"digits\"]\n + \"가각갂갃간갅갆갇갈갉갊갋갌갍갎갏감갑값갓갔강갖갗갘같갚갛개객갞갟갠갡갢갣갤갥갦갧갨갩갪갫갬갭갮갯갰갱갲갳갴갵갶갷갸갹갺갻갼갽갾갿걀걁걂걃걄걅걆걇걈\" # noqa: E501\n + \"걉걊걋걌걍걎걏걐걑걒걓걔걕걖걗걘걙걚걛걜걝걞걟걠걡걢걣걤걥걦걧걨걩걪걫걬걭걮걯거걱걲걳건걵걶걷걸걹걺걻걼걽걾걿검겁겂것겄겅겆겇겈겉겊겋게겍겎겏겐겑\" # noqa: E501\n + \"겒겓겔겕겖겗겘겙겚겛겜겝겞겟겠겡겢겣겤겥겦겧겨격겪겫견겭겮겯결겱겲겳겴겵겶겷겸겹겺겻겼경겾겿곀곁곂곃계곅곆곇곈곉곊곋곌곍곎곏곐곑곒곓곔곕곖곗곘곙곚\" # noqa: E501\n + \"곛곜곝곞곟고곡곢곣곤곥곦곧골곩곪곫곬곭곮곯곰곱곲곳곴공곶곷곸곹곺곻과곽곾곿관괁괂괃괄괅괆괇괈괉괊괋괌괍괎괏괐광괒괓괔괕괖괗괘괙괚괛괜괝괞괟괠괡괢괣\" # noqa: E501\n + \"괤괥괦괧괨괩괪괫괬괭괮괯괰괱괲괳괴괵괶괷괸괹괺괻괼괽괾괿굀굁굂굃굄굅굆굇굈굉굊굋굌굍굎굏교굑굒굓굔굕굖굗굘굙굚굛굜굝굞굟굠굡굢굣굤굥굦굧굨굩굪굫구\" # noqa: E501\n + \"국굮굯군굱굲굳굴굵굶굷굸굹굺굻굼굽굾굿궀궁궂궃궄궅궆궇궈궉궊궋권궍궎궏궐궑궒궓궔궕궖궗궘궙궚궛궜궝궞궟궠궡궢궣궤궥궦궧궨궩궪궫궬궭궮궯궰궱궲궳궴궵\" # noqa: E501\n + \"궶궷궸궹궺궻궼궽궾궿귀귁귂귃귄귅귆귇귈귉귊귋귌귍귎귏귐귑귒귓귔귕귖귗귘귙귚귛규귝귞귟균귡귢귣귤귥귦귧귨귩귪귫귬귭귮귯귰귱귲귳귴귵귶귷그극귺귻근귽귾\" # noqa: E501\n + \"귿글긁긂긃긄긅긆긇금급긊긋긌긍긎긏긐긑긒긓긔긕긖긗긘긙긚긛긜긝긞긟긠긡긢긣긤긥긦긧긨긩긪긫긬긭긮긯기긱긲긳긴긵긶긷길긹긺긻긼긽긾긿김깁깂깃깄깅깆깇\" # noqa: E501\n + \"깈깉깊깋까깍깎깏깐깑깒깓깔깕깖깗깘깙깚깛깜깝깞깟깠깡깢깣깤깥깦깧깨깩깪깫깬깭깮깯깰깱깲깳깴깵깶깷깸깹깺깻깼깽깾깿꺀꺁꺂꺃꺄꺅꺆꺇꺈꺉꺊꺋꺌꺍꺎꺏꺐\" # noqa: E501\n + \"꺑꺒꺓꺔꺕꺖꺗꺘꺙꺚꺛꺜꺝꺞꺟꺠꺡꺢꺣꺤꺥꺦꺧꺨꺩꺪꺫꺬꺭꺮꺯꺰꺱꺲꺳꺴꺵꺶꺷꺸꺹꺺꺻꺼꺽꺾꺿껀껁껂껃껄껅껆껇껈껉껊껋껌껍껎껏껐껑껒껓껔껕껖껗께껙\" # noqa: E501\n + \"껚껛껜껝껞껟껠껡껢껣껤껥껦껧껨껩껪껫껬껭껮껯껰껱껲껳껴껵껶껷껸껹껺껻껼껽껾껿꼀꼁꼂꼃꼄꼅꼆꼇꼈꼉꼊꼋꼌꼍꼎꼏꼐꼑꼒꼓꼔꼕꼖꼗꼘꼙꼚꼛꼜꼝꼞꼟꼠꼡꼢\" # noqa: E501\n + \"꼣꼤꼥꼦꼧꼨꼩꼪꼫꼬꼭꼮꼯꼰꼱꼲꼳꼴꼵꼶꼷꼸꼹꼺꼻꼼꼽꼾꼿꽀꽁꽂꽃꽄꽅꽆꽇꽈꽉꽊꽋꽌꽍꽎꽏꽐꽑꽒꽓꽔꽕꽖꽗꽘꽙꽚꽛꽜꽝꽞꽟꽠꽡꽢꽣꽤꽥꽦꽧꽨꽩꽪꽫\" # noqa: E501\n + \"꽬꽭꽮꽯꽰꽱꽲꽳꽴꽵꽶꽷꽸꽹꽺꽻꽼꽽꽾꽿꾀꾁꾂꾃꾄꾅꾆꾇꾈꾉꾊꾋꾌꾍꾎꾏꾐꾑꾒꾓꾔꾕꾖꾗꾘꾙꾚꾛꾜꾝꾞꾟꾠꾡꾢꾣꾤꾥꾦꾧꾨꾩꾪꾫꾬꾭꾮꾯꾰꾱꾲꾳꾴\" # noqa: E501\n + \"꾵꾶꾷꾸꾹꾺꾻꾼꾽꾾꾿꿀꿁꿂꿃꿄꿅꿆꿇꿈꿉꿊꿋꿌꿍꿎꿏꿐꿑꿒꿓꿔꿕꿖꿗꿘꿙꿚꿛꿜꿝꿞꿟꿠꿡꿢꿣꿤꿥꿦꿧꿨꿩꿪꿫꿬꿭꿮꿯꿰꿱꿲꿳꿴꿵꿶꿷꿸꿹꿺꿻꿼꿽\" # noqa: E501\n + \"꿾꿿뀀뀁뀂뀃뀄뀅뀆뀇뀈뀉뀊뀋뀌뀍뀎뀏뀐뀑뀒뀓뀔뀕뀖뀗뀘뀙뀚뀛뀜뀝뀞뀟뀠뀡뀢뀣뀤뀥뀦뀧뀨뀩뀪뀫뀬뀭뀮뀯뀰뀱뀲뀳뀴뀵뀶뀷뀸뀹뀺뀻뀼뀽뀾뀿끀끁끂끃끄끅끆\" # noqa: E501\n + \"끇끈끉끊끋끌끍끎끏끐끑끒끓끔끕끖끗끘끙끚끛끜끝끞끟끠끡끢끣끤끥끦끧끨끩끪끫끬끭끮끯끰끱끲끳끴끵끶끷끸끹끺끻끼끽끾끿낀낁낂낃낄낅낆낇낈낉낊낋낌낍낎낏\" # noqa: E501\n + \"낐낑낒낓낔낕낖낗나낙낚낛난낝낞낟날낡낢낣낤낥낦낧남납낪낫났낭낮낯낰낱낲낳내낵낶낷낸낹낺낻낼낽낾낿냀냁냂냃냄냅냆냇냈냉냊냋냌냍냎냏냐냑냒냓냔냕냖냗냘\" # noqa: E501\n + \"냙냚냛냜냝냞냟냠냡냢냣냤냥냦냧냨냩냪냫냬냭냮냯냰냱냲냳냴냵냶냷냸냹냺냻냼냽냾냿넀넁넂넃넄넅넆넇너넉넊넋넌넍넎넏널넑넒넓넔넕넖넗넘넙넚넛넜넝넞넟넠넡\" # noqa: E501\n + \"넢넣네넥넦넧넨넩넪넫넬넭넮넯넰넱넲넳넴넵넶넷넸넹넺넻넼넽넾넿녀녁녂녃년녅녆녇녈녉녊녋녌녍녎녏념녑녒녓녔녕녖녗녘녙녚녛녜녝녞녟녠녡녢녣녤녥녦녧녨녩녪\" # noqa: E501\n + \"녫녬녭녮녯녰녱녲녳녴녵녶녷노녹녺녻논녽녾녿놀놁놂놃놄놅놆놇놈놉놊놋놌농놎놏놐놑높놓놔놕놖놗놘놙놚놛놜놝놞놟놠놡놢놣놤놥놦놧놨놩놪놫놬놭놮놯놰놱놲놳\" # noqa: E501\n + \"놴놵놶놷놸놹놺놻놼놽놾놿뇀뇁뇂뇃뇄뇅뇆뇇뇈뇉뇊뇋뇌뇍뇎뇏뇐뇑뇒뇓뇔뇕뇖뇗뇘뇙뇚뇛뇜뇝뇞뇟뇠뇡뇢뇣뇤뇥뇦뇧뇨뇩뇪뇫뇬뇭뇮뇯뇰뇱뇲뇳뇴뇵뇶뇷뇸뇹뇺뇻뇼\" # noqa: E501\n + \"뇽뇾뇿눀눁눂눃누눅눆눇눈눉눊눋눌눍눎눏눐눑눒눓눔눕눖눗눘눙눚눛눜눝눞눟눠눡눢눣눤눥눦눧눨눩눪눫눬눭눮눯눰눱눲눳눴눵눶눷눸눹눺눻눼눽눾눿뉀뉁뉂뉃뉄뉅\" # noqa: E501\n + \"뉆뉇뉈뉉뉊뉋뉌뉍뉎뉏뉐뉑뉒뉓뉔뉕뉖뉗뉘뉙뉚뉛뉜뉝뉞뉟뉠뉡뉢뉣뉤뉥뉦뉧뉨뉩뉪뉫뉬뉭뉮뉯뉰뉱뉲뉳뉴뉵뉶뉷뉸뉹뉺뉻뉼뉽뉾뉿늀늁늂늃늄늅늆늇늈늉늊늋늌늍늎\" # noqa: E501\n + \"늏느늑늒늓는늕늖늗늘늙늚늛늜늝늞늟늠늡늢늣늤능늦늧늨늩늪늫늬늭늮늯늰늱늲늳늴늵늶늷늸늹늺늻늼늽늾늿닀닁닂닃닄닅닆닇니닉닊닋닌닍닎닏닐닑닒닓닔닕닖닗\" # noqa: E501\n + \"님닙닚닛닜닝닞닟닠닡닢닣다닥닦닧단닩닪닫달닭닮닯닰닱닲닳담답닶닷닸당닺닻닼닽닾닿대댁댂댃댄댅댆댇댈댉댊댋댌댍댎댏댐댑댒댓댔댕댖댗댘댙댚댛댜댝댞댟댠\" # noqa: E501\n + \"댡댢댣댤댥댦댧댨댩댪댫댬댭댮댯댰댱댲댳댴댵댶댷댸댹댺댻댼댽댾댿덀덁덂덃덄덅덆덇덈덉덊덋덌덍덎덏덐덑덒덓더덕덖덗던덙덚덛덜덝덞덟덠덡덢덣덤덥덦덧덨덩\" # noqa: E501\n + \"덪덫덬덭덮덯데덱덲덳덴덵덶덷델덹덺덻덼덽덾덿뎀뎁뎂뎃뎄뎅뎆뎇뎈뎉뎊뎋뎌뎍뎎뎏뎐뎑뎒뎓뎔뎕뎖뎗뎘뎙뎚뎛뎜뎝뎞뎟뎠뎡뎢뎣뎤뎥뎦뎧뎨뎩뎪뎫뎬뎭뎮뎯뎰뎱뎲\" # noqa: E501\n + \"뎳뎴뎵뎶뎷뎸뎹뎺뎻뎼뎽뎾뎿돀돁돂돃도독돆돇돈돉돊돋돌돍돎돏돐돑돒돓돔돕돖돗돘동돚돛돜돝돞돟돠돡돢돣돤돥돦돧돨돩돪돫돬돭돮돯돰돱돲돳돴돵돶돷돸돹돺돻\" # noqa: E501\n + \"돼돽돾돿됀됁됂됃됄됅됆됇됈됉됊됋됌됍됎됏됐됑됒됓됔됕됖됗되됙됚됛된됝됞됟될됡됢됣됤됥됦됧됨됩됪됫됬됭됮됯됰됱됲됳됴됵됶됷됸됹됺됻됼됽됾됿둀둁둂둃둄\" # noqa: E501\n + \"둅둆둇둈둉둊둋둌둍둎둏두둑둒둓둔둕둖둗둘둙둚둛둜둝둞둟둠둡둢둣둤둥둦둧둨둩둪둫둬둭둮둯둰둱둲둳둴둵둶둷둸둹둺둻둼둽둾둿뒀뒁뒂뒃뒄뒅뒆뒇뒈뒉뒊뒋뒌뒍\" # noqa: E501\n + \"뒎뒏뒐뒑뒒뒓뒔뒕뒖뒗뒘뒙뒚뒛뒜뒝뒞뒟뒠뒡뒢뒣뒤뒥뒦뒧뒨뒩뒪뒫뒬뒭뒮뒯뒰뒱뒲뒳뒴뒵뒶뒷뒸뒹뒺뒻뒼뒽뒾뒿듀듁듂듃듄듅듆듇듈듉듊듋듌듍듎듏듐듑듒듓듔듕듖\" # noqa: E501\n + \"듗듘듙듚듛드득듞듟든듡듢듣들듥듦듧듨듩듪듫듬듭듮듯듰등듲듳듴듵듶듷듸듹듺듻듼듽듾듿딀딁딂딃딄딅딆딇딈딉딊딋딌딍딎딏딐딑딒딓디딕딖딗딘딙딚딛딜딝딞딟\" # noqa: E501\n + \"딠딡딢딣딤딥딦딧딨딩딪딫딬딭딮딯따딱딲딳딴딵딶딷딸딹딺딻딼딽딾딿땀땁땂땃땄땅땆땇땈땉땊땋때땍땎땏땐땑땒땓땔땕땖땗땘땙땚땛땜땝땞땟땠땡땢땣땤땥땦땧땨\" # noqa: E501\n + \"땩땪땫땬땭땮땯땰땱땲땳땴땵땶땷땸땹땺땻땼땽땾땿떀떁떂떃떄떅떆떇떈떉떊떋떌떍떎떏떐떑떒떓떔떕떖떗떘떙떚떛떜떝떞떟떠떡떢떣떤떥떦떧떨떩떪떫떬떭떮떯떰떱\" # noqa: E501\n + \"떲떳떴떵떶떷떸떹떺떻떼떽떾떿뗀뗁뗂뗃뗄뗅뗆뗇뗈뗉뗊뗋뗌뗍뗎뗏뗐뗑뗒뗓뗔뗕뗖뗗뗘뗙뗚뗛뗜뗝뗞뗟뗠뗡뗢뗣뗤뗥뗦뗧뗨뗩뗪뗫뗬뗭뗮뗯뗰뗱뗲뗳뗴뗵뗶뗷뗸뗹뗺\" # noqa: E501\n + \"뗻뗼뗽뗾뗿똀똁똂똃똄똅똆똇똈똉똊똋똌똍똎똏또똑똒똓똔똕똖똗똘똙똚똛똜똝똞똟똠똡똢똣똤똥똦똧똨똩똪똫똬똭똮똯똰똱똲똳똴똵똶똷똸똹똺똻똼똽똾똿뙀뙁뙂뙃\" # noqa: E501\n + \"뙄뙅뙆뙇뙈뙉뙊뙋뙌뙍뙎뙏뙐뙑뙒뙓뙔뙕뙖뙗뙘뙙뙚뙛뙜뙝뙞뙟뙠뙡뙢뙣뙤뙥뙦뙧뙨뙩뙪뙫뙬뙭뙮뙯뙰뙱뙲뙳뙴뙵뙶뙷뙸뙹뙺뙻뙼뙽뙾뙿뚀뚁뚂뚃뚄뚅뚆뚇뚈뚉뚊뚋뚌\" # noqa: E501\n + \"뚍뚎뚏뚐뚑뚒뚓뚔뚕뚖뚗뚘뚙뚚뚛뚜뚝뚞뚟뚠뚡뚢뚣뚤뚥뚦뚧뚨뚩뚪뚫뚬뚭뚮뚯뚰뚱뚲뚳뚴뚵뚶뚷뚸뚹뚺뚻뚼뚽뚾뚿뛀뛁뛂뛃뛄뛅뛆뛇뛈뛉뛊뛋뛌뛍뛎뛏뛐뛑뛒뛓뛔뛕\" # noqa: E501\n + \"뛖뛗뛘뛙뛚뛛뛜뛝뛞뛟뛠뛡뛢뛣뛤뛥뛦뛧뛨뛩뛪뛫뛬뛭뛮뛯뛰뛱뛲뛳뛴뛵뛶뛷뛸뛹뛺뛻뛼뛽뛾뛿뜀뜁뜂뜃뜄뜅뜆뜇뜈뜉뜊뜋뜌뜍뜎뜏뜐뜑뜒뜓뜔뜕뜖뜗뜘뜙뜚뜛뜜뜝뜞\" # noqa: E501\n + \"뜟뜠뜡뜢뜣뜤뜥뜦뜧뜨뜩뜪뜫뜬뜭뜮뜯뜰뜱뜲뜳뜴뜵뜶뜷뜸뜹뜺뜻뜼뜽뜾뜿띀띁띂띃띄띅띆띇띈띉띊띋띌띍띎띏띐띑띒띓띔띕띖띗띘띙띚띛띜띝띞띟띠띡띢띣띤띥띦띧\" # noqa: E501\n + \"띨띩띪띫띬띭띮띯띰띱띲띳띴띵띶띷띸띹띺띻라락띾띿란랁랂랃랄랅랆랇랈랉랊랋람랍랎랏랐랑랒랓랔랕랖랗래랙랚랛랜랝랞랟랠랡랢랣랤랥랦랧램랩랪랫랬랭랮랯랰\" # noqa: E501\n + \"랱랲랳랴략랶랷랸랹랺랻랼랽랾랿럀럁럂럃럄럅럆럇럈량럊럋럌럍럎럏럐럑럒럓럔럕럖럗럘럙럚럛럜럝럞럟럠럡럢럣럤럥럦럧럨럩럪럫러럭럮럯런럱럲럳럴럵럶럷럸럹\" # noqa: E501\n + \"럺럻럼럽럾럿렀렁렂렃렄렅렆렇레렉렊렋렌렍렎렏렐렑렒렓렔렕렖렗렘렙렚렛렜렝렞렟렠렡렢렣려력렦렧련렩렪렫렬렭렮렯렰렱렲렳렴렵렶렷렸령렺렻렼렽렾렿례롁롂\" # noqa: E501\n + \"롃롄롅롆롇롈롉롊롋롌롍롎롏롐롑롒롓롔롕롖롗롘롙롚롛로록롞롟론롡롢롣롤롥롦롧롨롩롪롫롬롭롮롯롰롱롲롳롴롵롶롷롸롹롺롻롼롽롾롿뢀뢁뢂뢃뢄뢅뢆뢇뢈뢉뢊뢋\" # noqa: E501\n + \"뢌뢍뢎뢏뢐뢑뢒뢓뢔뢕뢖뢗뢘뢙뢚뢛뢜뢝뢞뢟뢠뢡뢢뢣뢤뢥뢦뢧뢨뢩뢪뢫뢬뢭뢮뢯뢰뢱뢲뢳뢴뢵뢶뢷뢸뢹뢺뢻뢼뢽뢾뢿룀룁룂룃룄룅룆룇룈룉룊룋료룍룎룏룐룑룒룓룔\" # noqa: E501\n + \"룕룖룗룘룙룚룛룜룝룞룟룠룡룢룣룤룥룦룧루룩룪룫룬룭룮룯룰룱룲룳룴룵룶룷룸룹룺룻룼룽룾룿뤀뤁뤂뤃뤄뤅뤆뤇뤈뤉뤊뤋뤌뤍뤎뤏뤐뤑뤒뤓뤔뤕뤖뤗뤘뤙뤚뤛뤜뤝\" # noqa: E501\n + \"뤞뤟뤠뤡뤢뤣뤤뤥뤦뤧뤨뤩뤪뤫뤬뤭뤮뤯뤰뤱뤲뤳뤴뤵뤶뤷뤸뤹뤺뤻뤼뤽뤾뤿륀륁륂륃륄륅륆륇륈륉륊륋륌륍륎륏륐륑륒륓륔륕륖륗류륙륚륛륜륝륞륟률륡륢륣륤륥륦\" # noqa: E501\n + \"륧륨륩륪륫륬륭륮륯륰륱륲륳르륵륶륷른륹륺륻를륽륾륿릀릁릂릃름릅릆릇릈릉릊릋릌릍릎릏릐릑릒릓릔릕릖릗릘릙릚릛릜릝릞릟릠릡릢릣릤릥릦릧릨릩릪릫리릭릮릯\" # noqa: E501\n + \"린릱릲릳릴릵릶릷릸릹릺릻림립릾릿맀링맂맃맄맅맆맇마막맊맋만맍많맏말맑맒맓맔맕맖맗맘맙맚맛맜망맞맟맠맡맢맣매맥맦맧맨맩맪맫맬맭맮맯맰맱맲맳맴맵맶맷맸\" # noqa: E501\n + \"맹맺맻맼맽맾맿먀먁먂먃먄먅먆먇먈먉먊먋먌먍먎먏먐먑먒먓먔먕먖먗먘먙먚먛먜먝먞먟먠먡먢먣먤먥먦먧먨먩먪먫먬먭먮먯먰먱먲먳먴먵먶먷머먹먺먻먼먽먾먿멀멁\" # noqa: E501\n + \"멂멃멄멅멆멇멈멉멊멋멌멍멎멏멐멑멒멓메멕멖멗멘멙멚멛멜멝멞멟멠멡멢멣멤멥멦멧멨멩멪멫멬멭멮멯며멱멲멳면멵멶멷멸멹멺멻멼멽멾멿몀몁몂몃몄명몆몇몈몉몊\" # noqa: E501\n + \"몋몌몍몎몏몐몑몒몓몔몕몖몗몘몙몚몛몜몝몞몟몠몡몢몣몤몥몦몧모목몪몫몬몭몮몯몰몱몲몳몴몵몶몷몸몹몺못몼몽몾몿뫀뫁뫂뫃뫄뫅뫆뫇뫈뫉뫊뫋뫌뫍뫎뫏뫐뫑뫒뫓\" # noqa: E501\n + \"뫔뫕뫖뫗뫘뫙뫚뫛뫜뫝뫞뫟뫠뫡뫢뫣뫤뫥뫦뫧뫨뫩뫪뫫뫬뫭뫮뫯뫰뫱뫲뫳뫴뫵뫶뫷뫸뫹뫺뫻뫼뫽뫾뫿묀묁묂묃묄묅묆묇묈묉묊묋묌묍묎묏묐묑묒묓묔묕묖묗묘묙묚묛묜\" # noqa: E501\n + \"묝묞묟묠묡묢묣묤묥묦묧묨묩묪묫묬묭묮묯묰묱묲묳무묵묶묷문묹묺묻물묽묾묿뭀뭁뭂뭃뭄뭅뭆뭇뭈뭉뭊뭋뭌뭍뭎뭏뭐뭑뭒뭓뭔뭕뭖뭗뭘뭙뭚뭛뭜뭝뭞뭟뭠뭡뭢뭣뭤뭥\" # noqa: E501\n + \"뭦뭧뭨뭩뭪뭫뭬뭭뭮뭯뭰뭱뭲뭳뭴뭵뭶뭷뭸뭹뭺뭻뭼뭽뭾뭿뮀뮁뮂뮃뮄뮅뮆뮇뮈뮉뮊뮋뮌뮍뮎뮏뮐뮑뮒뮓뮔뮕뮖뮗뮘뮙뮚뮛뮜뮝뮞뮟뮠뮡뮢뮣뮤뮥뮦뮧뮨뮩뮪뮫뮬뮭뮮\" # noqa: E501\n + \"뮯뮰뮱뮲뮳뮴뮵뮶뮷뮸뮹뮺뮻뮼뮽뮾뮿므믁믂믃믄믅믆믇믈믉믊믋믌믍믎믏믐믑믒믓믔믕믖믗믘믙믚믛믜믝믞믟믠믡믢믣믤믥믦믧믨믩믪믫믬믭믮믯믰믱믲믳믴믵믶믷\" # noqa: E501\n + \"미믹믺믻민믽믾믿밀밁밂밃밄밅밆밇밈밉밊밋밌밍밎및밐밑밒밓바박밖밗반밙밚받발밝밞밟밠밡밢밣밤밥밦밧밨방밪밫밬밭밮밯배백밲밳밴밵밶밷밸밹밺밻밼밽밾밿뱀\" # noqa: E501\n + \"뱁뱂뱃뱄뱅뱆뱇뱈뱉뱊뱋뱌뱍뱎뱏뱐뱑뱒뱓뱔뱕뱖뱗뱘뱙뱚뱛뱜뱝뱞뱟뱠뱡뱢뱣뱤뱥뱦뱧뱨뱩뱪뱫뱬뱭뱮뱯뱰뱱뱲뱳뱴뱵뱶뱷뱸뱹뱺뱻뱼뱽뱾뱿벀벁벂벃버벅벆벇번벉\" # noqa: E501\n + \"벊벋벌벍벎벏벐벑벒벓범법벖벗벘벙벚벛벜벝벞벟베벡벢벣벤벥벦벧벨벩벪벫벬벭벮벯벰벱벲벳벴벵벶벷벸벹벺벻벼벽벾벿변볁볂볃별볅볆볇볈볉볊볋볌볍볎볏볐병볒\" # noqa: E501\n + \"볓볔볕볖볗볘볙볚볛볜볝볞볟볠볡볢볣볤볥볦볧볨볩볪볫볬볭볮볯볰볱볲볳보복볶볷본볹볺볻볼볽볾볿봀봁봂봃봄봅봆봇봈봉봊봋봌봍봎봏봐봑봒봓봔봕봖봗봘봙봚봛\" # noqa: E501\n + \"봜봝봞봟봠봡봢봣봤봥봦봧봨봩봪봫봬봭봮봯봰봱봲봳봴봵봶봷봸봹봺봻봼봽봾봿뵀뵁뵂뵃뵄뵅뵆뵇뵈뵉뵊뵋뵌뵍뵎뵏뵐뵑뵒뵓뵔뵕뵖뵗뵘뵙뵚뵛뵜뵝뵞뵟뵠뵡뵢뵣뵤\" # noqa: E501\n + \"뵥뵦뵧뵨뵩뵪뵫뵬뵭뵮뵯뵰뵱뵲뵳뵴뵵뵶뵷뵸뵹뵺뵻뵼뵽뵾뵿부북붂붃분붅붆붇불붉붊붋붌붍붎붏붐붑붒붓붔붕붖붗붘붙붚붛붜붝붞붟붠붡붢붣붤붥붦붧붨붩붪붫붬붭\" # noqa: E501\n + \"붮붯붰붱붲붳붴붵붶붷붸붹붺붻붼붽붾붿뷀뷁뷂뷃뷄뷅뷆뷇뷈뷉뷊뷋뷌뷍뷎뷏뷐뷑뷒뷓뷔뷕뷖뷗뷘뷙뷚뷛뷜뷝뷞뷟뷠뷡뷢뷣뷤뷥뷦뷧뷨뷩뷪뷫뷬뷭뷮뷯뷰뷱뷲뷳뷴뷵뷶\" # noqa: E501\n + \"뷷뷸뷹뷺뷻뷼뷽뷾뷿븀븁븂븃븄븅븆븇븈븉븊븋브븍븎븏븐븑븒븓블븕븖븗븘븙븚븛븜븝븞븟븠븡븢븣븤븥븦븧븨븩븪븫븬븭븮븯븰븱븲븳븴븵븶븷븸븹븺븻븼븽븾븿\" # noqa: E501\n + \"빀빁빂빃비빅빆빇빈빉빊빋빌빍빎빏빐빑빒빓빔빕빖빗빘빙빚빛빜빝빞빟빠빡빢빣빤빥빦빧빨빩빪빫빬빭빮빯빰빱빲빳빴빵빶빷빸빹빺빻빼빽빾빿뺀뺁뺂뺃뺄뺅뺆뺇뺈\" # noqa: E501\n + \"뺉뺊뺋뺌뺍뺎뺏뺐뺑뺒뺓뺔뺕뺖뺗뺘뺙뺚뺛뺜뺝뺞뺟뺠뺡뺢뺣뺤뺥뺦뺧뺨뺩뺪뺫뺬뺭뺮뺯뺰뺱뺲뺳뺴뺵뺶뺷뺸뺹뺺뺻뺼뺽뺾뺿뻀뻁뻂뻃뻄뻅뻆뻇뻈뻉뻊뻋뻌뻍뻎뻏뻐뻑\" # noqa: E501\n + \"뻒뻓뻔뻕뻖뻗뻘뻙뻚뻛뻜뻝뻞뻟뻠뻡뻢뻣뻤뻥뻦뻧뻨뻩뻪뻫뻬뻭뻮뻯뻰뻱뻲뻳뻴뻵뻶뻷뻸뻹뻺뻻뻼뻽뻾뻿뼀뼁뼂뼃뼄뼅뼆뼇뼈뼉뼊뼋뼌뼍뼎뼏뼐뼑뼒뼓뼔뼕뼖뼗뼘뼙뼚\" # noqa: E501\n + \"뼛뼜뼝뼞뼟뼠뼡뼢뼣뼤뼥뼦뼧뼨뼩뼪뼫뼬뼭뼮뼯뼰뼱뼲뼳뼴뼵뼶뼷뼸뼹뼺뼻뼼뼽뼾뼿뽀뽁뽂뽃뽄뽅뽆뽇뽈뽉뽊뽋뽌뽍뽎뽏뽐뽑뽒뽓뽔뽕뽖뽗뽘뽙뽚뽛뽜뽝뽞뽟뽠뽡뽢뽣\" # noqa: E501\n + \"뽤뽥뽦뽧뽨뽩뽪뽫뽬뽭뽮뽯뽰뽱뽲뽳뽴뽵뽶뽷뽸뽹뽺뽻뽼뽽뽾뽿뾀뾁뾂뾃뾄뾅뾆뾇뾈뾉뾊뾋뾌뾍뾎뾏뾐뾑뾒뾓뾔뾕뾖뾗뾘뾙뾚뾛뾜뾝뾞뾟뾠뾡뾢뾣뾤뾥뾦뾧뾨뾩뾪뾫뾬\" # noqa: E501\n + \"뾭뾮뾯뾰뾱뾲뾳뾴뾵뾶뾷뾸뾹뾺뾻뾼뾽뾾뾿뿀뿁뿂뿃뿄뿅뿆뿇뿈뿉뿊뿋뿌뿍뿎뿏뿐뿑뿒뿓뿔뿕뿖뿗뿘뿙뿚뿛뿜뿝뿞뿟뿠뿡뿢뿣뿤뿥뿦뿧뿨뿩뿪뿫뿬뿭뿮뿯뿰뿱뿲뿳뿴뿵\" # noqa: E501\n + \"뿶뿷뿸뿹뿺뿻뿼뿽뿾뿿쀀쀁쀂쀃쀄쀅쀆쀇쀈쀉쀊쀋쀌쀍쀎쀏쀐쀑쀒쀓쀔쀕쀖쀗쀘쀙쀚쀛쀜쀝쀞쀟쀠쀡쀢쀣쀤쀥쀦쀧쀨쀩쀪쀫쀬쀭쀮쀯쀰쀱쀲쀳쀴쀵쀶쀷쀸쀹쀺쀻쀼쀽쀾\" # noqa: E501\n + \"쀿쁀쁁쁂쁃쁄쁅쁆쁇쁈쁉쁊쁋쁌쁍쁎쁏쁐쁑쁒쁓쁔쁕쁖쁗쁘쁙쁚쁛쁜쁝쁞쁟쁠쁡쁢쁣쁤쁥쁦쁧쁨쁩쁪쁫쁬쁭쁮쁯쁰쁱쁲쁳쁴쁵쁶쁷쁸쁹쁺쁻쁼쁽쁾쁿삀삁삂삃삄삅삆삇\" # noqa: E501\n + \"삈삉삊삋삌삍삎삏삐삑삒삓삔삕삖삗삘삙삚삛삜삝삞삟삠삡삢삣삤삥삦삧삨삩삪삫사삭삮삯산삱삲삳살삵삶삷삸삹삺삻삼삽삾삿샀상샂샃샄샅샆샇새색샊샋샌샍샎샏샐\" # noqa: E501\n + \"샑샒샓샔샕샖샗샘샙샚샛샜생샞샟샠샡샢샣샤샥샦샧샨샩샪샫샬샭샮샯샰샱샲샳샴샵샶샷샸샹샺샻샼샽샾샿섀섁섂섃섄섅섆섇섈섉섊섋섌섍섎섏섐섑섒섓섔섕섖섗섘섙\" # noqa: E501\n + \"섚섛서석섞섟선섡섢섣설섥섦섧섨섩섪섫섬섭섮섯섰성섲섳섴섵섶섷세섹섺섻센섽섾섿셀셁셂셃셄셅셆셇셈셉셊셋셌셍셎셏셐셑셒셓셔셕셖셗션셙셚셛셜셝셞셟셠셡셢\" # noqa: E501\n + \"셣셤셥셦셧셨셩셪셫셬셭셮셯셰셱셲셳셴셵셶셷셸셹셺셻셼셽셾셿솀솁솂솃솄솅솆솇솈솉솊솋소속솎솏손솑솒솓솔솕솖솗솘솙솚솛솜솝솞솟솠송솢솣솤솥솦솧솨솩솪솫\" # noqa: E501\n + \"솬솭솮솯솰솱솲솳솴솵솶솷솸솹솺솻솼솽솾솿쇀쇁쇂쇃쇄쇅쇆쇇쇈쇉쇊쇋쇌쇍쇎쇏쇐쇑쇒쇓쇔쇕쇖쇗쇘쇙쇚쇛쇜쇝쇞쇟쇠쇡쇢쇣쇤쇥쇦쇧쇨쇩쇪쇫쇬쇭쇮쇯쇰쇱쇲쇳쇴\" # noqa: E501\n + \"쇵쇶쇷쇸쇹쇺쇻쇼쇽쇾쇿숀숁숂숃숄숅숆숇숈숉숊숋숌숍숎숏숐숑숒숓숔숕숖숗수숙숚숛순숝숞숟술숡숢숣숤숥숦숧숨숩숪숫숬숭숮숯숰숱숲숳숴숵숶숷숸숹숺숻숼숽\" # noqa: E501\n + \"숾숿쉀쉁쉂쉃쉄쉅쉆쉇쉈쉉쉊쉋쉌쉍쉎쉏쉐쉑쉒쉓쉔쉕쉖쉗쉘쉙쉚쉛쉜쉝쉞쉟쉠쉡쉢쉣쉤쉥쉦쉧쉨쉩쉪쉫쉬쉭쉮쉯쉰쉱쉲쉳쉴쉵쉶쉷쉸쉹쉺쉻쉼쉽쉾쉿슀슁슂슃슄슅슆\" # noqa: E501\n + \"슇슈슉슊슋슌슍슎슏슐슑슒슓슔슕슖슗슘슙슚슛슜슝슞슟슠슡슢슣스슥슦슧슨슩슪슫슬슭슮슯슰슱슲슳슴습슶슷슸승슺슻슼슽슾슿싀싁싂싃싄싅싆싇싈싉싊싋싌싍싎싏\" # noqa: E501\n + \"싐싑싒싓싔싕싖싗싘싙싚싛시식싞싟신싡싢싣실싥싦싧싨싩싪싫심십싮싯싰싱싲싳싴싵싶싷싸싹싺싻싼싽싾싿쌀쌁쌂쌃쌄쌅쌆쌇쌈쌉쌊쌋쌌쌍쌎쌏쌐쌑쌒쌓쌔쌕쌖쌗쌘\" # noqa: E501\n + \"쌙쌚쌛쌜쌝쌞쌟쌠쌡쌢쌣쌤쌥쌦쌧쌨쌩쌪쌫쌬쌭쌮쌯쌰쌱쌲쌳쌴쌵쌶쌷쌸쌹쌺쌻쌼쌽쌾쌿썀썁썂썃썄썅썆썇썈썉썊썋썌썍썎썏썐썑썒썓썔썕썖썗썘썙썚썛썜썝썞썟썠썡\" # noqa: E501\n + \"썢썣썤썥썦썧써썩썪썫썬썭썮썯썰썱썲썳썴썵썶썷썸썹썺썻썼썽썾썿쎀쎁쎂쎃쎄쎅쎆쎇쎈쎉쎊쎋쎌쎍쎎쎏쎐쎑쎒쎓쎔쎕쎖쎗쎘쎙쎚쎛쎜쎝쎞쎟쎠쎡쎢쎣쎤쎥쎦쎧쎨쎩쎪\" # noqa: E501\n + \"쎫쎬쎭쎮쎯쎰쎱쎲쎳쎴쎵쎶쎷쎸쎹쎺쎻쎼쎽쎾쎿쏀쏁쏂쏃쏄쏅쏆쏇쏈쏉쏊쏋쏌쏍쏎쏏쏐쏑쏒쏓쏔쏕쏖쏗쏘쏙쏚쏛쏜쏝쏞쏟쏠쏡쏢쏣쏤쏥쏦쏧쏨쏩쏪쏫쏬쏭쏮쏯쏰쏱쏲쏳\" # noqa: E501\n + \"쏴쏵쏶쏷쏸쏹쏺쏻쏼쏽쏾쏿쐀쐁쐂쐃쐄쐅쐆쐇쐈쐉쐊쐋쐌쐍쐎쐏쐐쐑쐒쐓쐔쐕쐖쐗쐘쐙쐚쐛쐜쐝쐞쐟쐠쐡쐢쐣쐤쐥쐦쐧쐨쐩쐪쐫쐬쐭쐮쐯쐰쐱쐲쐳쐴쐵쐶쐷쐸쐹쐺쐻쐼\" # noqa: E501\n + \"쐽쐾쐿쑀쑁쑂쑃쑄쑅쑆쑇쑈쑉쑊쑋쑌쑍쑎쑏쑐쑑쑒쑓쑔쑕쑖쑗쑘쑙쑚쑛쑜쑝쑞쑟쑠쑡쑢쑣쑤쑥쑦쑧쑨쑩쑪쑫쑬쑭쑮쑯쑰쑱쑲쑳쑴쑵쑶쑷쑸쑹쑺쑻쑼쑽쑾쑿쒀쒁쒂쒃쒄쒅\" # noqa: E501\n + \"쒆쒇쒈쒉쒊쒋쒌쒍쒎쒏쒐쒑쒒쒓쒔쒕쒖쒗쒘쒙쒚쒛쒜쒝쒞쒟쒠쒡쒢쒣쒤쒥쒦쒧쒨쒩쒪쒫쒬쒭쒮쒯쒰쒱쒲쒳쒴쒵쒶쒷쒸쒹쒺쒻쒼쒽쒾쒿쓀쓁쓂쓃쓄쓅쓆쓇쓈쓉쓊쓋쓌쓍쓎\" # noqa: E501\n + \"쓏쓐쓑쓒쓓쓔쓕쓖쓗쓘쓙쓚쓛쓜쓝쓞쓟쓠쓡쓢쓣쓤쓥쓦쓧쓨쓩쓪쓫쓬쓭쓮쓯쓰쓱쓲쓳쓴쓵쓶쓷쓸쓹쓺쓻쓼쓽쓾쓿씀씁씂씃씄씅씆씇씈씉씊씋씌씍씎씏씐씑씒씓씔씕씖씗\" # noqa: E501\n + \"씘씙씚씛씜씝씞씟씠씡씢씣씤씥씦씧씨씩씪씫씬씭씮씯씰씱씲씳씴씵씶씷씸씹씺씻씼씽씾씿앀앁앂앃아악앆앇안앉않앋알앍앎앏앐앑앒앓암압앖앗았앙앚앛앜앝앞앟애\" # noqa: E501\n + \"액앢앣앤앥앦앧앨앩앪앫앬앭앮앯앰앱앲앳앴앵앶앷앸앹앺앻야약앾앿얀얁얂얃얄얅얆얇얈얉얊얋얌얍얎얏얐양얒얓얔얕얖얗얘얙얚얛얜얝얞얟얠얡얢얣얤얥얦얧얨얩\" # noqa: E501\n + \"얪얫얬얭얮얯얰얱얲얳어억얶얷언얹얺얻얼얽얾얿엀엁엂엃엄업없엇었엉엊엋엌엍엎엏에엑엒엓엔엕엖엗엘엙엚엛엜엝엞엟엠엡엢엣엤엥엦엧엨엩엪엫여역엮엯연엱엲\" # noqa: E501\n + \"엳열엵엶엷엸엹엺엻염엽엾엿였영옂옃옄옅옆옇예옉옊옋옌옍옎옏옐옑옒옓옔옕옖옗옘옙옚옛옜옝옞옟옠옡옢옣오옥옦옧온옩옪옫올옭옮옯옰옱옲옳옴옵옶옷옸옹옺옻\" # noqa: E501\n + \"옼옽옾옿와왁왂왃완왅왆왇왈왉왊왋왌왍왎왏왐왑왒왓왔왕왖왗왘왙왚왛왜왝왞왟왠왡왢왣왤왥왦왧왨왩왪왫왬왭왮왯왰왱왲왳왴왵왶왷외왹왺왻왼왽왾왿욀욁욂욃욄\" # noqa: E501\n + \"욅욆욇욈욉욊욋욌욍욎욏욐욑욒욓요욕욖욗욘욙욚욛욜욝욞욟욠욡욢욣욤욥욦욧욨용욪욫욬욭욮욯우욱욲욳운욵욶욷울욹욺욻욼욽욾욿움웁웂웃웄웅웆웇웈웉웊웋워웍\" # noqa: E501\n + \"웎웏원웑웒웓월웕웖웗웘웙웚웛웜웝웞웟웠웡웢웣웤웥웦웧웨웩웪웫웬웭웮웯웰웱웲웳웴웵웶웷웸웹웺웻웼웽웾웿윀윁윂윃위윅윆윇윈윉윊윋윌윍윎윏윐윑윒윓윔윕윖\" # noqa: E501\n + \"윗윘윙윚윛윜윝윞윟유육윢윣윤윥윦윧율윩윪윫윬윭윮윯윰윱윲윳윴융윶윷윸윹윺윻으윽윾윿은읁읂읃을읅읆읇읈읉읊읋음읍읎읏읐응읒읓읔읕읖읗의읙읚읛읜읝읞읟\" # noqa: E501\n + \"읠읡읢읣읤읥읦읧읨읩읪읫읬읭읮읯읰읱읲읳이익읶읷인읹읺읻일읽읾읿잀잁잂잃임입잆잇있잉잊잋잌잍잎잏자작잒잓잔잕잖잗잘잙잚잛잜잝잞잟잠잡잢잣잤장잦잧잨\" # noqa: E501\n + \"잩잪잫재잭잮잯잰잱잲잳잴잵잶잷잸잹잺잻잼잽잾잿쟀쟁쟂쟃쟄쟅쟆쟇쟈쟉쟊쟋쟌쟍쟎쟏쟐쟑쟒쟓쟔쟕쟖쟗쟘쟙쟚쟛쟜쟝쟞쟟쟠쟡쟢쟣쟤쟥쟦쟧쟨쟩쟪쟫쟬쟭쟮쟯쟰쟱\" # noqa: E501\n + \"쟲쟳쟴쟵쟶쟷쟸쟹쟺쟻쟼쟽쟾쟿저적젂젃전젅젆젇절젉젊젋젌젍젎젏점접젒젓젔정젖젗젘젙젚젛제젝젞젟젠젡젢젣젤젥젦젧젨젩젪젫젬젭젮젯젰젱젲젳젴젵젶젷져젹젺\" # noqa: E501\n + \"젻젼젽젾젿졀졁졂졃졄졅졆졇졈졉졊졋졌졍졎졏졐졑졒졓졔졕졖졗졘졙졚졛졜졝졞졟졠졡졢졣졤졥졦졧졨졩졪졫졬졭졮졯조족졲졳존졵졶졷졸졹졺졻졼졽졾졿좀좁좂좃\" # noqa: E501\n + \"좄종좆좇좈좉좊좋좌좍좎좏좐좑좒좓좔좕좖좗좘좙좚좛좜좝좞좟좠좡좢좣좤좥좦좧좨좩좪좫좬좭좮좯좰좱좲좳좴좵좶좷좸좹좺좻좼좽좾좿죀죁죂죃죄죅죆죇죈죉죊죋죌\" # noqa: E501\n + \"죍죎죏죐죑죒죓죔죕죖죗죘죙죚죛죜죝죞죟죠죡죢죣죤죥죦죧죨죩죪죫죬죭죮죯죰죱죲죳죴죵죶죷죸죹죺죻주죽죾죿준줁줂줃줄줅줆줇줈줉줊줋줌줍줎줏줐중줒줓줔줕\" # noqa: E501\n + \"줖줗줘줙줚줛줜줝줞줟줠줡줢줣줤줥줦줧줨줩줪줫줬줭줮줯줰줱줲줳줴줵줶줷줸줹줺줻줼줽줾줿쥀쥁쥂쥃쥄쥅쥆쥇쥈쥉쥊쥋쥌쥍쥎쥏쥐쥑쥒쥓쥔쥕쥖쥗쥘쥙쥚쥛쥜쥝쥞\" # noqa: E501\n + \"쥟쥠쥡쥢쥣쥤쥥쥦쥧쥨쥩쥪쥫쥬쥭쥮쥯쥰쥱쥲쥳쥴쥵쥶쥷쥸쥹쥺쥻쥼쥽쥾쥿즀즁즂즃즄즅즆즇즈즉즊즋즌즍즎즏즐즑즒즓즔즕즖즗즘즙즚즛즜증즞즟즠즡즢즣즤즥즦즧\" # noqa: E501\n + \"즨즩즪즫즬즭즮즯즰즱즲즳즴즵즶즷즸즹즺즻즼즽즾즿지직짂짃진짅짆짇질짉짊짋짌짍짎짏짐집짒짓짔징짖짗짘짙짚짛짜짝짞짟짠짡짢짣짤짥짦짧짨짩짪짫짬짭짮짯짰\" # noqa: E501\n + \"짱짲짳짴짵짶짷째짹짺짻짼짽짾짿쨀쨁쨂쨃쨄쨅쨆쨇쨈쨉쨊쨋쨌쨍쨎쨏쨐쨑쨒쨓쨔쨕쨖쨗쨘쨙쨚쨛쨜쨝쨞쨟쨠쨡쨢쨣쨤쨥쨦쨧쨨쨩쨪쨫쨬쨭쨮쨯쨰쨱쨲쨳쨴쨵쨶쨷쨸쨹\" # noqa: E501\n + \"쨺쨻쨼쨽쨾쨿쩀쩁쩂쩃쩄쩅쩆쩇쩈쩉쩊쩋쩌쩍쩎쩏쩐쩑쩒쩓쩔쩕쩖쩗쩘쩙쩚쩛쩜쩝쩞쩟쩠쩡쩢쩣쩤쩥쩦쩧쩨쩩쩪쩫쩬쩭쩮쩯쩰쩱쩲쩳쩴쩵쩶쩷쩸쩹쩺쩻쩼쩽쩾쩿쪀쪁쪂\" # noqa: E501\n + \"쪃쪄쪅쪆쪇쪈쪉쪊쪋쪌쪍쪎쪏쪐쪑쪒쪓쪔쪕쪖쪗쪘쪙쪚쪛쪜쪝쪞쪟쪠쪡쪢쪣쪤쪥쪦쪧쪨쪩쪪쪫쪬쪭쪮쪯쪰쪱쪲쪳쪴쪵쪶쪷쪸쪹쪺쪻쪼쪽쪾쪿쫀쫁쫂쫃쫄쫅쫆쫇쫈쫉쫊쫋\" # noqa: E501\n + \"쫌쫍쫎쫏쫐쫑쫒쫓쫔쫕쫖쫗쫘쫙쫚쫛쫜쫝쫞쫟쫠쫡쫢쫣쫤쫥쫦쫧쫨쫩쫪쫫쫬쫭쫮쫯쫰쫱쫲쫳쫴쫵쫶쫷쫸쫹쫺쫻쫼쫽쫾쫿쬀쬁쬂쬃쬄쬅쬆쬇쬈쬉쬊쬋쬌쬍쬎쬏쬐쬑쬒쬓쬔\" # noqa: E501\n + \"쬕쬖쬗쬘쬙쬚쬛쬜쬝쬞쬟쬠쬡쬢쬣쬤쬥쬦쬧쬨쬩쬪쬫쬬쬭쬮쬯쬰쬱쬲쬳쬴쬵쬶쬷쬸쬹쬺쬻쬼쬽쬾쬿쭀쭁쭂쭃쭄쭅쭆쭇쭈쭉쭊쭋쭌쭍쭎쭏쭐쭑쭒쭓쭔쭕쭖쭗쭘쭙쭚쭛쭜쭝\" # noqa: E501\n + \"쭞쭟쭠쭡쭢쭣쭤쭥쭦쭧쭨쭩쭪쭫쭬쭭쭮쭯쭰쭱쭲쭳쭴쭵쭶쭷쭸쭹쭺쭻쭼쭽쭾쭿쮀쮁쮂쮃쮄쮅쮆쮇쮈쮉쮊쮋쮌쮍쮎쮏쮐쮑쮒쮓쮔쮕쮖쮗쮘쮙쮚쮛쮜쮝쮞쮟쮠쮡쮢쮣쮤쮥쮦\" # noqa: E501\n + \"쮧쮨쮩쮪쮫쮬쮭쮮쮯쮰쮱쮲쮳쮴쮵쮶쮷쮸쮹쮺쮻쮼쮽쮾쮿쯀쯁쯂쯃쯄쯅쯆쯇쯈쯉쯊쯋쯌쯍쯎쯏쯐쯑쯒쯓쯔쯕쯖쯗쯘쯙쯚쯛쯜쯝쯞쯟쯠쯡쯢쯣쯤쯥쯦쯧쯨쯩쯪쯫쯬쯭쯮쯯\" # noqa: E501\n + \"쯰쯱쯲쯳쯴쯵쯶쯷쯸쯹쯺쯻쯼쯽쯾쯿찀찁찂찃찄찅찆찇찈찉찊찋찌찍찎찏찐찑찒찓찔찕찖찗찘찙찚찛찜찝찞찟찠찡찢찣찤찥찦찧차착찪찫찬찭찮찯찰찱찲찳찴찵찶찷참\" # noqa: E501\n + \"찹찺찻찼창찾찿챀챁챂챃채책챆챇챈챉챊챋챌챍챎챏챐챑챒챓챔챕챖챗챘챙챚챛챜챝챞챟챠챡챢챣챤챥챦챧챨챩챪챫챬챭챮챯챰챱챲챳챴챵챶챷챸챹챺챻챼챽챾챿첀첁\" # noqa: E501\n + \"첂첃첄첅첆첇첈첉첊첋첌첍첎첏첐첑첒첓첔첕첖첗처척첚첛천첝첞첟철첡첢첣첤첥첦첧첨첩첪첫첬청첮첯첰첱첲첳체첵첶첷첸첹첺첻첼첽첾첿쳀쳁쳂쳃쳄쳅쳆쳇쳈쳉쳊\" # noqa: E501\n + \"쳋쳌쳍쳎쳏쳐쳑쳒쳓쳔쳕쳖쳗쳘쳙쳚쳛쳜쳝쳞쳟쳠쳡쳢쳣쳤쳥쳦쳧쳨쳩쳪쳫쳬쳭쳮쳯쳰쳱쳲쳳쳴쳵쳶쳷쳸쳹쳺쳻쳼쳽쳾쳿촀촁촂촃촄촅촆촇초촉촊촋촌촍촎촏촐촑촒촓\" # noqa: E501\n + \"촔촕촖촗촘촙촚촛촜총촞촟촠촡촢촣촤촥촦촧촨촩촪촫촬촭촮촯촰촱촲촳촴촵촶촷촸촹촺촻촼촽촾촿쵀쵁쵂쵃쵄쵅쵆쵇쵈쵉쵊쵋쵌쵍쵎쵏쵐쵑쵒쵓쵔쵕쵖쵗쵘쵙쵚쵛최\" # noqa: E501\n + \"쵝쵞쵟쵠쵡쵢쵣쵤쵥쵦쵧쵨쵩쵪쵫쵬쵭쵮쵯쵰쵱쵲쵳쵴쵵쵶쵷쵸쵹쵺쵻쵼쵽쵾쵿춀춁춂춃춄춅춆춇춈춉춊춋춌춍춎춏춐춑춒춓추축춖춗춘춙춚춛출춝춞춟춠춡춢춣춤춥\" # noqa: E501\n + \"춦춧춨충춪춫춬춭춮춯춰춱춲춳춴춵춶춷춸춹춺춻춼춽춾춿췀췁췂췃췄췅췆췇췈췉췊췋췌췍췎췏췐췑췒췓췔췕췖췗췘췙췚췛췜췝췞췟췠췡췢췣췤췥췦췧취췩췪췫췬췭췮\" # noqa: E501\n + \"췯췰췱췲췳췴췵췶췷췸췹췺췻췼췽췾췿츀츁츂츃츄츅츆츇츈츉츊츋츌츍츎츏츐츑츒츓츔츕츖츗츘츙츚츛츜츝츞츟츠측츢츣츤츥츦츧츨츩츪츫츬츭츮츯츰츱츲츳츴층츶츷\" # noqa: E501\n + \"츸츹츺츻츼츽츾츿칀칁칂칃칄칅칆칇칈칉칊칋칌칍칎칏칐칑칒칓칔칕칖칗치칙칚칛친칝칞칟칠칡칢칣칤칥칦칧침칩칪칫칬칭칮칯칰칱칲칳카칵칶칷칸칹칺칻칼칽칾칿캀\" # noqa: E501\n + \"캁캂캃캄캅캆캇캈캉캊캋캌캍캎캏캐캑캒캓캔캕캖캗캘캙캚캛캜캝캞캟캠캡캢캣캤캥캦캧캨캩캪캫캬캭캮캯캰캱캲캳캴캵캶캷캸캹캺캻캼캽캾캿컀컁컂컃컄컅컆컇컈컉\" # noqa: E501\n + \"컊컋컌컍컎컏컐컑컒컓컔컕컖컗컘컙컚컛컜컝컞컟컠컡컢컣커컥컦컧컨컩컪컫컬컭컮컯컰컱컲컳컴컵컶컷컸컹컺컻컼컽컾컿케켁켂켃켄켅켆켇켈켉켊켋켌켍켎켏켐켑켒\" # noqa: E501\n + \"켓켔켕켖켗켘켙켚켛켜켝켞켟켠켡켢켣켤켥켦켧켨켩켪켫켬켭켮켯켰켱켲켳켴켵켶켷켸켹켺켻켼켽켾켿콀콁콂콃콄콅콆콇콈콉콊콋콌콍콎콏콐콑콒콓코콕콖콗콘콙콚콛\" # noqa: E501\n + \"콜콝콞콟콠콡콢콣콤콥콦콧콨콩콪콫콬콭콮콯콰콱콲콳콴콵콶콷콸콹콺콻콼콽콾콿쾀쾁쾂쾃쾄쾅쾆쾇쾈쾉쾊쾋쾌쾍쾎쾏쾐쾑쾒쾓쾔쾕쾖쾗쾘쾙쾚쾛쾜쾝쾞쾟쾠쾡쾢쾣쾤\" # noqa: E501\n + \"쾥쾦쾧쾨쾩쾪쾫쾬쾭쾮쾯쾰쾱쾲쾳쾴쾵쾶쾷쾸쾹쾺쾻쾼쾽쾾쾿쿀쿁쿂쿃쿄쿅쿆쿇쿈쿉쿊쿋쿌쿍쿎쿏쿐쿑쿒쿓쿔쿕쿖쿗쿘쿙쿚쿛쿜쿝쿞쿟쿠쿡쿢쿣쿤쿥쿦쿧쿨쿩쿪쿫쿬쿭\" # noqa: E501\n + \"쿮쿯쿰쿱쿲쿳쿴쿵쿶쿷쿸쿹쿺쿻쿼쿽쿾쿿퀀퀁퀂퀃퀄퀅퀆퀇퀈퀉퀊퀋퀌퀍퀎퀏퀐퀑퀒퀓퀔퀕퀖퀗퀘퀙퀚퀛퀜퀝퀞퀟퀠퀡퀢퀣퀤퀥퀦퀧퀨퀩퀪퀫퀬퀭퀮퀯퀰퀱퀲퀳퀴퀵퀶\" # noqa: E501\n + \"퀷퀸퀹퀺퀻퀼퀽퀾퀿큀큁큂큃큄큅큆큇큈큉큊큋큌큍큎큏큐큑큒큓큔큕큖큗큘큙큚큛큜큝큞큟큠큡큢큣큤큥큦큧큨큩큪큫크큭큮큯큰큱큲큳클큵큶큷큸큹큺큻큼큽큾큿\" # noqa: E501\n + \"킀킁킂킃킄킅킆킇킈킉킊킋킌킍킎킏킐킑킒킓킔킕킖킗킘킙킚킛킜킝킞킟킠킡킢킣키킥킦킧킨킩킪킫킬킭킮킯킰킱킲킳킴킵킶킷킸킹킺킻킼킽킾킿타탁탂탃탄탅탆탇탈\" # noqa: E501\n + \"탉탊탋탌탍탎탏탐탑탒탓탔탕탖탗탘탙탚탛태택탞탟탠탡탢탣탤탥탦탧탨탩탪탫탬탭탮탯탰탱탲탳탴탵탶탷탸탹탺탻탼탽탾탿턀턁턂턃턄턅턆턇턈턉턊턋턌턍턎턏턐턑\" # noqa: E501\n + \"턒턓턔턕턖턗턘턙턚턛턜턝턞턟턠턡턢턣턤턥턦턧턨턩턪턫턬턭턮턯터턱턲턳턴턵턶턷털턹턺턻턼턽턾턿텀텁텂텃텄텅텆텇텈텉텊텋테텍텎텏텐텑텒텓텔텕텖텗텘텙텚\" # noqa: E501\n + \"텛템텝텞텟텠텡텢텣텤텥텦텧텨텩텪텫텬텭텮텯텰텱텲텳텴텵텶텷텸텹텺텻텼텽텾텿톀톁톂톃톄톅톆톇톈톉톊톋톌톍톎톏톐톑톒톓톔톕톖톗톘톙톚톛톜톝톞톟토톡톢톣\" # noqa: E501\n + \"톤톥톦톧톨톩톪톫톬톭톮톯톰톱톲톳톴통톶톷톸톹톺톻톼톽톾톿퇀퇁퇂퇃퇄퇅퇆퇇퇈퇉퇊퇋퇌퇍퇎퇏퇐퇑퇒퇓퇔퇕퇖퇗퇘퇙퇚퇛퇜퇝퇞퇟퇠퇡퇢퇣퇤퇥퇦퇧퇨퇩퇪퇫퇬\" # noqa: E501\n + \"퇭퇮퇯퇰퇱퇲퇳퇴퇵퇶퇷퇸퇹퇺퇻퇼퇽퇾퇿툀툁툂툃툄툅툆툇툈툉툊툋툌툍툎툏툐툑툒툓툔툕툖툗툘툙툚툛툜툝툞툟툠툡툢툣툤툥툦툧툨툩툪툫투툭툮툯툰툱툲툳툴툵\" # noqa: E501\n + \"툶툷툸툹툺툻툼툽툾툿퉀퉁퉂퉃퉄퉅퉆퉇퉈퉉퉊퉋퉌퉍퉎퉏퉐퉑퉒퉓퉔퉕퉖퉗퉘퉙퉚퉛퉜퉝퉞퉟퉠퉡퉢퉣퉤퉥퉦퉧퉨퉩퉪퉫퉬퉭퉮퉯퉰퉱퉲퉳퉴퉵퉶퉷퉸퉹퉺퉻퉼퉽퉾\" # noqa: E501\n + \"퉿튀튁튂튃튄튅튆튇튈튉튊튋튌튍튎튏튐튑튒튓튔튕튖튗튘튙튚튛튜튝튞튟튠튡튢튣튤튥튦튧튨튩튪튫튬튭튮튯튰튱튲튳튴튵튶튷트특튺튻튼튽튾튿틀틁틂틃틄틅틆틇\" # noqa: E501\n + \"틈틉틊틋틌틍틎틏틐틑틒틓틔틕틖틗틘틙틚틛틜틝틞틟틠틡틢틣틤틥틦틧틨틩틪틫틬틭틮틯티틱틲틳틴틵틶틷틸틹틺틻틼틽틾틿팀팁팂팃팄팅팆팇팈팉팊팋파팍팎팏판\" # noqa: E501\n + \"팑팒팓팔팕팖팗팘팙팚팛팜팝팞팟팠팡팢팣팤팥팦팧패팩팪팫팬팭팮팯팰팱팲팳팴팵팶팷팸팹팺팻팼팽팾팿퍀퍁퍂퍃퍄퍅퍆퍇퍈퍉퍊퍋퍌퍍퍎퍏퍐퍑퍒퍓퍔퍕퍖퍗퍘퍙\" # noqa: E501\n + \"퍚퍛퍜퍝퍞퍟퍠퍡퍢퍣퍤퍥퍦퍧퍨퍩퍪퍫퍬퍭퍮퍯퍰퍱퍲퍳퍴퍵퍶퍷퍸퍹퍺퍻퍼퍽퍾퍿펀펁펂펃펄펅펆펇펈펉펊펋펌펍펎펏펐펑펒펓펔펕펖펗페펙펚펛펜펝펞펟펠펡펢\" # noqa: E501\n + \"펣펤펥펦펧펨펩펪펫펬펭펮펯펰펱펲펳펴펵펶펷편펹펺펻펼펽펾펿폀폁폂폃폄폅폆폇폈평폊폋폌폍폎폏폐폑폒폓폔폕폖폗폘폙폚폛폜폝폞폟폠폡폢폣폤폥폦폧폨폩폪폫\" # noqa: E501\n + \"포폭폮폯폰폱폲폳폴폵폶폷폸폹폺폻폼폽폾폿퐀퐁퐂퐃퐄퐅퐆퐇퐈퐉퐊퐋퐌퐍퐎퐏퐐퐑퐒퐓퐔퐕퐖퐗퐘퐙퐚퐛퐜퐝퐞퐟퐠퐡퐢퐣퐤퐥퐦퐧퐨퐩퐪퐫퐬퐭퐮퐯퐰퐱퐲퐳퐴\" # noqa: E501\n + \"퐵퐶퐷퐸퐹퐺퐻퐼퐽퐾퐿푀푁푂푃푄푅푆푇푈푉푊푋푌푍푎푏푐푑푒푓푔푕푖푗푘푙푚푛표푝푞푟푠푡푢푣푤푥푦푧푨푩푪푫푬푭푮푯푰푱푲푳푴푵푶푷푸푹푺푻푼푽\" # noqa: E501\n + \"푾푿풀풁풂풃풄풅풆풇품풉풊풋풌풍풎풏풐풑풒풓풔풕풖풗풘풙풚풛풜풝풞풟풠풡풢풣풤풥풦풧풨풩풪풫풬풭풮풯풰풱풲풳풴풵풶풷풸풹풺풻풼풽풾풿퓀퓁퓂퓃퓄퓅퓆\" # noqa: E501\n + \"퓇퓈퓉퓊퓋퓌퓍퓎퓏퓐퓑퓒퓓퓔퓕퓖퓗퓘퓙퓚퓛퓜퓝퓞퓟퓠퓡퓢퓣퓤퓥퓦퓧퓨퓩퓪퓫퓬퓭퓮퓯퓰퓱퓲퓳퓴퓵퓶퓷퓸퓹퓺퓻퓼퓽퓾퓿픀픁픂픃프픅픆픇픈픉픊픋플픍픎픏\" # noqa: E501\n + \"픐픑픒픓픔픕픖픗픘픙픚픛픜픝픞픟픠픡픢픣픤픥픦픧픨픩픪픫픬픭픮픯픰픱픲픳픴픵픶픷픸픹픺픻피픽픾픿핀핁핂핃필핅핆핇핈핉핊핋핌핍핎핏핐핑핒핓핔핕핖핗하\" # noqa: E501\n + \"학핚핛한핝핞핟할핡핢핣핤핥핦핧함합핪핫핬항핮핯핰핱핲핳해핵핶핷핸핹핺핻핼핽핾핿햀햁햂햃햄햅햆햇했행햊햋햌햍햎햏햐햑햒햓햔햕햖햗햘햙햚햛햜햝햞햟햠햡\" # noqa: E501\n + \"햢햣햤향햦햧햨햩햪햫햬햭햮햯햰햱햲햳햴햵햶햷햸햹햺햻햼햽햾햿헀헁헂헃헄헅헆헇허헉헊헋헌헍헎헏헐헑헒헓헔헕헖헗험헙헚헛헜헝헞헟헠헡헢헣헤헥헦헧헨헩헪\" # noqa: E501\n + \"헫헬헭헮헯헰헱헲헳헴헵헶헷헸헹헺헻헼헽헾헿혀혁혂혃현혅혆혇혈혉혊혋혌혍혎혏혐협혒혓혔형혖혗혘혙혚혛혜혝혞혟혠혡혢혣혤혥혦혧혨혩혪혫혬혭혮혯혰혱혲혳\" # noqa: E501\n + \"혴혵혶혷호혹혺혻혼혽혾혿홀홁홂홃홄홅홆홇홈홉홊홋홌홍홎홏홐홑홒홓화확홖홗환홙홚홛활홝홞홟홠홡홢홣홤홥홦홧홨황홪홫홬홭홮홯홰홱홲홳홴홵홶홷홸홹홺홻홼\" # noqa: E501\n + \"홽홾홿횀횁횂횃횄횅횆횇횈횉횊횋회획횎횏횐횑횒횓횔횕횖횗횘횙횚횛횜횝횞횟횠횡횢횣횤횥횦횧효횩횪횫횬횭횮횯횰횱횲횳횴횵횶횷횸횹횺횻횼횽횾횿훀훁훂훃후훅\" # noqa: E501\n + \"훆훇훈훉훊훋훌훍훎훏훐훑훒훓훔훕훖훗훘훙훚훛훜훝훞훟훠훡훢훣훤훥훦훧훨훩훪훫훬훭훮훯훰훱훲훳훴훵훶훷훸훹훺훻훼훽훾훿휀휁휂휃휄휅휆휇휈휉휊휋휌휍휎\" # noqa: E501\n + \"휏휐휑휒휓휔휕휖휗휘휙휚휛휜휝휞휟휠휡휢휣휤휥휦휧휨휩휪휫휬휭휮휯휰휱휲휳휴휵휶휷휸휹휺휻휼휽휾휿흀흁흂흃흄흅흆흇흈흉흊흋흌흍흎흏흐흑흒흓흔흕흖흗\" # noqa: E501\n + \"흘흙흚흛흜흝흞흟흠흡흢흣흤흥흦흧흨흩흪흫희흭흮흯흰흱흲흳흴흵흶흷흸흹흺흻흼흽흾흿힀힁힂힃힄힅힆힇히힉힊힋힌힍힎힏힐힑힒힓힔힕힖힗힘힙힚힛힜힝힞힟힠\" # noqa: E501\n + \"힡힢힣\"\n + _BASE_VOCABS[\"punctuation\"]\n + \"。・〜°—、「」『』【】゛》《〉〈\" # punctuation\n + _BASE_VOCABS[\"currency\"]\n + \"₩\"\n)\n\nVOCABS[\"simplified_chinese\"] = (\n _BASE_VOCABS[\"digits\"]\n + \"㐀㐁㐂㐃㐄㐅㐆㐇㐈㐉㐊㐋㐌㐍㐎㐏㐐㐑㐒㐓㐔㐕㐖㐗㐘㐙㐚㐛㐜㐝㐞㐟㐠㐡㐢㐣㐤㐥㐦㐧㐨㐩㐪㐫㐬㐭㐮㐯㐰㐱㐲㐳㐴㐵㐶㐷㐸㐹㐺㐻㐼㐽㐾㐿㑀㑁㑂\" # noqa: E501\n + \"㑄㑅㑆㑇㑈㑉㑊㑋㑌㑍㑎㑏㑐㑑㑒㑓㑔㑕㑖㑗㑘㑙㑚㑛㑜㑝㑞㑟㑠㑡㑢㑣㑤㑥㑦㑧㑨㑩㑪㑫㑬㑭㑮㑯㑰㑱㑲㑳㑴㑵㑶㑷㑸㑹㑺㑻㑼㑽㑾㑿㒀㒁㒂㒃㒄㒅㒆\" # noqa: E501\n + \"㒇㒈㒉㒊㒋㒌㒍㒎㒏㒐㒑㒒㒓㒔㒕㒖㒗㒘㒙㒚㒛㒜㒝㒞㒟㒠㒡㒢㒣㒤㒥㒦㒧㒨㒩㒪㒫㒬㒭㒮㒯㒰㒱㒲㒳㒴㒵㒶㒷㒸㒹㒺㒻㒼㒽㒾㒿㓀㓁㓂㓃㓄㓅㓆㓇㓈㓉\" # noqa: E501\n + \"㓊㓋㓌㓍㓎㓏㓐㓑㓒㓓㓔㓕㓖㓗㓘㓙㓚㓛㓜㓝㓞㓟㓠㓡㓢㓣㓤㓥㓦㓧㓨㓩㓪㓫㓬㓭㓮㓯㓰㓱㓲㓳㓴㓵㓶㓷㓸㓹㓺㓻㓼㓽㓾㓿㔀㔁㔂㔃㔄㔅㔆㔇㔈㔉㔊㔋㔌\" # noqa: E501\n + \"㔍㔎㔏㔐㔑㔒㔓㔔㔕㔖㔗㔘㔙㔚㔛㔜㔝㔞㔟㔠㔡㔢㔣㔤㔥㔦㔧㔨㔩㔪㔫㔬㔭㔮㔯㔰㔱㔲㔳㔴㔵㔶㔷㔸㔹㔺㔻㔼㔽㔾㔿㕀㕁㕂㕃㕄㕅㕆㕇㕈㕉㕊㕋㕌㕍㕎㕏\" # noqa: E501\n + \"㕐㕑㕒㕓㕔㕕㕖㕗㕘㕙㕚㕛㕜㕝㕞㕟㕠㕡㕢㕣㕤㕥㕦㕧㕨㕩㕪㕫㕬㕭㕮㕯㕰㕱㕲㕳㕴㕵㕶㕷㕸㕹㕺㕻㕼㕽㕾㕿㖀㖁㖂㖃㖄㖅㖆㖇㖈㖉㖊㖋㖌㖍㖎㖏㖐㖑㖒\" # noqa: E501\n + \"㖓㖔㖕㖖㖗㖘㖙㖚㖛㖜㖝㖞㖟㖠㖡㖢㖣㖤㖥㖦㖧㖨㖩㖪㖫㖬㖭㖮㖯㖰㖱㖲㖳㖴㖵㖶㖷㖸㖹㖺㖻㖼㖽㖾㖿㗀㗁㗂㗃㗄㗅㗆㗇㗈㗉㗊㗋㗌㗍㗎㗏㗐㗑㗒㗓㗔㗕\" # noqa: E501\n + \"㗖㗗㗘㗙㗚㗛㗜㗝㗞㗟㗠㗡㗢㗣㗤㗥㗦㗧㗨㗩㗪㗫㗬㗭㗮㗯㗰㗱㗲㗳㗴㗵㗶㗷㗸㗹㗺㗻㗼㗽㗾㗿㘀㘁㘂㘃㘄㘅㘆㘇㘈㘉㘊㘋㘌㘍㘎㘏㘐㘑㘒㘓㘔㘕㘖㘗㘘\" # noqa: E501\n + \"㘙㘚㘛㘜㘝㘞㘟㘠㘡㘢㘣㘤㘥㘦㘧㘨㘩㘪㘫㘬㘭㘮㘯㘰㘱㘲㘳㘴㘵㘶㘷㘸㘹㘺㘻㘼㘽㘾㘿㙀㙁㙂㙃㙄㙅㙆㙇㙈㙉㙊㙋㙌㙍㙎㙏㙐㙑㙒㙓㙔㙕㙖㙗㙘㙙㙚㙛\" # noqa: E501\n + \"㙜㙝㙞㙟㙠㙡㙢㙣㙤㙥㙦㙧㙨㙩㙪㙫㙬㙭㙮㙯㙰㙱㙲㙳㙴㙵㙶㙷㙸㙹㙺㙻㙼㙽㙾㙿㚀㚁㚂㚃㚄㚅㚆㚇㚈㚉㚊㚋㚌㚍㚎㚏㚐㚑㚒㚓㚔㚕㚖㚗㚘㚙㚚㚛㚜㚝㚞\" # noqa: E501\n + \"㚟㚠㚡㚢㚣㚤㚥㚦㚧㚨㚩㚪㚫㚬㚭㚮㚯㚰㚱㚲㚳㚴㚵㚶㚷㚸㚹㚺㚻㚼㚽㚾㚿㛀㛁㛂㛃㛄㛅㛆㛇㛈㛉㛊㛋㛌㛍㛎㛏㛐㛑㛒㛓㛔㛕㛖㛗㛘㛙㛚㛛㛜㛝㛞㛟㛠㛡\" # noqa: E501\n + \"㛢㛣㛤㛥㛦㛧㛨㛩㛪㛫㛬㛭㛮㛯㛰㛱㛲㛳㛴㛵㛶㛷㛸㛹㛺㛻㛼㛽㛾㛿㜀㜁㜂㜃㜄㜅㜆㜇㜈㜉㜊㜋㜌㜍㜎㜏㜐㜑㜒㜓㜔㜕㜖㜗㜘㜙㜚㜛㜜㜝㜞㜟㜠㜡㜢㜣㜤\" # noqa: E501\n + \"㜥㜦㜧㜨㜩㜪㜫㜬㜭㜮㜯㜰㜱㜲㜳㜴㜵㜶㜷㜸㜹㜺㜻㜼㜽㜾㜿㝀㝁㝂㝃㝄㝅㝆㝇㝈㝉㝊㝋㝌㝍㝎㝏㝐㝑㝒㝓㝔㝕㝖㝗㝘㝙㝚㝛㝜㝝㝞㝟㝠㝡㝢㝣㝤㝥㝦㝧\" # noqa: E501\n + \"㝨㝩㝪㝫㝬㝭㝮㝯㝰㝱㝲㝳㝴㝵㝶㝷㝸㝹㝺㝻㝼㝽㝾㝿㞀㞁㞂㞃㞄㞅㞆㞇㞈㞉㞊㞋㞌㞍㞎㞏㞐㞑㞒㞓㞔㞕㞖㞗㞘㞙㞚㞛㞜㞝㞞㞟㞠㞡㞢㞣㞤㞥㞦㞧㞨㞩㞪\" # noqa: E501\n + \"㞫㞬㞭㞮㞯㞰㞱㞲㞳㞴㞵㞶㞷㞸㞹㞺㞻㞼㞽㞾㞿㟀㟁㟂㟃㟄㟅㟆㟇㟈㟉㟊㟋㟌㟍㟎㟏㟐㟑㟒㟓㟔㟕㟖㟗㟘㟙㟚㟛㟜㟝㟞㟟㟠㟡㟢㟣㟤㟥㟦㟧㟨㟩㟪㟫㟬㟭\" # noqa: E501\n + \"㟮㟯㟰㟱㟲㟳㟴㟵㟶㟷㟸㟹㟺㟻㟼㟽㟾㟿㠀㠁㠂㠃㠄㠅㠆㠇㠈㠉㠊㠋㠌㠍㠎㠏㠐㠑㠒㠓㠔㠕㠖㠗㠘㠙㠚㠛㠜㠝㠞㠟㠠㠡㠢㠣㠤㠥㠦㠧㠨㠩㠪㠫㠬㠭㠮㠯㠰\" # noqa: E501\n + \"㠱㠲㠳㠴㠵㠶㠷㠸㠹㠺㠻㠼㠽㠾㠿㡀㡁㡂㡃㡄㡅㡆㡇㡈㡉㡊㡋㡌㡍㡎㡏㡐㡑㡒㡓㡔㡕㡖㡗㡘㡙㡚㡛㡜㡝㡞㡟㡠㡡㡢㡣㡤㡥㡦㡧㡨㡩㡪㡫㡬㡭㡮㡯㡰㡱㡲㡳\" # noqa: E501\n + \"㡴㡵㡶㡷㡸㡹㡺㡻㡼㡽㡾㡿㢀㢁㢂㢃㢄㢅㢆㢇㢈㢉㢊㢋㢌㢍㢎㢏㢐㢑㢒㢓㢔㢕㢖㢗㢘㢙㢚㢛㢜㢝㢞㢟㢠㢡㢢㢣㢤㢥㢦㢧㢨㢩㢪㢫㢬㢭㢮㢯㢰㢱㢲㢳㢴㢵㢶\" # noqa: E501\n + \"㢷㢸㢹㢺㢻㢼㢽㢾㢿㣀㣁㣂㣃㣄㣅㣆㣇㣈㣉㣊㣋㣌㣍㣎㣏㣐㣑㣒㣓㣔㣕㣖㣗㣘㣙㣚㣛㣜㣝㣞㣟㣠㣡㣢㣣㣤㣥㣦㣧㣨㣩㣪㣫㣬㣭㣮㣯㣰㣱㣲㣳㣴㣵㣶㣷㣸㣹\" # noqa: E501\n + \"㣺㣻㣼㣽㣾㣿㤀㤁㤂㤃㤄㤅㤆㤇㤈㤉㤊㤋㤌㤍㤎㤏㤐㤑㤒㤓㤔㤕㤖㤗㤘㤙㤚㤛㤜㤝㤞㤟㤠㤡㤢㤣㤤㤥㤦㤧㤨㤩㤪㤫㤬㤭㤮㤯㤰㤱㤲㤳㤴㤵㤶㤷㤸㤹㤺㤻㤼\" # noqa: E501\n + \"㤽㤾㤿㥀㥁㥂㥃㥄㥅㥆㥇㥈㥉㥊㥋㥌㥍㥎㥏㥐㥑㥒㥓㥔㥕㥖㥗㥘㥙㥚㥛㥜㥝㥞㥟㥠㥡㥢㥣㥤㥥㥦㥧㥨㥩㥪㥫㥬㥭㥮㥯㥰㥱㥲㥳㥴㥵㥶㥷㥸㥹㥺㥻㥼㥽㥾㥿\" # noqa: E501\n + \"㦀㦁㦂㦃㦄㦅㦆㦇㦈㦉㦊㦋㦌㦍㦎㦏㦐㦑㦒㦓㦔㦕㦖㦗㦘㦙㦚㦛㦜㦝㦞㦟㦠㦡㦢㦣㦤㦥㦦㦧㦨㦩㦪㦫㦬㦭㦮㦯㦰㦱㦲㦳㦴㦵㦶㦷㦸㦹㦺㦻㦼㦽㦾㦿㧀㧁㧂\" # noqa: E501\n + \"㧃㧄㧅㧆㧇㧈㧉㧊㧋㧌㧍㧎㧏㧐㧑㧒㧓㧔㧕㧖㧗㧘㧙㧚㧛㧜㧝㧞㧟㧠㧡㧢㧣㧤㧥㧦㧧㧨㧩㧪㧫㧬㧭㧮㧯㧰㧱㧲㧳㧴㧵㧶㧷㧸㧹㧺㧻㧼㧽㧾㧿㨀㨁㨂㨃㨄㨅\" # noqa: E501\n + \"㨆㨇㨈㨉㨊㨋㨌㨍㨎㨏㨐㨑㨒㨓㨔㨕㨖㨗㨘㨙㨚㨛㨜㨝㨞㨟㨠㨡㨢㨣㨤㨥㨦㨧㨨㨩㨪㨫㨬㨭㨮㨯㨰㨱㨲㨳㨴㨵㨶㨷㨸㨹㨺㨻㨼㨽㨾㨿㩀㩁㩂㩃㩄㩅㩆㩇㩈\" # noqa: E501\n + \"㩉㩊㩋㩌㩍㩎㩏㩐㩑㩒㩓㩔㩕㩖㩗㩘㩙㩚㩛㩜㩝㩞㩟㩠㩡㩢㩣㩤㩥㩦㩧㩨㩩㩪㩫㩬㩭㩮㩯㩰㩱㩲㩳㩴㩵㩶㩷㩸㩹㩺㩻㩼㩽㩾㩿㪀㪁㪂㪃㪄㪅㪆㪇㪈㪉㪊㪋\" # noqa: E501\n + \"㪌㪍㪎㪏㪐㪑㪒㪓㪔㪕㪖㪗㪘㪙㪚㪛㪜㪝㪞㪟㪠㪡㪢㪣㪤㪥㪦㪧㪨㪩㪪㪫㪬㪭㪮㪯㪰㪱㪲㪳㪴㪵㪶㪷㪸㪹㪺㪻㪼㪽㪾㪿㫀㫁㫂㫃㫄㫅㫆㫇㫈㫉㫊㫋㫌㫍㫎\" # noqa: E501\n + \"㫏㫐㫑㫒㫓㫔㫕㫖㫗㫘㫙㫚㫛㫜㫝㫞㫟㫠㫡㫢㫣㫤㫥㫦㫧㫨㫩㫪㫫㫬㫭㫮㫯㫰㫱㫲㫳㫴㫵㫶㫷㫸㫹㫺㫻㫼㫽㫾㫿㬀㬁㬂㬃㬄㬅㬆㬇㬈㬉㬊㬋㬌㬍㬎㬏㬐㬑\" # noqa: E501\n + \"㬒㬓㬔㬕㬖㬗㬘㬙㬚㬛㬜㬝㬞㬟㬠㬡㬢㬣㬤㬥㬦㬧㬨㬩㬪㬫㬬㬭㬮㬯㬰㬱㬲㬳㬴㬵㬶㬷㬸㬹㬺㬻㬼㬽㬾㬿㭀㭁㭂㭃㭄㭅㭆㭇㭈㭉㭊㭋㭌㭍㭎㭏㭐㭑㭒㭓㭔\" # noqa: E501\n + \"㭕㭖㭗㭘㭙㭚㭛㭜㭝㭞㭟㭠㭡㭢㭣㭤㭥㭦㭧㭨㭩㭪㭫㭬㭭㭮㭯㭰㭱㭲㭳㭴㭵㭶㭷㭸㭹㭺㭻㭼㭽㭾㭿㮀㮁㮂㮃㮄㮅㮆㮇㮈㮉㮊㮋㮌㮍㮎㮏㮐㮑㮒㮓㮔㮕㮖㮗\" # noqa: E501\n + \"㮘㮙㮚㮛㮜㮝㮞㮟㮠㮡㮢㮣㮤㮥㮦㮧㮨㮩㮪㮫㮬㮭㮮㮯㮰㮱㮲㮳㮴㮵㮶㮷㮸㮹㮺㮻㮼㮽㮾㮿㯀㯁㯂㯃㯄㯅㯆㯇㯈㯉㯊㯋㯌㯍㯎㯏㯐㯑㯒㯓㯔㯕㯖㯗㯘㯙㯚\" # noqa: E501\n + \"㯛㯜㯝㯞㯟㯠㯡㯢㯣㯤㯥㯦㯧㯨㯩㯪㯫㯬㯭㯮㯯㯰㯱㯲㯳㯴㯵㯶㯷㯸㯹㯺㯻㯼㯽㯾㯿㰀㰁㰂㰃㰄㰅㰆㰇㰈㰉㰊㰋㰌㰍㰎㰏㰐㰑㰒㰓㰔㰕㰖㰗㰘㰙㰚㰛㰜㰝\" # noqa: E501\n + \"㰞㰟㰠㰡㰢㰣㰤㰥㰦㰧㰨㰩㰪㰫㰬㰭㰮㰯㰰㰱㰲㰳㰴㰵㰶㰷㰸㰹㰺㰻㰼㰽㰾㰿㱀㱁㱂㱃㱄㱅㱆㱇㱈㱉㱊㱋㱌㱍㱎㱏㱐㱑㱒㱓㱔㱕㱖㱗㱘㱙㱚㱛㱜㱝㱞㱟㱠\" # noqa: E501\n + \"㱡㱢㱣㱤㱥㱦㱧㱨㱩㱪㱫㱬㱭㱮㱯㱰㱱㱲㱳㱴㱵㱶㱷㱸㱹㱺㱻㱼㱽㱾㱿㲀㲁㲂㲃㲄㲅㲆㲇㲈㲉㲊㲋㲌㲍㲎㲏㲐㲑㲒㲓㲔㲕㲖㲗㲘㲙㲚㲛㲜㲝㲞㲟㲠㲡㲢㲣\" # noqa: E501\n + \"㲤㲥㲦㲧㲨㲩㲪㲫㲬㲭㲮㲯㲰㲱㲲㲳㲴㲵㲶㲷㲸㲹㲺㲻㲼㲽㲾㲿㳀㳁㳂㳃㳄㳅㳆㳇㳈㳉㳊㳋㳌㳍㳎㳏㳐㳑㳒㳓㳔㳕㳖㳗㳘㳙㳚㳛㳜㳝㳞㳟㳠㳡㳢㳣㳤㳥㳦\" # noqa: E501\n + \"㳧㳨㳩㳪㳫㳬㳭㳮㳯㳰㳱㳲㳳㳴㳵㳶㳷㳸㳹㳺㳻㳼㳽㳾㳿㴀㴁㴂㴃㴄㴅㴆㴇㴈㴉㴊㴋㴌㴍㴎㴏㴐㴑㴒㴓㴔㴕㴖㴗㴘㴙㴚㴛㴜㴝㴞㴟㴠㴡㴢㴣㴤㴥㴦㴧㴨㴩\" # noqa: E501\n + \"㴪㴫㴬㴭㴮㴯㴰㴱㴲㴳㴴㴵㴶㴷㴸㴹㴺㴻㴼㴽㴾㴿㵀㵁㵂㵃㵄㵅㵆㵇㵈㵉㵊㵋㵌㵍㵎㵏㵐㵑㵒㵓㵔㵕㵖㵗㵘㵙㵚㵛㵜㵝㵞㵟㵠㵡㵢㵣㵤㵥㵦㵧㵨㵩㵪㵫㵬\" # noqa: E501\n + \"㵭㵮㵯㵰㵱㵲㵳㵴㵵㵶㵷㵸㵹㵺㵻㵼㵽㵾㵿㶀㶁㶂㶃㶄㶅㶆㶇㶈㶉㶊㶋㶌㶍㶎㶏㶐㶑㶒㶓㶔㶕㶖㶗㶘㶙㶚㶛㶜㶝㶞㶟㶠㶡㶢㶣㶤㶥㶦㶧㶨㶩㶪㶫㶬㶭㶮㶯\" # noqa: E501\n + \"㶰㶱㶲㶳㶴㶵㶶㶷㶸㶹㶺㶻㶼㶽㶾㶿㷀㷁㷂㷃㷄㷅㷆㷇㷈㷉㷊㷋㷌㷍㷎㷏㷐㷑㷒㷓㷔㷕㷖㷗㷘㷙㷚㷛㷜㷝㷞㷟㷠㷡㷢㷣㷤㷥㷦㷧㷨㷩㷪㷫㷬㷭㷮㷯㷰㷱㷲\" # noqa: E501\n + \"㷳㷴㷵㷶㷷㷸㷹㷺㷻㷼㷽㷾㷿㸀㸁㸂㸃㸄㸅㸆㸇㸈㸉㸊㸋㸌㸍㸎㸏㸐㸑㸒㸓㸔㸕㸖㸗㸘㸙㸚㸛㸜㸝㸞㸟㸠㸡㸢㸣㸤㸥㸦㸧㸨㸩㸪㸫㸬㸭㸮㸯㸰㸱㸲㸳㸴㸵\" # noqa: E501\n + \"㸶㸷㸸㸹㸺㸻㸼㸽㸾㸿㹀㹁㹂㹃㹄㹅㹆㹇㹈㹉㹊㹋㹌㹍㹎㹏㹐㹑㹒㹓㹔㹕㹖㹗㹘㹙㹚㹛㹜㹝㹞㹟㹠㹡㹢㹣㹤㹥㹦㹧㹨㹩㹪㹫㹬㹭㹮㹯㹰㹱㹲㹳㹴㹵㹶㹷㹸\" # noqa: E501\n + \"㹹㹺㹻㹼㹽㹾㹿㺀㺁㺂㺃㺄㺅㺆㺇㺈㺉㺊㺋㺌㺍㺎㺏㺐㺑㺒㺓㺔㺕㺖㺗㺘㺙㺚㺛㺜㺝㺞㺟㺠㺡㺢㺣㺤㺥㺦㺧㺨㺩㺪㺫㺬㺭㺮㺯㺰㺱㺲㺳㺴㺵㺶㺷㺸㺹㺺㺻\" # noqa: E501\n + \"㺼㺽㺾㺿㻀㻁㻂㻃㻄㻅㻆㻇㻈㻉㻊㻋㻌㻍㻎㻏㻐㻑㻒㻓㻔㻕㻖㻗㻘㻙㻚㻛㻜㻝㻞㻟㻠㻡㻢㻣㻤㻥㻦㻧㻨㻩㻪㻫㻬㻭㻮㻯㻰㻱㻲㻳㻴㻵㻶㻷㻸㻹㻺㻻㻼㻽㻾\" # noqa: E501\n + \"㻿㼀㼁㼂㼃㼄㼅㼆㼇㼈㼉㼊㼋㼌㼍㼎㼏㼐㼑㼒㼓㼔㼕㼖㼗㼘㼙㼚㼛㼜㼝㼞㼟㼠㼡㼢㼣㼤㼥㼦㼧㼨㼩㼪㼫㼬㼭㼮㼯㼰㼱㼲㼳㼴㼵㼶㼷㼸㼹㼺㼻㼼㼽㼾㼿㽀㽁\" # noqa: E501\n + \"㽂㽃㽄㽅㽆㽇㽈㽉㽊㽋㽌㽍㽎㽏㽐㽑㽒㽓㽔㽕㽖㽗㽘㽙㽚㽛㽜㽝㽞㽟㽠㽡㽢㽣㽤㽥㽦㽧㽨㽩㽪㽫㽬㽭㽮㽯㽰㽱㽲㽳㽴㽵㽶㽷㽸㽹㽺㽻㽼㽽㽾㽿㾀㾁㾂㾃㾄\" # noqa: E501\n + \"㾅㾆㾇㾈㾉㾊㾋㾌㾍㾎㾏㾐㾑㾒㾓㾔㾕㾖㾗㾘㾙㾚㾛㾜㾝㾞㾟㾠㾡㾢㾣㾤㾥㾦㾧㾨㾩㾪㾫㾬㾭㾮㾯㾰㾱㾲㾳㾴㾵㾶㾷㾸㾹㾺㾻㾼㾽㾾㾿㿀㿁㿂㿃㿄㿅㿆㿇\" # noqa: E501\n + \"㿈㿉㿊㿋㿌㿍㿎㿏㿐㿑㿒㿓㿔㿕㿖㿗㿘㿙㿚㿛㿜㿝㿞㿟㿠㿡㿢㿣㿤㿥㿦㿧㿨㿩㿪㿫㿬㿭㿮㿯㿰㿱㿲㿳㿴㿵㿶㿷㿸㿹㿺㿻㿼㿽㿾㿿䀀䀁䀂䀃䀄䀅䀆䀇䀈䀉䀊\" # noqa: E501\n + \"䀋䀌䀍䀎䀏䀐䀑䀒䀓䀔䀕䀖䀗䀘䀙䀚䀛䀜䀝䀞䀟䀠䀡䀢䀣䀤䀥䀦䀧䀨䀩䀪䀫䀬䀭䀮䀯䀰䀱䀲䀳䀴䀵䀶䀷䀸䀹䀺䀻䀼䀽䀾䀿䁀䁁䁂䁃䁄䁅䁆䁇䁈䁉䁊䁋䁌䁍\" # noqa: E501\n + \"䁎䁏䁐䁑䁒䁓䁔䁕䁖䁗䁘䁙䁚䁛䁜䁝䁞䁟䁠䁡䁢䁣䁤䁥䁦䁧䁨䁩䁪䁫䁬䁭䁮䁯䁰䁱䁲䁳䁴䁵䁶䁷䁸䁹䁺䁻䁼䁽䁾䁿䂀䂁䂂䂃䂄䂅䂆䂇䂈䂉䂊䂋䂌䂍䂎䂏䂐\" # noqa: E501\n + \"䂑䂒䂓䂔䂕䂖䂗䂘䂙䂚䂛䂜䂝䂞䂟䂠䂡䂢䂣䂤䂥䂦䂧䂨䂩䂪䂫䂬䂭䂮䂯䂰䂱䂲䂳䂴䂵䂶䂷䂸䂹䂺䂻䂼䂽䂾䂿䃀䃁䃂䃃䃄䃅䃆䃇䃈䃉䃊䃋䃌䃍䃎䃏䃐䃑䃒䃓\" # noqa: E501\n + \"䃔䃕䃖䃗䃘䃙䃚䃛䃜䃝䃞䃟䃠䃡䃢䃣䃤䃥䃦䃧䃨䃩䃪䃫䃬䃭䃮䃯䃰䃱䃲䃳䃴䃵䃶䃷䃸䃹䃺䃻䃼䃽䃾䃿䄀䄁䄂䄃䄄䄅䄆䄇䄈䄉䄊䄋䄌䄍䄎䄏䄐䄑䄒䄓䄔䄕䄖\" # noqa: E501\n + \"䄗䄘䄙䄚䄛䄜䄝䄞䄟䄠䄡䄢䄣䄤䄥䄦䄧䄨䄩䄪䄫䄬䄭䄮䄯䄰䄱䄲䄳䄴䄵䄶䄷䄸䄹䄺䄻䄼䄽䄾䄿䅀䅁䅂䅃䅄䅅䅆䅇䅈䅉䅊䅋䅌䅍䅎䅏䅐䅑䅒䅓䅔䅕䅖䅗䅘䅙\" # noqa: E501\n + \"䅚䅛䅜䅝䅞䅟䅠䅡䅢䅣䅤䅥䅦䅧䅨䅩䅪䅫䅬䅭䅮䅯䅰䅱䅲䅳䅴䅵䅶䅷䅸䅹䅺䅻䅼䅽䅾䅿䆀䆁䆂䆃䆄䆅䆆䆇䆈䆉䆊䆋䆌䆍䆎䆏䆐䆑䆒䆓䆔䆕䆖䆗䆘䆙䆚䆛䆜\" # noqa: E501\n + \"䆝䆞䆟䆠䆡䆢䆣䆤䆥䆦䆧䆨䆩䆪䆫䆬䆭䆮䆯䆰䆱䆲䆳䆴䆵䆶䆷䆸䆹䆺䆻䆼䆽䆾䆿䇀䇁䇂䇃䇄䇅䇆䇇䇈䇉䇊䇋䇌䇍䇎䇏䇐䇑䇒䇓䇔䇕䇖䇗䇘䇙䇚䇛䇜䇝䇞䇟\" # noqa: E501\n + \"䇠䇡䇢䇣䇤䇥䇦䇧䇨䇩䇪䇫䇬䇭䇮䇯䇰䇱䇲䇳䇴䇵䇶䇷䇸䇹䇺䇻䇼䇽䇾䇿䈀䈁䈂䈃䈄䈅䈆䈇䈈䈉䈊䈋䈌䈍䈎䈏䈐䈑䈒䈓䈔䈕䈖䈗䈘䈙䈚䈛䈜䈝䈞䈟䈠䈡䈢\" # noqa: E501\n + \"䈣䈤䈥䈦䈧䈨䈩䈪䈫䈬䈭䈮䈯䈰䈱䈲䈳䈴䈵䈶䈷䈸䈹䈺䈻䈼䈽䈾䈿䉀䉁䉂䉃䉄䉅䉆䉇䉈䉉䉊䉋䉌䉍䉎䉏䉐䉑䉒䉓䉔䉕䉖䉗䉘䉙䉚䉛䉜䉝䉞䉟䉠䉡䉢䉣䉤䉥\" # noqa: E501\n + \"䉦䉧䉨䉩㑃䉪䉫䉬䉭䉮䉯䉰䉱䉲䉳䉴䉵䉶䉷䉸䉹䉺䉻䉼䉽䉾䉿䊀䊁䊂䊃䊄䊅䊆䊇䊈䊉䊊䊋䊌䊍䊎䊏䊐䊑䊒䊓䊔䊕䊖䊗䊘䊙䊚䊛䊜䊝䊞䊟䊠䊡䊢䊣䊤䊥䊦䊧\" # noqa: E501\n + \"䊨䊩䊪䊫䊬䊭䊮䊯䊰䊱䊲䊳䊴䊵䊶䊷䊸䊹䊺䊻䊼䊽䊾䊿䋀䋁䋂䋃䋄䋅䋆䋇䋈䋉䋊䋋䋌䋍䋎䋏䋐䋑䋒䋓䋔䋕䋖䋗䋘䋙䋚䋛䋜䋝䋞䋟䋠䋡䋢䋣䋤䋥䋦䋧䋨䋩䋪\" # noqa: E501\n + \"䋫䋬䋭䋮䋯䋰䋱䋲䋳䋴䋵䋶䋷䋸䋹䋺䋻䋼䋽䋾䋿䌀䌁䌂䌃䌄䌅䌆䌇䌈䌉䌊䌋䌌䌍䌎䌏䌐䌑䌒䌓䌔䌕䌖䌗䌘䌙䌚䌛䌜䌝䌞䌟䌠䌡䌢䌣䌤䌥䌦䌧䌨䌩䌪䌫䌬䌭\" # noqa: E501\n + \"䌮䌯䌰䌱䌲䌳䌴䌵䌶䌷䌸䌹䌺䌻䌼䌽䌾䌿䍀䍁䍂䍃䍄䍅䍆䍇䍈䍉䍊䍋䍌䍍䍎䍏䍐䍑䍒䍓䍔䍕䍖䍗䍘䍙䍚䍛䍜䍝䍞䍟䍠䍡䍢䍣䍤䍥䍦䍧䍨䍩䍪䍫䍬䍭䍮䍯䍰\" # noqa: E501\n + \"䍱䍲䍳䍴䍵䍶䍷䍸䍹䍺䍻䍼䍽䍾䍿䎀䎁䎂䎃䎄䎅䎆䎇䎈䎉䎊䎋䎌䎍䎎䎏䎐䎑䎒䎓䎔䎕䎖䎗䎘䎙䎚䎛䎜䎝䎞䎟䎠䎡䎢䎣䎤䎥䎦䎧䎨䎩䎪䎫䎬䎭䎮䎯䎰䎱䎲䎳\" # noqa: E501\n + \"䎴䎵䎶䎷䎸䎹䎺䎻䎼䎽䎾䎿䏀䏁䏂䏃䏄䏅䏆䏇䏈䏉䏊䏋䏌䏍䏎䏏䏐䏑䏒䏓䏔䏕䏖䏗䏘䏙䏚䏛䏜䏝䏞䏟䏠䏡䏢䏣䏤䏥䏦䏧䏨䏩䏪䏫䏬䏭䏮䏯䏰䏱䏲䏳䏴䏵䏶\" # noqa: E501\n + \"䏷䏸䏹䏺䏻䏼䏽䏾䏿䐀䐁䐂䐃䐄䐅䐆䐇䐈䐉䐊䐋䐌䐍䐎䐏䐐䐑䐒䐓䐔䐕䐖䐗䐘䐙䐚䐛䐜䐝䐞䐟䐠䐡䐢䐣䐤䐥䐦䐧䐨䐩䐪䐫䐬䐭䐮䐯䐰䐱䐲䐳䐴䐵䐶䐷䐸䐹\" # noqa: E501\n + \"䐺䐻䐼䐽䐾䐿䑀䑁䑂䑃䑄䑅䑆䑇䑈䑉䑊䑋䑌䑍䑎䑏䑐䑑䑒䑓䑔䑕䑖䑗䑘䑙䑚䑛䑜䑝䑞䑟䑠䑡䑢䑣䑤䑥䑦䑧䑨䑩䑪䑫䑬䑭䑮䑯䑰䑱䑲䑳䑴䑵䑶䑷䑸䑹䑺䑻䑼\" # noqa: E501\n + \"䑽䑾䑿䒀䒁䒂䒃䒄䒅䒆䒇䒈䒉䒊䒋䒌䒍䒎䒏䒐䒑䒒䒓䒔䒕䒖䒗䒘䒙䒚䒛䒜䒝䒞䒟䒠䒡䒢䒣䒤䒥䒦䒧䒨䒩䒪䒫䒬䒭䒮䒯䒰䒱䒲䒳䒴䒵䒶䒷䒸䒹䒺䒻䒼䒽䒾䒿\" # noqa: E501\n + \"䓀䓁䓂䓃䓄䓅䓆䓇䓈䓉䓊䓋䓌䓍䓎䓏䓐䓑䓒䓓䓔䓕䓖䓗䓘䓙䓚䓛䓜䓝䓞䓟䓠䓡䓢䓣䓤䓥䓦䓧䓨䓩䓪䓫䓬䓭䓮䓯䓰䓱䓲䓳䓴䓵䓶䓷䓸䓹䓺䓻䓼䓽䓾䓿䔀䔁䔂\" # noqa: E501\n + \"䔃䔄䔅䔆䔇䔈䔉䔊䔋䔌䔍䔎䔏䔐䔑䔒䔓䔔䔕䔖䔗䔘䔙䔚䔛䔜䔝䔞䔟䔠䔡䔢䔣䔤䔥䔦䔧䔨䔩䔪䔫䔬䔭䔮䔯䔰䔱䔲䔳䔴䔵䔶䔷䔸䔹䔺䔻䔼䔽䔾䔿䕀䕁䕂䕃䕄䕅\" # noqa: E501\n + \"䕆䕇䕈䕉䕊䕋䕌䕍䕎䕏䕐䕑䕒䕓䕔䕕䕖䕗䕘䕙䕚䕛䕜䕝䕞䕟䕠䕡䕢䕣䕤䕥䕦䕧䕨䕩䕪䕫䕬䕭䕮䕯䕰䕱䕲䕳䕴䕵䕶䕷䕸䕹䕺䕻䕼䕽䕾䕿䖀䖁䖂䖃䖄䖅䖆䖇䖈\" # noqa: E501\n + \"䖉䖊䖋䖌䖍䖎䖏䖐䖑䖒䖓䖔䖕䖖䖗䖘䖙䖚䖛䖜䖝䖞䖟䖠䖡䖢䖣䖤䖥䖦䖧䖨䖩䖪䖫䖬䖭䖮䖯䖰䖱䖲䖳䖴䖵䖶䖷䖸䖹䖺䖻䖼䖽䖾䖿䗀䗁䗂䗃䗄䗅䗆䗇䗈䗉䗊䗋\" # noqa: E501\n + \"䗌䗍䗎䗏䗐䗑䗒䗓䗔䗕䗖䗗䗘䗙䗚䗛䗜䗝䗞䗟䗠䗡䗢䗣䗤䗥䗦䗧䗨䗩䗪䗫䗬䗭䗮䗯䗰䗱䗲䗳䗴䗵䗶䗷䗸䗹䗺䗻䗼䗽䗾䗿䘀䘁䘂䘃䘄䘅䘆䘇䘈䘉䘊䘋䘌䘍䘎\" # noqa: E501\n + \"䘏䘐䘑䘒䘓䘔䘕䘖䘗䘘䘙䘚䘛䘜䘝䘞䘟䘠䘡䘢䘣䘤䘥䘦䘧䘨䘩䘪䘫䘬䘭䘮䘯䘰䘱䘲䘳䘴䘵䘶䘷䘸䘹䘺䘻䘼䘽䘾䘿䙀䙁䙂䙃䙄䙅䙆䙇䙈䙉䙊䙋䙌䙍䙎䙏䙐䙑\" # noqa: E501\n + \"䙒䙓䙔䙕䙖䙗䙘䙙䙚䙛䙜䙝䙞䙟䙠䙡䙢䙣䙤䙥䙦䙧䙨䙩䙪䙫䙬䙭䙮䙯䙰䙱䙲䙳䙴䙵䙶䙷䙸䙹䙺䙻䙼䙽䙾䙿䚀䚁䚂䚃䚄䚅䚆䚇䚈䚉䚊䚋䚌䚍䚎䚏䚐䚑䚒䚓䚔\" # noqa: E501\n + \"䚕䚖䚗䚘䚙䚚䚛䚜䚝䚞䚟䚠䚡䚢䚣䚤䚥䚦䚧䚨䚩䚪䚫䚬䚭䚮䚯䚰䚱䚲䚳䚴䚵䚶䚷䚸䚹䚺䚻䚼䚽䚾䚿䛀䛁䛂䛃䛄䛅䛆䛇䛈䛉䛊䛋䛌䛍䛎䛏䛐䛑䛒䛓䛔䛕䛖䛗\" # noqa: E501\n + \"䛘䛙䛚䛛䛜䛝䛞䛟䛠䛡䛢䛣䛤䛥䛦䛧䛨䛩䛪䛫䛬䛭䛮䛯䛰䛱䛲䛳䛴䛵䛶䛷䛸䛹䛺䛻䛼䛽䛾䛿䜀䜁䜂䜃䜄䜅䜆䜇䜈䜉䜊䜋䜌䜍䜎䜏䜐䜑䜒䜓䜔䜕䜖䜗䜘䜙䜚\" # noqa: E501\n + \"䜛䜜䜝䜞䜟䜠䜡䜢䜣䜤䜥䜦䜧䜨䜩䜪䜫䜬䜭䜮䜯䜰䜱䜲䜳䜴䜵䜶䜷䜸䜹䜺䜻䜼䜽䜾䜿䝀䝁䝂䝃䝄䝅䝆䝇䝈䝉䝊䝋䝌䝍䝎䝏䝐䝑䝒䝓䝔䝕䝖䝗䝘䝙䝚䝛䝜䝝\" # noqa: E501\n + \"䝞䝟䝠䝡䝢䝣䝤䝥䝦䝧䝨䝩䝪䝫䝬䝭䝮䝯䝰䝱䝲䝳䝴䝵䝶䝷䝸䝹䝺䝻䝼䝽䝾䝿䞀䞁䞂䞃䞄䞅䞆䞇䞈䞉䞊䞋䞌䞍䞎䞏䞐䞑䞒䞓䞔䞕䞖䞗䞘䞙䞚䞛䞜䞝䞞䞟䞠\" # noqa: E501\n + \"䞡䞢䞣䞤䞥䞦䞧䞨䞩䞪䞫䞬䞭䞮䞯䞰䞱䞲䞳䞴䞵䞶䞷䞸䞹䞺䞻䞼䞽䞾䞿䟀䟁䟂䟃䟄䟅䟆䟇䟈䟉䟊䟋䟌䟍䟎䟏䟐䟑䟒䟓䟔䟕䟖䟗䟘䟙䟚䟛䟜䟝䟞䟟䟠䟡䟢䟣\" # noqa: E501\n + \"䟤䟥䟦䟧䟨䟩䟪䟫䟬䟭䟮䟯䟰䟱䟲䟳䟴䟵䟶䟷䟸䟹䟺䟻䟼䟽䟾䟿䠀䠁䠂䠃䠄䠅䠆䠇䠈䠉䠊䠋䠌䠍䠎䠏䠐䠑䠒䠓䠔䠕䠖䠗䠘䠙䠚䠛䠜䠝䠞䠟䠠䠡䠢䠣䠤䠥䠦\" # noqa: E501\n + \"䠧䠨䠩䠪䠫䠬䠭䠮䠯䠰䠱䠲䠳䠴䠵䠶䠷䠸䠹䠺䠻䠼䠽䠾䠿䡀䡁䡂䡃䡄䡅䡆䡇䡈䡉䡊䡋䡌䡍䡎䡏䡐䡑䡒䡓䡔䡕䡖䡗䡘䡙䡚䡛䡜䡝䡞䡟䡠䡡䡢䡣䡤䡥䡦䡧䡨䡩\" # noqa: E501\n + \"䡪䡫䡬䡭䡮䡯䡰䡱䡲䡳䡴䡵䡶䡷䡸䡹䡺䡻䡼䡽䡾䡿䢀䢁䢂䢃䢄䢅䢆䢇䢈䢉䢊䢋䢌䢍䢎䢏䢐䢑䢒䢓䢔䢕䢖䢗䢘䢙䢚䢛䢜䢝䢞䢟䢠䢡䢢䢣䢤䢥䢦䢧䢨䢩䢪䢫䢬\" # noqa: E501\n + \"䢭䢮䢯䢰䢱䢲䢳䢴䢵䢶䢷䢸䢹䢺䢻䢼䢽䢾䢿䣀䣁䣂䣃䣄䣅䣆䣇䣈䣉䣊䣋䣌䣍䣎䣏䣐䣑䣒䣓䣔䣕䣖䣗䣘䣙䣚䣛䣜䣝䣞䣟䣠䣡䣢䣣䣤䣥䣦䣧䣨䣩䣪䣫䣬䣭䣮䣯\" # noqa: E501\n + \"䣰䣱䣲䣳䣴䣵䣶䣷䣸䣹䣺䣻䣼䣽䣾䣿䤀䤁䤂䤃䤄䤅䤆䤇䤈䤉䤊䤋䤌䤍䤎䤏䤐䤑䤒䤓䤔䤕䤖䤗䤘䤙䤚䤛䤜䤝䤞䤟䤠䤡䤢䤣䤤䤥䤦䤧䤨䤩䤪䤫䤬䤭䤮䤯䤰䤱䤲\" # noqa: E501\n + \"䤳䤴䤵䤶䤷䤸䤹䤺䤻䤼䤽䤾䤿䥀䥁䥂䥃䥄䥅䥆䥇䥈䥉䥊䥋䥌䥍䥎䥏䥐䥑䥒䥓䥔䥕䥖䥗䥘䥙䥚䥛䥜䥝䥞䥟䥠䥡䥢䥣䥤䥥䥦䥧䥨䥩䥪䥫䥬䥭䥮䥯䥰䥱䥲䥳䥴䥵\" # noqa: E501\n + \"䥶䥷䥸䥹䥺䥻䥼䥽䥾䥿䦀䦁䦂䦃䦄䦅䦆䦇䦈䦉䦊䦋䦌䦍䦎䦏䦐䦑䦒䦓䦔䦕䦖䦗䦘䦙䦚䦛䦜䦝䦞䦟䦠䦡䦢䦣䦤䦥䦦䦧䦨䦩䦪䦫䦬䦭䦮䦯䦰䦱䦲䦳䦴䦵䦶䦷䦸\" # noqa: E501\n + \"䦹䦺䦻䦼䦽䦾䦿䧀䧁䧂䧃䧄䧅䧆䧇䧈䧉䧊䧋䧌䧍䧎䧏䧐䧑䧒䧓䧔䧕䧖䧗䧘䧙䧚䧛䧜䧝䧞䧟䧠䧡䧢䧣䧤䧥䧦䧧䧨䧩䧪䧫䧬䧭䧮䧯䧰䧱䧲䧳䧴䧵䧶䧷䧸䧹䧺䧻\" # noqa: E501\n + \"䧼䧽䧾䧿䨀䨁䨂䨃䨄䨅䨆䨇䨈䨉䨊䨋䨌䨍䨎䨏䨐䨑䨒䨓䨔䨕䨖䨗䨘䨙䨚䨛䨜䨝䨞䨟䨠䨡䨢䨣䨤䨥䨦䨧䨨䨩䨪䨫䨬䨭䨮䨯䨰䨱䨲䨳䨴䨵䨶䨷䨸䨹䨺䨻䨼䨽䨾\" # noqa: E501\n + \"䨿䩀䩁䩂䩃䩄䩅䩆䩇䩈䩉䩊䩋䩌䩍䩎䩏䩐䩑䩒䩓䩔䩕䩖䩗䩘䩙䩚䩛䩜䩝䩞䩟䩠䩡䩢䩣䩤䩥䩦䩧䩨䩩䩪䩫䩬䩭䩮䩯䩰䩱䩲䩳䩴䩵䩶䩷䩸䩹䩺䩻䩼䩽䩾䩿䪀䪁\" # noqa: E501\n + \"䪂䪃䪄䪅䪆䪇䪈䪉䪊䪋䪌䪍䪎䪏䪐䪑䪒䪓䪔䪕䪖䪗䪘䪙䪚䪛䪜䪝䪞䪟䪠䪡䪢䪣䪤䪥䪦䪧䪨䪩䪪䪫䪬䪭䪮䪯䪰䪱䪲䪳䪴䪵䪶䪷䪸䪹䪺䪻䪼䪽䪾䪿䫀䫁䫂䫃䫄\" # noqa: E501\n + \"䫅䫆䫇䫈䫉䫊䫋䫌䫍䫎䫏䫐䫑䫒䫓䫔䫕䫖䫗䫘䫙䫚䫛䫜䫝䫞䫟䫠䫡䫢䫣䫤䫥䫦䫧䫨䫩䫪䫫䫬䫭䫮䫯䫰䫱䫲䫳䫴䫵䫶䫷䫸䫹䫺䫻䫼䫽䫾䫿䬀䬁䬂䬃䬄䬅䬆䬇\" # noqa: E501\n + \"䬈䬉䬊䬋䬌䬍䬎䬏䬐䬑䬒䬓䬔䬕䬖䬗䬘䬙䬚䬛䬜䬝䬞䬟䬠䬡䬢䬣䬤䬥䬦䬧䬨䬩䬪䬫䬬䬭䬮䬯䬰䬱䬲䬳䬴䬵䬶䬷䬸䬹䬺䬻䬼䬽䬾䬿䭀䭁䭂䭃䭄䭅䭆䭇䭈䭉䭊\" # noqa: E501\n + \"䭋䭌䭍䭎䭏䭐䭑䭒䭓䭔䭕䭖䭗䭘䭙䭚䭛䭜䭝䭞䭟䭠䭡䭢䭣䭤䭥䭦䭧䭨䭩䭪䭫䭬䭭䭮䭯䭰䭱䭲䭳䭴䭵䭶䭷䭸䭹䭺䭻䭼䭽䭾䭿䮀䮁䮂䮃䮄䮅䮆䮇䮈䮉䮊䮋䮌䮍\" # noqa: E501\n + \"䮎䮏䮐䮑䮒䮓䮔䮕䮖䮗䮘䮙䮚䮛䮜䮝䮞䮟䮠䮡䮢䮣䮤䮥䮦䮧䮨䮩䮪䮫䮬䮭䮮䮯䮰䮱䮲䮳䮴䮵䮶䮷䮸䮹䮺䮻䮼䮽䮾䮿䯀䯁䯂䯃䯄䯅䯆䯇䯈䯉䯊䯋䯌䯍䯎䯏䯐\" # noqa: E501\n + \"䯑䯒䯓䯔䯕䯖䯗䯘䯙䯚䯛䯜䯝䯞䯟䯠䯡䯢䯣䯤䯥䯦䯧䯨䯩䯪䯫䯬䯭䯮䯯䯰䯱䯲䯳䯴䯵䯶䯷䯸䯹䯺䯻䯼䯽䯾䯿䰀䰁䰂䰃䰄䰅䰆䰇䰈䰉䰊䰋䰌䰍䰎䰏䰐䰑䰒䰓\" # noqa: E501\n + \"䰔䰕䰖䰗䰘䰙䰚䰛䰜䰝䰞䰟䰠䰡䰢䰣䰤䰥䰦䰧䰨䰩䰪䰫䰬䰭䰮䰯䰰䰱䰲䰳䰴䰵䰶䰷䰸䰹䰺䰻䰼䰽䰾䰿䱀䱁䱂䱃䱄䱅䱆䱇䱈䱉䱊䱋䱌䱍䱎䱏䱐䱑䱒䱓䱔䱕䱖\" # noqa: E501\n + \"䱗䱘䱙䱚䱛䱜䱝䱞䱟䱠䱡䱢䱣䱤䱥䱦䱧䱨䱩䱪䱫䱬䱭䱮䱯䱰䱱䱲䱳䱴䱵䱶䱷䱸䱹䱺䱻䱼䱽䱾䱿䲀䲁䲂䲃䲄䲅䲆䲇䲈䲉䲊䲋䲌䲍䲎䲏䲐䲑䲒䲓䲔䲕䲖䲗䲘䲙\" # noqa: E501\n + \"䲚䲛䲜䲝䲞䲟䲠䲡䲢䲣䲤䲥䲦䲧䲨䲩䲪䲫䲬䲭䲮䲯䲰䲱䲲䲳䲴䲵䲶䲷䲸䲹䲺䲻䲼䲽䲾䲿䳀䳁䳂䳃䳄䳅䳆䳇䳈䳉䳊䳋䳌䳍䳎䳏䳐䳑䳒䳓䳔䳕䳖䳗䳘䳙䳚䳛䳜\" # noqa: E501\n + \"䳝䳞䳟䳠䳡䳢䳣䳤䳥䳦䳧䳨䳩䳪䳫䳬䳭䳮䳯䳰䳱䳲䳳䳴䳵䳶䳷䳸䳹䳺䳻䳼䳽䳾䳿䴀䴁䴂䴃䴄䴅䴆䴇䴈䴉䴊䴋䴌䴍䴎䴏䴐䴑䴒䴓䴔䴕䴖䴗䴘䴙䴚䴛䴜䴝䴞䴟\" # noqa: E501\n + \"䴠䴡䴢䴣䴤䴥䴦䴧䴨䴩䴪䴫䴬䴭䴮䴯䴰䴱䴲䴳䴴䴵䴶䴷䴸䴹䴺䴻䴼䴽䴾䴿䵀䵁䵂䵃䵄䵅䵆䵇䵈䵉䵊䵋䵌䵍䵎䵏䵐䵑䵒䵓䵔䵕䵖䵗䵘䵙䵚䵛䵜䵝䵞䵟䵠䵡䵢\" # noqa: E501\n + \"䵣䵤䵥䵦䵧䵨䵩䵪䵫䵬䵭䵮䵯䵰䵱䵲䵳䵴䵵䵶䵷䵸䵹䵺䵻䵼䵽䵾䵿䶀䶁䶂䶃䶄䶅䶆䶇䶈䶉䶊䶋䶌䶍䶎䶏䶐䶑䶒䶓䶔䶕䶖䶗䶘䶙䶚䶛䶜䶝䶞䶟䶠䶡䶢䶣䶤䶥\" # noqa: E501\n + \"䶦䶧䶨䶩䶪䶫䶬䶭䶮䶯䶰䶱䶲䶳䶴䶵䶶䶷䶸䶹䶺䶻䶼䶽䶾䶿\"\n + _BASE_VOCABS[\"punctuation\"]\n + \"。・〜°—、「」『』【】゛》《〉〈\" # punctuation\n + _BASE_VOCABS[\"currency\"]\n)\n\n# Multi-lingual\nVOCABS[\"multilingual\"] = \"\".join(\n dict.fromkeys(\n # latin_based\n VOCABS[\"english\"]\n + VOCABS[\"albanian\"]\n + VOCABS[\"afrikaans\"]\n + VOCABS[\"azerbaijani\"]\n + VOCABS[\"basque\"]\n + VOCABS[\"bosnian\"]\n + VOCABS[\"catalan\"]\n + VOCABS[\"croatian\"]\n + VOCABS[\"czech\"]\n + VOCABS[\"danish\"]\n + VOCABS[\"dutch\"]\n + VOCABS[\"estonian\"]\n + VOCABS[\"esperanto\"]\n + VOCABS[\"french\"]\n + VOCABS[\"finnish\"]\n + VOCABS[\"frisian\"]\n + VOCABS[\"galician\"]\n + VOCABS[\"german\"]\n + VOCABS[\"hausa\"]\n + VOCABS[\"hungarian\"]\n + VOCABS[\"icelandic\"]\n + VOCABS[\"indonesian\"]\n + VOCABS[\"irish\"]\n + VOCABS[\"italian\"]\n + VOCABS[\"latvian\"]\n + VOCABS[\"lithuanian\"]\n + VOCABS[\"luxembourgish\"]\n + VOCABS[\"maori\"]\n + VOCABS[\"malagasy\"]\n + VOCABS[\"malay\"]\n + VOCABS[\"maltese\"]\n + VOCABS[\"montenegrin\"]\n + VOCABS[\"norwegian\"]\n + VOCABS[\"polish\"]\n + VOCABS[\"portuguese\"]\n + VOCABS[\"quechua\"]\n + VOCABS[\"romanian\"]\n + VOCABS[\"scottish_gaelic\"]\n + VOCABS[\"serbian_latin\"]\n + VOCABS[\"slovak\"]\n + VOCABS[\"slovene\"]\n + VOCABS[\"somali\"]\n + VOCABS[\"spanish\"]\n + VOCABS[\"swahili\"]\n + VOCABS[\"swedish\"]\n + VOCABS[\"tagalog\"]\n + VOCABS[\"turkish\"]\n + VOCABS[\"uzbek_latin\"]\n + VOCABS[\"vietnamese\"]\n + VOCABS[\"welsh\"]\n + VOCABS[\"yoruba\"]\n + VOCABS[\"zulu\"]\n + \"§\" # paragraph sign\n # cyrillic_based\n + VOCABS[\"russian\"]\n + VOCABS[\"belarusian\"]\n + VOCABS[\"ukrainian\"]\n + VOCABS[\"tatar\"]\n + VOCABS[\"tajik\"]\n + VOCABS[\"kazakh\"]\n + VOCABS[\"kyrgyz\"]\n + VOCABS[\"bulgarian\"]\n + VOCABS[\"macedonian\"]\n + VOCABS[\"mongolian\"]\n + VOCABS[\"yakut\"]\n + VOCABS[\"serbian_cyrillic\"]\n + VOCABS[\"uzbek_cyrillic\"]\n # greek\n + VOCABS[\"greek\"]\n # hebrew\n + VOCABS[\"hebrew\"]\n )\n)\n"
},
{
"path": "pyproject.toml",
"content": "[build-system]\nrequires = [\"setuptools\", \"wheel\"]\nbuild-backend = \"setuptools.build_meta\"\n\n[project]\nname = \"onnxtr\"\ndescription = \"Onnx Text Recognition (OnnxTR): docTR Onnx-Wrapper for high-performance OCR on documents.\"\nauthors = [{name = \"Felix Dittrich\", email = \"felixdittrich92@gmail.com\"}]\nmaintainers = [\n {name = \"Felix Dittrich\"},\n]\nreadme = \"README.md\"\nrequires-python = \">=3.10.0,<4\"\nlicense = {file = \"LICENSE\"}\nkeywords=[\"OCR\", \"deep learning\", \"computer vision\", \"onnx\", \"text detection\", \"text recognition\", \"docTR\", \"document analysis\", \"document processing\", \"document AI\"]\nclassifiers=[\n \"Development Status :: 4 - Beta\",\n \"Intended Audience :: Developers\",\n \"Intended Audience :: Education\",\n \"Intended Audience :: Science/Research\",\n \"License :: OSI Approved :: Apache Software License\",\n \"Natural Language :: English\",\n \"Operating System :: OS Independent\",\n \"Programming Language :: Python :: 3\",\n \"Programming Language :: Python :: 3.10\",\n \"Programming Language :: Python :: 3.11\",\n \"Programming Language :: Python :: 3.12\",\n \"Topic :: Scientific/Engineering :: Artificial Intelligence\",\n]\ndynamic = [\"version\"]\ndependencies = [\n # For proper typing, mypy needs numpy>=1.20.0 (cf. https://github.com/numpy/numpy/pull/16515)\n # Additional typing support is brought by numpy>=1.22.4, but core build sticks to >=1.16.0\n \"numpy>=1.16.0,<3.0.0\",\n \"scipy>=1.4.0,<2.0.0\",\n \"pypdfium2>=4.11.0,<6.0.0\",\n \"pyclipper>=1.2.0,<2.0.0\",\n \"rapidfuzz>=3.0.0,<4.0.0\",\n \"langdetect>=1.0.9,<2.0.0\",\n \"huggingface-hub>=0.23.0,<2.0.0\",\n \"Pillow>=9.2.0\",\n \"defusedxml>=0.7.0\",\n \"anyascii>=0.3.2\",\n \"tqdm>=4.30.0\",\n]\n\n[project.optional-dependencies]\ncpu = [\n \"onnxruntime>=1.18.0\",\n \"opencv-python>=4.5.0,<5.0.0\",\n]\ngpu = [\n \"onnxruntime-gpu>=1.18.0\",\n \"opencv-python>=4.5.0,<5.0.0\",\n]\nopenvino = [\n \"onnxruntime-openvino>=1.18.0\",\n \"opencv-python>=4.5.0,<5.0.0\",\n]\ncpu-headless = [\n \"onnxruntime>=1.18.0\",\n \"opencv-python-headless>=4.5.0,<5.0.0\",\n]\ngpu-headless = [\n \"onnxruntime-gpu>=1.18.0\",\n \"opencv-python-headless>=4.5.0,<5.0.0\",\n]\nopenvino-headless = [\n \"onnxruntime-openvino>=1.18.0\",\n \"opencv-python-headless>=4.5.0,<5.0.0\",\n]\nhtml = [\n \"weasyprint>=55.0\",\n]\nviz = [\n \"matplotlib>=3.1.0\",\n \"mplcursors>=0.3\",\n]\ntesting = [\n \"pytest>=5.3.2\",\n \"coverage[toml]>=4.5.4\",\n \"requests>=2.20.0\",\n \"pytest-memray>=1.7.0\",\n \"psutil>=7.0.0\",\n]\nquality = [\n \"ruff>=0.1.5\",\n \"mypy>=0.812\",\n \"pre-commit>=2.17.0\",\n]\ndev = [\n # Runtime\n \"onnxruntime>=1.18.0\",\n \"opencv-python>=4.5.0,<5.0.0\",\n # HTML\n \"weasyprint>=55.0\",\n # Visualization\n \"matplotlib>=3.1.0\",\n \"mplcursors>=0.3\",\n # Testing\n \"pytest>=5.3.2\",\n \"coverage[toml]>=4.5.4\",\n \"requests>=2.20.0\",\n \"pytest-memray>=1.7.0\",\n \"psutil>=7.0.0\",\n # Quality\n \"ruff>=0.1.5\",\n \"mypy>=0.812\",\n \"pre-commit>=2.17.0\",\n]\n\n[project.urls]\nrepository = \"https://github.com/felixdittrich92/OnnxTR\"\ntracker = \"https://github.com/felixdittrich92/OnnxTR/issues\"\nchangelog = \"https://github.com/felixdittrich92/OnnxTR/releases\"\n\n[tool.setuptools]\nzip-safe = true\n\n[tool.setuptools.packages.find]\nexclude = [\"docs*\", \"tests*\", \"scripts*\", \"demo*\"]\n\n[tool.setuptools.package-data]\nonnxtr = [\"py.typed\"]\n\n[tool.mypy]\nfiles = \"onnxtr/\"\nshow_error_codes = true\npretty = true\nwarn_unused_ignores = true\nwarn_redundant_casts = true\nno_implicit_optional = true\ncheck_untyped_defs = true\nimplicit_reexport = false\n\n[[tool.mypy.overrides]]\nmodule = [\n \"onnxruntime.*\",\n\t\"PIL.*\",\n\t\"scipy.*\",\n\t\"cv2.*\",\n\t\"matplotlib.*\",\n \"numpy.*\",\n\t\"pyclipper.*\",\n\t\"mplcursors.*\",\n\t\"defusedxml.*\",\n\t\"weasyprint.*\",\n\t\"pypdfium2.*\",\n\t\"langdetect.*\",\n \"huggingface_hub.*\",\n \"rapidfuzz.*\",\n \"anyascii.*\",\n \"tqdm.*\",\n]\nignore_missing_imports = true\n\n[tool.ruff]\nexclude = [\".git\", \"venv*\", \"build\", \"**/__init__.py\"]\nline-length = 120\ntarget-version = \"py310\"\npreview=true\n\n[tool.ruff.lint]\nselect = [\n # https://docs.astral.sh/ruff/rules/\n \"E\", \"W\", \"F\", \"I\", \"N\", \"Q\", \"C4\", \"T10\", \"LOG\",\n \"D101\", \"D103\", \"D201\",\"D202\",\"D207\",\"D208\",\"D214\",\"D215\",\"D300\",\"D301\",\"D417\", \"D419\", \"D207\" # pydocstyle\n]\nignore = [\"E402\", \"E203\", \"F403\", \"E731\", \"N812\", \"N817\", \"C408\", \"LOG015\"]\n\n[tool.ruff.lint.isort]\nknown-first-party = [\"onnxtr\", \"utils\"]\nknown-third-party = [\"onnxruntime\", \"cv2\"]\n\n[tool.ruff.lint.per-file-ignores]\n\"onnxtr/models/**.py\" = [\"N806\", \"F841\"]\n\"tests/**.py\" = [\"D\"]\n\"scripts/**.py\" = [\"D\"]\n\"demo/**.py\" = [\"D\"]\n\".github/**.py\" = [\"D\"]\n\n\n[tool.ruff.lint.flake8-quotes]\ndocstring-quotes = \"double\"\n\n[tool.coverage.run]\nsource = [\"onnxtr\"]\n"
},
{
"path": "scripts/convert_to_float16.py",
"content": "# Copyright (C) 2021-2026, Mindee | Felix Dittrich.\n\n# This program is licensed under the Apache License 2.0.\n# See LICENSE or go to for full license details.\n\ntry:\n from onnxconverter_common import auto_convert_mixed_precision\nexcept ImportError:\n raise ImportError(\"Failed to import onnxconverter_common. Please install `pip install onnxconverter-common`.\")\n\n# Check GPU availability\nimport onnxruntime\n\nif onnxruntime.get_device() != \"GPU\":\n raise RuntimeError(\n \"Please install OnnxTR with GPU support to run this script. \"\n + \"`pip install onnxtr[gpu]` or `pip install -e .[gpu]`\"\n )\n\nimport argparse\nimport time\nfrom tempfile import TemporaryDirectory\nfrom typing import Any\n\nimport numpy as np\nimport onnx\n\nfrom onnxtr.models import classification, detection, recognition\nfrom onnxtr.models.classification.zoo import ORIENTATION_ARCHS\nfrom onnxtr.models.detection.zoo import ARCHS as DETECTION_ARCHS\nfrom onnxtr.models.recognition.zoo import ARCHS as RECOGNITION_ARCHS\n\n\ndef _load_model(arch: str, model_path: str | None = None) -> Any:\n if arch in DETECTION_ARCHS:\n model = detection.__dict__[arch]() if model_path is None else detection.__dict__[arch](model_path)\n elif args.arch in RECOGNITION_ARCHS:\n model = recognition.__dict__[arch]() if model_path is None else recognition.__dict__[arch](model_path)\n elif args.arch in ORIENTATION_ARCHS:\n model = classification.__dict__[arch]() if model_path is None else classification.__dict__[arch](model_path)\n else:\n raise ValueError(f\"Unknown architecture {arch}\")\n return model\n\n\ndef _latency_check(args: Any, size: tuple[int], model: Any, img_tensor: np.ndarray) -> None:\n # Warmup\n for _ in range(10):\n _ = model(img_tensor)\n\n timings = []\n\n # Evaluation runs\n for _ in range(args.it):\n start_ts = time.perf_counter()\n _ = model(img_tensor)\n timings.append(time.perf_counter() - start_ts)\n\n _timings = np.array(timings)\n print(f\"{args.arch} ({args.it} runs on ({size}) inputs)\")\n print(f\"mean {1000 * _timings.mean():.2f}ms, std {1000 * _timings.std():.2f}ms\")\n\n\ndef _validate(fp32_in: list[np.ndarray], fp16_in: list[np.ndarray]) -> bool:\n assert fp32_in[0].shape == fp16_in[0].shape, \"Input shapes are not the same\"\n # print mean difference between fp32 and fp16 inputs\n if np.abs(fp32_in[0] - fp16_in[0]).mean() > 1e-3:\n print(\n f\"Mean difference between fp32 and fp16 inputs: {np.abs(fp32_in[0] - fp16_in[0]).mean()} \"\n + \"-> YOU MAY EXPECT DIFFERING RESULTS\"\n )\n return True # NOTE: Only warning, not error\n\n\ndef main(args):\n model_float32 = _load_model(args.arch, model_path=args.input_model if args.input_model else None)\n size = (1, *model_float32.cfg[\"input_shape\"])\n\n img_tensor = np.random.rand(*size).astype(np.float32)\n\n with TemporaryDirectory() as temp_dir:\n model_fp16_path = f\"{temp_dir}/model_fp16.onnx\"\n input_feed = {model_float32.runtime_inputs.name: img_tensor}\n model_float16 = auto_convert_mixed_precision(\n # NOTE: keep_io_types=True is required to keep the input/output type as float32\n onnx.load(str(model_float32.model_path)),\n input_feed,\n validate_fn=_validate,\n keep_io_types=True,\n )\n onnx.save(model_float16, model_fp16_path)\n model_fp16 = _load_model(args.arch, model_fp16_path)\n\n # Latency check\n _latency_check(args, size, model_float32, img_tensor)\n _latency_check(args, size, model_fp16, img_tensor)\n\n onnx.save(model_float16, args.arch + \"_fp16.onnx\")\n print(f\"FP16 model saved at {args.arch}_fp16.onnx\")\n print(\"Attention: FP16 converted models can only run on GPU devices.\")\n\n\nif __name__ == \"__main__\":\n parser = argparse.ArgumentParser(\n description=\"OnnxTR FP32 to FP16 conversion\",\n formatter_class=argparse.ArgumentDefaultsHelpFormatter,\n )\n parser.add_argument(\n \"arch\",\n type=str,\n choices=DETECTION_ARCHS + RECOGNITION_ARCHS + ORIENTATION_ARCHS,\n help=\"Architecture to convert\",\n )\n parser.add_argument(\"--input_model\", type=str, help=\"Path to the input model\", required=False)\n parser.add_argument(\"--it\", type=int, default=1000, help=\"Number of iterations to run\")\n args = parser.parse_args()\n\n main(args)\n"
},
{
"path": "scripts/evaluate.py",
"content": "# Copyright (C) 2021-2026, Mindee | Felix Dittrich.\n\n# This program is licensed under the Apache License 2.0.\n# See LICENSE or go to for full license details.\n\ntry:\n from doctr.version import __version__\n\n print(f\"DocTR version: {__version__}\")\nexcept ImportError:\n raise ImportError(\"Failed to import `doctr`. Please install `pip install python-doctr[torch]`.\")\n\nimport os\nimport time\nfrom typing import Any\n\nimport numpy as np\nfrom doctr import datasets\nfrom doctr import transforms as T\nfrom doctr.utils.metrics import LocalizationConfusion, OCRMetric, TextMatch\nfrom tqdm import tqdm\n\nfrom onnxtr.models import EngineConfig, ocr_predictor\nfrom onnxtr.utils.geometry import extract_crops, extract_rcrops\n\n\ndef _pct(val):\n return \"N/A\" if val is None else f\"{val:.2%}\"\n\n\ndef main(args):\n if not args.rotation:\n args.eval_straight = True\n\n if args.profiling:\n os.environ[\"ONNXTR_MULTIPROCESSING_DISABLE\"] = \"TRUE\"\n try:\n import memray\n import yappi\n except ImportError:\n raise ImportError(\"Please install yappi and memray to enable profiling - `pip install yappi memray`.\")\n yappi.set_clock_type(\"cpu\")\n # Drop memray profile and flamegraph if they already exist\n if os.path.exists(\"memray_profile.bin\"):\n os.remove(\"memray_profile.bin\")\n if os.path.exists(\"memray_flamegraph.html\"):\n os.remove(\"memray_flamegraph.html\")\n memray_tracker = memray.Tracker(\"memray_profile.bin\")\n memray_tracker.__enter__()\n\n input_shape = (args.size, args.size)\n\n # We define a transformation function which does transform the annotation\n # to the required format for the Resize transformation\n def _transform(img, target):\n boxes = target[\"boxes\"]\n transformed_img, transformed_boxes = T.Resize(\n input_shape, preserve_aspect_ratio=args.keep_ratio, symmetric_pad=args.symmetric_pad\n )(img, boxes)\n return transformed_img, {\"boxes\": transformed_boxes, \"labels\": target[\"labels\"]}\n\n predictor = ocr_predictor(\n args.detection,\n args.recognition,\n reco_bs=args.batch_size,\n preserve_aspect_ratio=False, # we handle the transformation directly in the dataset so this is set to False\n symmetric_pad=False, # we handle the transformation directly in the dataset so this is set to False\n assume_straight_pages=not args.rotation,\n load_in_8_bit=args.load_8bit,\n det_engine_cfg=EngineConfig(providers=[\"CPUExecutionProvider\"]) if args.force_cpu else None,\n reco_engine_cfg=EngineConfig(providers=[\"CPUExecutionProvider\"]) if args.force_cpu else None,\n clf_engine_cfg=EngineConfig(providers=[\"CPUExecutionProvider\"]) if args.force_cpu else None,\n )\n\n # Load the dataset\n train_set = datasets.__dict__[args.dataset](\n train=True,\n download=True,\n use_polygons=not args.eval_straight,\n sample_transforms=_transform,\n )\n val_set = datasets.__dict__[args.dataset](\n train=False,\n download=True,\n use_polygons=not args.eval_straight,\n sample_transforms=_transform,\n )\n sets = [train_set, val_set]\n\n reco_metric = TextMatch()\n\n det_metric = LocalizationConfusion(iou_thresh=args.iou, use_polygons=not args.eval_straight)\n e2e_metric = OCRMetric(iou_thresh=args.iou, use_polygons=not args.eval_straight)\n\n sample_idx = 0\n extraction_fn = extract_crops if args.eval_straight else extract_rcrops\n\n timings = []\n\n # Warmup\n print(\"Warming up the model...\")\n dummy_img = np.zeros((args.size, args.size, 3), dtype=np.uint8)\n for _ in range(5):\n _ = predictor([dummy_img])\n print(\"Warmup done.\\n\")\n\n for dataset in sets:\n for page, target in tqdm(dataset):\n if hasattr(page, \"numpy\"):\n page = page.numpy()\n\n if page.ndim == 3 and page.shape[0] in [1, 3]:\n page = np.moveaxis(page, 0, -1)\n\n if page.dtype != np.uint8:\n page = (page * 255).astype(np.uint8) if np.max(page) <= 1 else page.astype(np.uint8)\n\n # GT\n gt_boxes = target[\"boxes\"]\n gt_labels = target[\"labels\"]\n\n # Forward\n if args.profiling:\n yappi.start()\n start_ts = time.perf_counter()\n out = predictor(page[None, ...])\n timings.append(time.perf_counter() - start_ts)\n if args.profiling:\n yappi.stop()\n\n crops = extraction_fn(page, gt_boxes, channels_last=True)\n reco_out = predictor.reco_predictor(crops)\n\n reco_words: Any = []\n if len(reco_out):\n reco_words, _ = zip(*reco_out)\n\n # Unpack preds\n pred_boxes: list[list[Any]] = []\n pred_labels: list[str] = []\n for page in out.pages:\n height, width = page.dimensions\n for block in page.blocks:\n for line in block.lines:\n for word in line.words:\n if not args.rotation:\n (a, b), (c, d) = word.geometry\n else:\n (\n [x1, y1],\n [x2, y2],\n [x3, y3],\n [x4, y4],\n ) = word.geometry\n if np.issubdtype(gt_boxes.dtype, np.integer):\n if not args.rotation:\n pred_boxes.append([\n int(a * width),\n int(b * height),\n int(c * width),\n int(d * height),\n ])\n else:\n if args.eval_straight:\n pred_boxes.append([\n int(width * min(x1, x2, x3, x4)),\n int(height * min(y1, y2, y3, y4)),\n int(width * max(x1, x2, x3, x4)),\n int(height * max(y1, y2, y3, y4)),\n ])\n else:\n pred_boxes.append([\n [int(x1 * width), int(y1 * height)],\n [int(x2 * width), int(y2 * height)],\n [int(x3 * width), int(y3 * height)],\n [int(x4 * width), int(y4 * height)],\n ])\n else:\n if not args.rotation:\n pred_boxes.append([a, b, c, d])\n else:\n if args.eval_straight:\n pred_boxes.append([\n min(x1, x2, x3, x4),\n min(y1, y2, y3, y4),\n max(x1, x2, x3, x4),\n max(y1, y2, y3, y4),\n ])\n else:\n pred_boxes.append([[x1, y1], [x2, y2], [x3, y3], [x4, y4]])\n pred_labels.append(word.value)\n\n # Update the metric\n det_metric.update(gt_boxes, np.asarray(pred_boxes))\n reco_metric.update(gt_labels, reco_words)\n e2e_metric.update(gt_boxes, np.asarray(pred_boxes), gt_labels, pred_labels)\n\n # Loop break\n sample_idx += 1\n if isinstance(args.samples, int) and args.samples == sample_idx:\n break\n if isinstance(args.samples, int) and args.samples == sample_idx:\n break\n\n # Unpack aggregated metrics\n print(f\"Model Evaluation (model= {args.detection} + {args.recognition}, dataset={args.dataset})\")\n recall, precision, mean_iou = det_metric.summary()\n print(f\"Text Detection - Recall: {_pct(recall)}, Precision: {_pct(precision)}, Mean IoU: {_pct(mean_iou)}\")\n acc = reco_metric.summary()\n print(f\"Text Recognition - Accuracy: {_pct(acc['raw'])} (unicase: {_pct(acc['unicase'])})\")\n recall, precision, mean_iou = e2e_metric.summary()\n print(\n f\"OCR - Recall: {_pct(recall['raw'])} (unicase: {_pct(recall['unicase'])}), \"\n f\"Precision: {_pct(precision['raw'])} (unicase: {_pct(precision['unicase'])}), Mean IoU: {_pct(mean_iou)}\\n\"\n )\n print(f\"Number of samples: {sample_idx}\")\n print(f\"Total inference time: {np.sum(timings):.2f} sec\")\n print(f\"Average inference time per sample: {np.mean(timings):.6f} sec\")\n\n if args.profiling:\n import subprocess\n\n memray_tracker.__exit__(None, None, None)\n\n with open(\"yappi_profile.stats\", \"w\") as f:\n yappi.get_func_stats().print_all(out=f)\n\n print(\"Profiling complete. Generating memray flamegraph and stats...\")\n subprocess.run([\"memray\", \"flamegraph\", \"memray_profile.bin\", \"-o\", \"memray_flamegraph.html\"])\n subprocess.run([\"memray\", \"stats\", \"memray_profile.bin\"])\n\n\ndef parse_args():\n import argparse\n\n parser = argparse.ArgumentParser(\n description=\"OnnxTR end-to-end evaluation\", formatter_class=argparse.ArgumentDefaultsHelpFormatter\n )\n\n parser.add_argument(\"detection\", type=str, help=\"Text detection model to use for analysis\")\n parser.add_argument(\"recognition\", type=str, help=\"Text recognition model to use for analysis\")\n parser.add_argument(\"--iou\", type=float, default=0.5, help=\"IoU threshold to match a pair of boxes\")\n parser.add_argument(\"--dataset\", type=str, default=\"FUNSD\", help=\"choose a dataset: FUNSD, CORD\")\n parser.add_argument(\"--rotation\", dest=\"rotation\", action=\"store_true\", help=\"run rotated OCR + postprocessing\")\n parser.add_argument(\"-b\", \"--batch_size\", type=int, default=32, help=\"batch size for recognition\")\n parser.add_argument(\"--size\", type=int, default=1024, help=\"model input size, H = W\")\n parser.add_argument(\"--keep_ratio\", action=\"store_true\", help=\"keep the aspect ratio of the input image\")\n parser.add_argument(\"--symmetric_pad\", action=\"store_true\", help=\"pad the image symmetrically\")\n parser.add_argument(\"--samples\", type=int, default=None, help=\"evaluate only on the N first samples\")\n parser.add_argument(\n \"--eval-straight\",\n action=\"store_true\",\n help=\"evaluate on straight pages with straight bbox (to use the quick and light metric)\",\n )\n parser.add_argument(\"--load_8bit\", action=\"store_true\", help=\"load model in 8bit mode\")\n parser.add_argument(\"--force-cpu\", action=\"store_true\", help=\"force CPU execution\")\n parser.add_argument(\"--profiling\", action=\"store_true\", help=\"enable profiling\")\n args = parser.parse_args()\n\n return args\n\n\nif __name__ == \"__main__\":\n args = parse_args()\n main(args)\n"
},
{
"path": "scripts/latency.py",
"content": "# Copyright (C) 2021-2026, Mindee | Felix Dittrich.\n\n# This program is licensed under the Apache License 2.0.\n# See LICENSE or go to for full license details.\n\nimport argparse\nimport time\n\nimport numpy as np\n\nfrom onnxtr.models import classification, detection, recognition\nfrom onnxtr.models.classification.zoo import ORIENTATION_ARCHS\nfrom onnxtr.models.detection.zoo import ARCHS as DETECTION_ARCHS\nfrom onnxtr.models.recognition.zoo import ARCHS as RECOGNITION_ARCHS\n\n\ndef main(args):\n if args.arch in DETECTION_ARCHS:\n model = detection.__dict__[args.arch](load_in_8_bit=args.load8bit)\n elif args.arch in RECOGNITION_ARCHS:\n model = recognition.__dict__[args.arch](load_in_8_bit=args.load8bit)\n elif args.arch in ORIENTATION_ARCHS:\n model = classification.__dict__[args.arch](load_in_8_bit=args.load8bit)\n else:\n raise ValueError(f\"Unknown architecture {args.arch}\")\n\n size = (1, *model.cfg[\"input_shape\"])\n img_tensor = np.random.rand(*size).astype(np.float32)\n\n # Warmup\n for _ in range(10):\n _ = model(img_tensor)\n\n timings = []\n\n # Evaluation runs\n for _ in range(args.it):\n start_ts = time.perf_counter()\n _ = model(img_tensor)\n timings.append(time.perf_counter() - start_ts)\n\n _timings = np.array(timings)\n print(f\"{args.arch} ({args.it} runs on ({size}) inputs)\")\n print(f\"mean {1000 * _timings.mean():.2f}ms, std {1000 * _timings.std():.2f}ms\")\n\n\nif __name__ == \"__main__\":\n parser = argparse.ArgumentParser(\n description=\"OnnxTR latency benchmark\",\n formatter_class=argparse.ArgumentDefaultsHelpFormatter,\n )\n parser.add_argument(\n \"arch\",\n type=str,\n choices=DETECTION_ARCHS + RECOGNITION_ARCHS + ORIENTATION_ARCHS,\n help=\"Architecture to benchmark\",\n )\n parser.add_argument(\"--load8bit\", action=\"store_true\", help=\"Load the 8-bit quantized model\")\n parser.add_argument(\"--it\", type=int, default=1000, help=\"Number of iterations to run\")\n args = parser.parse_args()\n\n main(args)\n"
},
{
"path": "scripts/quantize.py",
"content": "import argparse\nimport os\nimport time\nfrom enum import Enum\n\nimport numpy as np\nimport onnxruntime\nfrom onnxruntime.quantization import CalibrationDataReader, QuantFormat, QuantType, quantize_dynamic, quantize_static\n\nfrom onnxtr.io.image import read_img_as_numpy\nfrom onnxtr.models.preprocessor import PreProcessor\nfrom onnxtr.utils.geometry import shape_translate\n\n\nclass TaskShapes(Enum):\n \"\"\"Enum class to define the shapes of the input tensors for different tasks\"\"\"\n\n crop_orientation = (256, 256)\n page_orientation = (512, 512)\n detection = (1024, 1024)\n recognition = (32, 128)\n\n\nclass CalibrationDataLoader(CalibrationDataReader):\n def __init__(self, calibration_image_folder: str, model_path: str, task_shape: tuple[int]):\n self.enum_data = None\n self.preprocessor = PreProcessor(output_size=task_shape, batch_size=1)\n self.dataset = [\n self.preprocessor(\n np.expand_dims(read_img_as_numpy(os.path.join(calibration_image_folder, img_file)), axis=0)\n )\n for img_file in os.listdir(calibration_image_folder)[:500] # limit to 500 images\n ]\n\n session = onnxruntime.InferenceSession(model_path, None)\n self.input_name = session.get_inputs()[0].name\n self.datasize = len(self.dataset)\n\n def get_next(self):\n if self.enum_data is None:\n self.enum_data = iter([\n {self.input_name: shape_translate(input_data[0], format=\"BCHW\")} for input_data in self.dataset\n ])\n return next(self.enum_data, None)\n\n def rewind(self):\n self.enum_data = None\n\n\ndef benchmark(calibration_image_folder: str, model_path: str, task_shape: tuple[int]):\n session = onnxruntime.InferenceSession(model_path)\n input_name = session.get_inputs()[0].name\n output_name = [output.name for output in session.get_outputs()]\n dataset = CalibrationDataLoader(calibration_image_folder, model_path, task_shape)\n sample = shape_translate(dataset.dataset[0][0], format=\"BCHW\") # take 1 sample for benchmarking\n\n total = 0.0\n runs = 10\n # Warming up\n _ = session.run(output_name, {input_name: sample})\n for _ in range(runs):\n start = time.perf_counter()\n _ = session.run(output_name, {input_name: sample})\n end = (time.perf_counter() - start) * 1000\n total += end\n print(f\"{end:.2f}ms\")\n total /= runs\n print(f\"Avg: {total:.2f}ms\")\n\n\ndef benchmark_mean_diff(\n calibration_image_folder: str, model_path: str, quantized_model_path: str, task_shape: tuple[int]\n):\n \"\"\"Check the mean difference between the original and quantized model\"\"\"\n session = onnxruntime.InferenceSession(model_path)\n quantized_session = onnxruntime.InferenceSession(quantized_model_path)\n input_name = session.get_inputs()[0].name\n output_name = [output.name for output in session.get_outputs()]\n quantized_output_name = [output.name for output in quantized_session.get_outputs()]\n dataset = CalibrationDataLoader(calibration_image_folder, model_path, task_shape)\n sample = shape_translate(dataset.dataset[0][0], format=\"BCHW\") # take 1 sample for benchmarking\n\n output = session.run(output_name, {input_name: sample})[0]\n quantized_output = quantized_session.run(quantized_output_name, {input_name: sample})[0]\n\n mean_diff = np.mean(np.abs(output - quantized_output))\n print(f\"Mean difference between original and quantized model: {mean_diff:.2f}\")\n\n\ndef main(args):\n input_model_path = args.input_model\n calibration_dataset_path = args.calibrate_dataset\n if args.task == \"crop_orientation\":\n task_shape = TaskShapes.crop_orientation.value\n elif args.task == \"page_orientation\":\n task_shape = TaskShapes.page_orientation.value\n elif args.task == \"detection\":\n task_shape = TaskShapes.detection.value\n else:\n task_shape = TaskShapes.recognition.value\n print(f\"Task: {args.task} | Task shape: {task_shape}\")\n\n dr = CalibrationDataLoader(calibration_dataset_path, input_model_path, task_shape)\n base_model_name = input_model_path.split(\"/\")[-1].split(\"-\")[0]\n static_out_name = base_model_name + \"_static_8_bit.onnx\"\n dynamic_out_name = base_model_name + \"_dynamic_8_bit.onnx\"\n\n print(\"benchmarking fp32 model...\")\n benchmark(calibration_dataset_path, input_model_path, task_shape)\n\n # Calibrate and quantize model\n # Turn off model optimization during quantization\n if \"parseq\" not in input_model_path: # Skip static quantization for Parseq\n print(\"Calibrating and quantizing model static...\")\n try:\n quantize_static(\n input_model_path,\n static_out_name,\n dr,\n quant_format=args.quant_format,\n weight_type=QuantType.QInt8,\n activation_type=QuantType.QUInt8,\n reduce_range=True,\n )\n except Exception:\n print(\"Error during static quantization --> Change weight_type also to QUInt8\")\n quantize_static(\n input_model_path,\n static_out_name,\n dr,\n quant_format=args.quant_format,\n weight_type=QuantType.QUInt8,\n activation_type=QuantType.QUInt8,\n reduce_range=True,\n )\n\n print(\"benchmarking static int8 model...\")\n benchmark(calibration_dataset_path, static_out_name, task_shape)\n\n print(\"benchmarking mean difference between fp32 and static int8 model...\")\n benchmark_mean_diff(calibration_dataset_path, input_model_path, static_out_name, task_shape)\n\n print(\"Calibrated and quantized static model saved.\")\n\n if \"sar\" not in input_model_path: # Skip dynamic quantization for SAR_ResNet31\n print(\"Dynamic int 8 quantization...\")\n quantize_dynamic(\n input_model_path,\n dynamic_out_name,\n weight_type=QuantType.QUInt8,\n )\n print(\"Dynamic model saved.\")\n\n print(\"benchmarking dynamic int8 model...\")\n benchmark(calibration_dataset_path, dynamic_out_name, task_shape)\n\n print(\"benchmarking mean difference between fp32 and dynamic int8 model...\")\n benchmark_mean_diff(calibration_dataset_path, input_model_path, dynamic_out_name, task_shape)\n\n\nif __name__ == \"__main__\":\n parser = argparse.ArgumentParser(\n description=\"OnnxTR script to quantize models and benchmark the quantized models\",\n formatter_class=argparse.ArgumentDefaultsHelpFormatter,\n )\n parser.add_argument(\"--input_model\", required=True, help=\"input model\")\n parser.add_argument(\n \"--task\",\n required=True,\n type=str,\n choices=[\"crop_orientation\", \"page_orientation\", \"detection\", \"recognition\"],\n help=\"task shape\",\n )\n parser.add_argument(\n \"--calibrate_dataset\",\n type=str,\n required=True,\n help=\"calibration data set (word crop images for recognition, crop_orientation else page images for detection, page_orientation)\", # noqa\n )\n parser.add_argument(\n \"--quant_format\",\n default=QuantFormat.QDQ,\n type=QuantFormat.from_string,\n choices=list(QuantFormat),\n )\n args = parser.parse_args()\n\n main(args)\n"
},
{
"path": "setup.py",
"content": "# Copyright (C) 2021-2026, Mindee | Felix Dittrich.\n\n# This program is licensed under the Apache License 2.0.\n# See LICENSE or go to for full license details.\n\nimport os\nfrom pathlib import Path\n\nfrom setuptools import setup\n\nPKG_NAME = \"onnxtr\"\nVERSION = os.getenv(\"BUILD_VERSION\", \"0.8.2a0\")\n\n\nif __name__ == \"__main__\":\n print(f\"Building wheel {PKG_NAME}-{VERSION}\")\n\n # Dynamically set the __version__ attribute\n cwd = Path(__file__).parent.absolute()\n with open(cwd.joinpath(\"onnxtr\", \"version.py\"), \"w\", encoding=\"utf-8\") as f:\n f.write(f\"__version__ = '{VERSION}'\\n\")\n\n setup(name=PKG_NAME, version=VERSION)\n"
},
{
"path": "tests/common/test_contrib.py",
"content": "import numpy as np\nimport pytest\n\nfrom onnxtr.contrib import artefacts\nfrom onnxtr.contrib.base import _BasePredictor\nfrom onnxtr.io import DocumentFile\n\n\ndef test_base_predictor():\n # check that we need to provide either a url or a model_path\n with pytest.raises(ValueError):\n _ = _BasePredictor(batch_size=2)\n\n predictor = _BasePredictor(batch_size=2, url=artefacts.default_cfgs[\"yolov8_artefact\"][\"url\"])\n # check that we need to implement preprocess and postprocess\n with pytest.raises(NotImplementedError):\n predictor.preprocess(np.zeros((10, 10, 3)))\n with pytest.raises(NotImplementedError):\n predictor.postprocess([np.zeros((10, 10, 3))], [[np.zeros((10, 10, 3))]])\n\n\ndef test_artefact_detector(mock_artefact_image_stream):\n doc = DocumentFile.from_images([mock_artefact_image_stream])\n detector = artefacts.ArtefactDetector(batch_size=2, conf_threshold=0.5, iou_threshold=0.5)\n results = detector(doc)\n assert isinstance(results, list) and len(results) == 1 and isinstance(results[0], list)\n assert all(isinstance(artefact, dict) for artefact in results[0])\n # check result keys\n assert all(key in results[0][0] for key in [\"label\", \"confidence\", \"box\"])\n assert all(len(artefact[\"box\"]) == 4 for artefact in results[0])\n assert all(isinstance(coord, int) for box in results[0] for coord in box[\"box\"])\n assert all(isinstance(artefact[\"confidence\"], float) for artefact in results[0])\n assert all(isinstance(artefact[\"label\"], str) for artefact in results[0])\n # check results for the mock image are 9 artefacts\n assert len(results[0]) == 9\n # test visualization non-blocking for tests\n detector.show(block=False)\n"
},
{
"path": "tests/common/test_core.py",
"content": "import pytest\n\nimport onnxtr\nfrom onnxtr.file_utils import requires_package\n\n\ndef test_version():\n assert len(onnxtr.__version__.split(\".\")) == 3\n\n\ndef test_requires_package():\n requires_package(\"numpy\") # availbable\n with pytest.raises(ImportError): # not available\n requires_package(\"non_existent_package\")\n"
},
{
"path": "tests/common/test_engine_cfg.py",
"content": "import gc\n\nimport numpy as np\nimport psutil\nimport pytest\nfrom onnxruntime import RunOptions, SessionOptions\n\nfrom onnxtr import models\nfrom onnxtr.io import Document\nfrom onnxtr.models import EngineConfig, detection, recognition\nfrom onnxtr.models.predictor import OCRPredictor\n\n\ndef _get_rss_mb():\n gc.collect()\n process = psutil.Process()\n return process.memory_info().rss / (1024 * 1024)\n\n\ndef _test_predictor(predictor):\n # Output checks\n assert isinstance(predictor, OCRPredictor)\n\n doc = [np.zeros((1024, 1024, 3), dtype=np.uint8)]\n out = predictor(doc)\n # Document\n assert isinstance(out, Document)\n\n # The input doc has 1 page\n assert len(out.pages) == 1\n # Dimension check\n with pytest.raises(ValueError):\n input_page = (255 * np.random.rand(1, 256, 512, 3)).astype(np.uint8)\n _ = predictor([input_page])\n\n\n@pytest.mark.parametrize(\n \"det_arch, reco_arch\",\n [[det_arch, reco_arch] for det_arch, reco_arch in zip(detection.zoo.ARCHS, recognition.zoo.ARCHS)],\n)\ndef test_engine_cfg(det_arch, reco_arch):\n session_options = SessionOptions()\n session_options.enable_cpu_mem_arena = False\n engine_cfg = EngineConfig(\n providers=[\"CPUExecutionProvider\"],\n session_options=session_options,\n )\n\n assert engine_cfg.__repr__() == \"EngineConfig(providers=['CPUExecutionProvider'])\"\n\n # Model\n predictor = models.ocr_predictor(\n det_arch, reco_arch, det_engine_cfg=engine_cfg, reco_engine_cfg=engine_cfg, clf_engine_cfg=engine_cfg\n )\n assert predictor.det_predictor.model.providers == [\"CPUExecutionProvider\"]\n assert not predictor.det_predictor.model.session_options.enable_cpu_mem_arena\n assert predictor.reco_predictor.model.providers == [\"CPUExecutionProvider\"]\n assert not predictor.reco_predictor.model.session_options.enable_cpu_mem_arena\n _test_predictor(predictor)\n\n # passing model instance directly\n det_model = detection.__dict__[det_arch](engine_cfg=engine_cfg)\n assert det_model.providers == [\"CPUExecutionProvider\"]\n assert not det_model.session_options.enable_cpu_mem_arena\n\n reco_model = recognition.__dict__[reco_arch](engine_cfg=engine_cfg)\n assert reco_model.providers == [\"CPUExecutionProvider\"]\n assert not reco_model.session_options.enable_cpu_mem_arena\n\n predictor = models.ocr_predictor(det_model, reco_model)\n assert predictor.det_predictor.model.providers == [\"CPUExecutionProvider\"]\n assert not predictor.det_predictor.model.session_options.enable_cpu_mem_arena\n assert predictor.reco_predictor.model.providers == [\"CPUExecutionProvider\"]\n assert not predictor.reco_predictor.model.session_options.enable_cpu_mem_arena\n _test_predictor(predictor)\n\n det_predictor = models.detection_predictor(det_arch, engine_cfg=engine_cfg)\n assert det_predictor.model.providers == [\"CPUExecutionProvider\"]\n assert not det_predictor.model.session_options.enable_cpu_mem_arena\n\n reco_predictor = models.recognition_predictor(reco_arch, engine_cfg=engine_cfg)\n assert reco_predictor.model.providers == [\"CPUExecutionProvider\"]\n assert not reco_predictor.model.session_options.enable_cpu_mem_arena\n\n\ndef test_cpu_memory_arena_shrinkage_enabled():\n session_options = SessionOptions()\n session_options.enable_mem_pattern = False\n session_options.enable_cpu_mem_arena = True\n\n enable_shrinkage = False\n\n providers = [(\"CPUExecutionProvider\", {\"arena_extend_strategy\": \"kSameAsRequested\"})]\n\n def enable_arena_shrinkage(run_options: \"RunOptions\") -> \"RunOptions\":\n if enable_shrinkage:\n run_options.add_run_config_entry(\"memory.enable_memory_arena_shrinkage\", \"cpu:0\")\n assert run_options.get_run_config_entry(\"memory.enable_memory_arena_shrinkage\") == \"cpu:0\"\n return run_options\n\n engine_cfg = EngineConfig(\n providers=providers,\n session_options=session_options,\n run_options_provider=enable_arena_shrinkage,\n )\n\n predictor = models.ocr_predictor(\n det_engine_cfg=engine_cfg,\n reco_engine_cfg=engine_cfg,\n clf_engine_cfg=engine_cfg,\n detect_orientation=True,\n )\n\n assert predictor.det_predictor.model.providers == providers\n assert predictor.det_predictor.model.session_options.enable_cpu_mem_arena\n assert predictor.reco_predictor.model.providers == providers\n assert predictor.reco_predictor.model.session_options.enable_cpu_mem_arena\n\n rng = np.random.RandomState(seed=42)\n sample = rng.randint(0, 256, (1024, 1024, 3), dtype=np.uint8)\n\n start_rss = _get_rss_mb()\n\n predictor([sample])\n increased_rss = _get_rss_mb()\n\n assert increased_rss > start_rss\n\n enable_shrinkage = True\n\n predictor([sample])\n decreased_rss = _get_rss_mb()\n\n assert increased_rss > decreased_rss\n"
},
{
"path": "tests/common/test_headers.py",
"content": "\"\"\"Test for python files copyright headers.\"\"\"\n\nfrom datetime import datetime\nfrom pathlib import Path\n\n\ndef test_copyright_header():\n copyright_header = \"\".join([\n f\"# Copyright (C) {2021}-{datetime.now().year}, Mindee | Felix Dittrich.\\n\\n\",\n \"# This program is licensed under the Apache License 2.0.\\n\",\n \"# See LICENSE or go to for full license details.\\n\",\n ])\n excluded_files = [\"__init__.py\", \"version.py\"]\n invalid_files = []\n locations = [\".github\", \"onnxtr\"]\n\n for location in locations:\n for source_path in Path(__file__).parent.parent.parent.joinpath(location).rglob(\"*.py\"):\n if source_path.name not in excluded_files:\n source_path_content = source_path.read_text()\n if copyright_header not in source_path_content:\n invalid_files.append(source_path)\n assert len(invalid_files) == 0, f\"Invalid copyright header in the following files: {invalid_files}\"\n"
},
{
"path": "tests/common/test_io.py",
"content": "from io import BytesIO\nfrom pathlib import Path\n\nimport numpy as np\nimport pytest\nimport requests\n\nfrom onnxtr import io\n\n\ndef _check_doc_content(doc_tensors, num_pages):\n # 1 doc of 8 pages\n assert len(doc_tensors) == num_pages\n assert all(isinstance(page, np.ndarray) for page in doc_tensors)\n assert all(page.dtype == np.uint8 for page in doc_tensors)\n\n\ndef test_read_pdf(mock_pdf):\n doc = io.read_pdf(mock_pdf)\n _check_doc_content(doc, 2)\n\n # Test with Path\n doc = io.read_pdf(Path(mock_pdf))\n _check_doc_content(doc, 2)\n\n with open(mock_pdf, \"rb\") as f:\n doc = io.read_pdf(f.read())\n _check_doc_content(doc, 2)\n\n # Wrong input type\n with pytest.raises(TypeError):\n _ = io.read_pdf(123)\n\n # Wrong path\n with pytest.raises(FileNotFoundError):\n _ = io.read_pdf(\"my_imaginary_file.pdf\")\n\n\ndef test_read_img_as_numpy(tmpdir_factory, mock_pdf):\n # Wrong input type\n with pytest.raises(TypeError):\n _ = io.read_img_as_numpy(123)\n\n # Non-existing file\n with pytest.raises(FileNotFoundError):\n io.read_img_as_numpy(\"my_imaginary_file.jpg\")\n\n # Invalid image\n with pytest.raises(ValueError):\n io.read_img_as_numpy(str(mock_pdf))\n\n # From path\n url = \"https://doctr-static.mindee.com/models?id=v0.2.1/Grace_Hopper.jpg&src=0\"\n file = BytesIO(requests.get(url).content)\n tmp_path = str(tmpdir_factory.mktemp(\"data\").join(\"mock_img_file.jpg\"))\n with open(tmp_path, \"wb\") as f:\n f.write(file.getbuffer())\n\n # Path & stream\n with open(tmp_path, \"rb\") as f:\n page_stream = io.read_img_as_numpy(f.read())\n\n for page in (io.read_img_as_numpy(tmp_path), page_stream):\n # Data type\n assert isinstance(page, np.ndarray)\n assert page.dtype == np.uint8\n # Shape\n assert page.shape == (606, 517, 3)\n\n # RGB\n bgr_page = io.read_img_as_numpy(tmp_path, rgb_output=False)\n assert np.all(page == bgr_page[..., ::-1])\n\n # Resize\n target_size = (200, 150)\n resized_page = io.read_img_as_numpy(tmp_path, target_size)\n assert resized_page.shape[:2] == target_size\n\n\ndef test_read_html():\n url = \"https://www.google.com\"\n pdf_stream = io.read_html(url)\n assert isinstance(pdf_stream, bytes)\n\n\ndef test_document_file(mock_pdf, mock_artefact_image_stream):\n pages = io.DocumentFile.from_images([mock_artefact_image_stream])\n _check_doc_content(pages, 1)\n\n assert isinstance(io.DocumentFile.from_pdf(mock_pdf), list)\n assert isinstance(io.DocumentFile.from_url(\"https://www.google.com\"), list)\n\n\ndef test_pdf(mock_pdf):\n pages = io.DocumentFile.from_pdf(mock_pdf)\n\n # As images\n num_pages = 2\n _check_doc_content(pages, num_pages)\n"
},
{
"path": "tests/common/test_io_elements.py",
"content": "from xml.etree.ElementTree import ElementTree\n\nimport numpy as np\nimport pytest\n\nfrom onnxtr.io import elements\n\n\ndef _mock_words(size=(1.0, 1.0), offset=(0, 0), confidence=0.9, objectness_score=0.9, polygons=False):\n box_word_elements = [\n elements.Word(\n \"hello\",\n confidence,\n ((offset[0], offset[1]), (size[0] / 2 + offset[0], size[1] / 2 + offset[1])),\n objectness_score,\n {\"value\": 0, \"confidence\": None},\n ),\n elements.Word(\n \"world\",\n confidence,\n ((size[0] / 2 + offset[0], size[1] / 2 + offset[1]), (size[0] + offset[0], size[1] + offset[1])),\n objectness_score,\n {\"value\": 0, \"confidence\": None},\n ),\n ]\n polygons_word_elements = [\n elements.Word(\n \"hello\",\n confidence,\n # (x1, y1), (x2, y2), (x3, y3), (x4, y4) with shape (4, 2)\n np.array([\n [offset[0], offset[1]],\n [size[0] / 2 + offset[0], offset[1]],\n [size[0] / 2 + offset[0], size[1] / 2 + offset[1]],\n [offset[0], size[1] / 2 + offset[1]],\n ]),\n objectness_score,\n {\"value\": 0, \"confidence\": None},\n ),\n elements.Word(\n \"world\",\n confidence,\n # (x1, y1), (x2, y2), (x3, y3), (x4, y4) with shape (4, 2)\n np.array([\n [size[0] / 2 + offset[0], size[1] / 2 + offset[1]],\n [size[0] + offset[0], size[1] / 2 + offset[1]],\n [size[0] + offset[0], size[1] + offset[1]],\n [size[0] / 2 + offset[0], size[1] + offset[1]],\n ]),\n objectness_score,\n {\"value\": 0, \"confidence\": None},\n ),\n ]\n return polygons_word_elements if polygons else box_word_elements\n\n\ndef _mock_artefacts(size=(1, 1), offset=(0, 0), confidence=0.8):\n sub_size = (size[0] / 2, size[1] / 2)\n return [\n elements.Artefact(\n \"qr_code\", confidence, ((offset[0], offset[1]), (sub_size[0] + offset[0], sub_size[1] + offset[1]))\n ),\n elements.Artefact(\n \"qr_code\",\n confidence,\n ((sub_size[0] + offset[0], sub_size[1] + offset[1]), (size[0] + offset[0], size[1] + offset[1])),\n ),\n ]\n\n\ndef _mock_lines(size=(1, 1), offset=(0, 0), polygons=False):\n sub_size = (size[0] / 2, size[1] / 2)\n return [\n elements.Line(_mock_words(size=sub_size, offset=offset, polygons=polygons)),\n elements.Line(\n _mock_words(size=sub_size, offset=(offset[0] + sub_size[0], offset[1] + sub_size[1]), polygons=polygons)\n ),\n ]\n\n\ndef _mock_blocks(size=(1, 1), offset=(0, 0), polygons=False):\n sub_size = (size[0] / 4, size[1] / 4)\n return [\n elements.Block(\n _mock_lines(size=sub_size, offset=offset, polygons=polygons),\n _mock_artefacts(size=sub_size, offset=(offset[0] + sub_size[0], offset[1] + sub_size[1])),\n ),\n elements.Block(\n _mock_lines(\n size=sub_size, offset=(offset[0] + 2 * sub_size[0], offset[1] + 2 * sub_size[1]), polygons=polygons\n ),\n _mock_artefacts(size=sub_size, offset=(offset[0] + 3 * sub_size[0], offset[1] + 3 * sub_size[1])),\n ),\n ]\n\n\ndef _mock_pages(block_size=(1, 1), block_offset=(0, 0), polygons=False):\n return [\n elements.Page(\n np.random.randint(0, 255, (300, 200, 3), dtype=np.uint8),\n _mock_blocks(block_size, block_offset, polygons),\n 0,\n (300, 200),\n {\"value\": 0.0, \"confidence\": 1.0},\n {\"value\": \"EN\", \"confidence\": 0.8},\n ),\n elements.Page(\n np.random.randint(0, 255, (500, 1000, 3), dtype=np.uint8),\n _mock_blocks(block_size, block_offset),\n 1,\n (500, 1000),\n {\"value\": 0.15, \"confidence\": 0.8},\n {\"value\": \"FR\", \"confidence\": 0.7},\n ),\n ]\n\n\ndef test_element():\n with pytest.raises(KeyError):\n elements.Element(sub_elements=[1])\n\n\ndef test_word():\n word_str = \"hello\"\n conf = 0.8\n geom = ((0, 0), (1, 1))\n objectness_score = 0.9\n crop_orientation = {\"value\": 0, \"confidence\": None}\n word = elements.Word(word_str, conf, geom, objectness_score, crop_orientation)\n\n # Attribute checks\n assert word.value == word_str\n assert word.confidence == conf\n assert word.geometry == geom\n assert word.objectness_score == objectness_score\n assert word.crop_orientation == crop_orientation\n\n # Render\n assert word.render() == word_str\n\n # Export\n assert word.export() == {\n \"value\": word_str,\n \"confidence\": conf,\n \"geometry\": geom,\n \"objectness_score\": objectness_score,\n \"crop_orientation\": crop_orientation,\n }\n\n # Repr\n assert word.__repr__() == f\"Word(value='hello', confidence={conf:.2})\"\n\n # Class method\n state_dict = {\n \"value\": \"there\",\n \"confidence\": 0.1,\n \"geometry\": ((0, 0), (0.5, 0.5)),\n \"objectness_score\": objectness_score,\n \"crop_orientation\": crop_orientation,\n }\n word = elements.Word.from_dict(state_dict)\n assert word.export() == state_dict\n\n\ndef test_line():\n geom = ((0, 0), (0.5, 0.5))\n objectness_score = 0.9\n words = _mock_words(size=geom[1], offset=geom[0])\n line = elements.Line(words)\n\n # Attribute checks\n assert len(line.words) == len(words)\n assert all(isinstance(w, elements.Word) for w in line.words)\n assert line.geometry == geom\n assert line.objectness_score == objectness_score\n\n # Render\n assert line.render() == \"hello world\"\n\n # Export\n assert line.export() == {\n \"words\": [w.export() for w in words],\n \"geometry\": geom,\n \"objectness_score\": objectness_score,\n }\n\n # Repr\n words_str = \" \" * 4 + \",\\n \".join(repr(word) for word in words) + \",\"\n assert line.__repr__() == f\"Line(\\n (words): [\\n{words_str}\\n ]\\n)\"\n\n # Ensure that words repr does't span on several lines when there are none\n assert repr(elements.Line([], ((0, 0), (1, 1)))) == \"Line(\\n (words): []\\n)\"\n\n # from dict\n state_dict = {\n \"words\": [\n {\n \"value\": \"there\",\n \"confidence\": 0.1,\n \"geometry\": ((0, 0), (1.0, 1.0)),\n \"objectness_score\": objectness_score,\n \"crop_orientation\": {\"value\": 0, \"confidence\": None},\n }\n ],\n \"geometry\": ((0, 0), (1.0, 1.0)),\n \"objectness_score\": objectness_score,\n }\n line = elements.Line.from_dict(state_dict)\n assert line.export() == state_dict\n\n\ndef test_artefact():\n artefact_type = \"qr_code\"\n conf = 0.8\n geom = ((0, 0), (1, 1))\n artefact = elements.Artefact(artefact_type, conf, geom)\n\n # Attribute checks\n assert artefact.type == artefact_type\n assert artefact.confidence == conf\n assert artefact.geometry == geom\n\n # Render\n assert artefact.render() == \"[QR_CODE]\"\n\n # Export\n assert artefact.export() == {\"type\": artefact_type, \"confidence\": conf, \"geometry\": geom}\n\n # Repr\n assert artefact.__repr__() == f\"Artefact(type='{artefact_type}', confidence={conf:.2})\"\n\n\ndef test_block():\n geom = ((0, 0), (1, 1))\n sub_size = (geom[1][0] / 2, geom[1][0] / 2)\n objectness_score = 0.9\n lines = _mock_lines(size=sub_size, offset=geom[0])\n artefacts = _mock_artefacts(size=sub_size, offset=sub_size)\n block = elements.Block(lines, artefacts)\n\n # Attribute checks\n assert len(block.lines) == len(lines)\n assert len(block.artefacts) == len(artefacts)\n assert all(isinstance(w, elements.Line) for w in block.lines)\n assert all(isinstance(a, elements.Artefact) for a in block.artefacts)\n assert block.geometry == geom\n\n # Render\n assert block.render() == \"hello world\\nhello world\"\n\n # Export\n assert block.export() == {\n \"lines\": [line.export() for line in lines],\n \"artefacts\": [artefact.export() for artefact in artefacts],\n \"geometry\": geom,\n \"objectness_score\": objectness_score,\n }\n\n\ndef test_page():\n page = np.zeros((300, 200, 3), dtype=np.uint8)\n page_idx = 0\n page_size = (300, 200)\n orientation = {\"value\": 0.0, \"confidence\": 0.0}\n language = {\"value\": \"EN\", \"confidence\": 0.8}\n blocks = _mock_blocks()\n page = elements.Page(page, blocks, page_idx, page_size, orientation, language)\n\n # Attribute checks\n assert len(page.blocks) == len(blocks)\n assert all(isinstance(b, elements.Block) for b in page.blocks)\n assert isinstance(page.page, np.ndarray)\n assert page.page_idx == page_idx\n assert page.dimensions == page_size\n assert page.orientation == orientation\n assert page.language == language\n\n # Render\n assert page.render() == \"hello world\\nhello world\\n\\nhello world\\nhello world\"\n\n # Export\n assert page.export() == {\n \"blocks\": [b.export() for b in blocks],\n \"page_idx\": page_idx,\n \"dimensions\": page_size,\n \"orientation\": orientation,\n \"language\": language,\n }\n\n # Export XML\n assert (\n isinstance(page.export_as_xml(), tuple)\n and isinstance(page.export_as_xml()[0], (bytes, bytearray))\n and isinstance(page.export_as_xml()[1], ElementTree)\n )\n\n # Repr\n assert \"\\n\".join(repr(page).split(\"\\n\")[:2]) == f\"Page(\\n dimensions={page_size!r}\"\n\n # Show\n page.show(block=False)\n\n # Synthesize\n img = page.synthesize()\n assert isinstance(img, np.ndarray)\n assert img.shape == (*page_size, 3)\n\n\ndef test_document():\n pages = _mock_pages()\n doc = elements.Document(pages)\n\n # Attribute checks\n assert len(doc.pages) == len(pages)\n assert all(isinstance(p, elements.Page) for p in doc.pages)\n\n # Render\n page_export = \"hello world\\nhello world\\n\\nhello world\\nhello world\"\n assert doc.render() == f\"{page_export}\\n\\n\\n\\n{page_export}\"\n\n # Export\n assert doc.export() == {\"pages\": [p.export() for p in pages]}\n\n # Export XML\n xml_output = doc.export_as_xml()\n assert isinstance(xml_output, list) and len(xml_output) == len(pages)\n # Check that the XML is well-formed in hOCR format\n for xml_bytes, xml_tree in xml_output:\n assert isinstance(xml_bytes, bytes)\n assert isinstance(xml_tree, ElementTree)\n root = xml_tree.getroot()\n assert root.tag == \"html\"\n assert root[0].tag == \"head\"\n assert root[1].tag == \"body\"\n assert root[1][0].tag == \"div\" and root[1][0].attrib[\"class\"] == \"ocr_page\"\n for block in root[1][0]:\n assert block.tag == \"div\" and block.attrib[\"class\"] == \"ocr_carea\"\n assert block[0].tag == \"p\" and block[0].attrib[\"class\"] == \"ocr_par\"\n for line in block[0]:\n assert line.tag == \"span\" and line.attrib[\"class\"] == \"ocr_line\"\n for word in line:\n assert word.tag == \"span\" and word.attrib[\"class\"] == \"ocrx_word\"\n\n # Show\n doc.show(block=False)\n\n # Synthesize\n img_list = doc.synthesize()\n assert isinstance(img_list, list) and len(img_list) == len(pages)\n"
},
{
"path": "tests/common/test_models.py",
"content": "from io import BytesIO\n\nimport cv2\nimport numpy as np\nimport pytest\nimport requests\n\nfrom onnxtr.io import reader\nfrom onnxtr.models._utils import estimate_orientation, get_language\nfrom onnxtr.utils import geometry\n\n\n@pytest.fixture(scope=\"function\")\ndef mock_image(tmpdir_factory):\n url = \"https://doctr-static.mindee.com/models?id=v0.2.1/bitmap30.png&src=0\"\n file = BytesIO(requests.get(url).content)\n tmp_path = str(tmpdir_factory.mktemp(\"data\").join(\"mock_bitmap.jpg\"))\n with open(tmp_path, \"wb\") as f:\n f.write(file.getbuffer())\n image = reader.read_img_as_numpy(tmp_path)\n return image\n\n\n@pytest.fixture(scope=\"function\")\ndef mock_bitmap(mock_image):\n bitmap = np.squeeze(cv2.cvtColor(mock_image, cv2.COLOR_BGR2GRAY) / 255.0)\n bitmap = np.expand_dims(bitmap, axis=-1)\n return bitmap\n\n\ndef test_estimate_orientation(mock_image, mock_bitmap, mock_tilted_payslip):\n assert estimate_orientation(mock_image * 0) == 0\n\n # test binarized image\n angle = estimate_orientation(mock_bitmap)\n assert abs(angle) - 30 < 1.0\n\n angle = estimate_orientation(mock_bitmap * 255)\n assert abs(angle) - 30.0 < 1.0\n\n angle = estimate_orientation(mock_image)\n assert abs(angle) - 30.0 < 1.0\n\n rotated = geometry.rotate_image(mock_image, angle)\n angle_rotated = estimate_orientation(rotated)\n assert abs(angle_rotated) == 0\n\n mock_tilted_payslip = reader.read_img_as_numpy(mock_tilted_payslip)\n assert estimate_orientation(mock_tilted_payslip) == -30\n\n rotated = geometry.rotate_image(mock_tilted_payslip, -30, expand=True)\n angle_rotated = estimate_orientation(rotated)\n assert abs(angle_rotated) < 1.0\n with pytest.raises(AssertionError):\n estimate_orientation(np.ones((10, 10, 10)))\n\n # test with general_page_orientation\n assert estimate_orientation(mock_bitmap, (90, 0.9)) in range(140, 160)\n\n rotated = geometry.rotate_image(mock_tilted_payslip, -30)\n assert estimate_orientation(rotated, (0, 0.9)) in range(-10, 10)\n\n assert estimate_orientation(mock_image, (0, 0.9)) - 30 < 1.0\n\n # Aspect Ratio Independence (Portrait vs Landscape)\n # Pad the tilted image to be very tall (Portrait)\n portrait_img = cv2.copyMakeBorder(mock_tilted_payslip, 500, 500, 0, 0, cv2.BORDER_CONSTANT, value=[0, 0, 0])\n # Pad the tilted image to be very wide (Landscape)\n landscape_img = cv2.copyMakeBorder(mock_tilted_payslip, 0, 0, 500, 500, cv2.BORDER_CONSTANT, value=[0, 0, 0])\n\n assert abs(estimate_orientation(portrait_img) - (-30)) <= 1.0\n assert abs(estimate_orientation(landscape_img) - (-30)) <= 1.0\n\n # Perpendicular Noise Test\n vertical_noise = np.zeros((1000, 1000, 3), dtype=np.uint8)\n cv2.line(vertical_noise, (500, 100), (500, 900), (255, 255, 255), 10)\n assert estimate_orientation(vertical_noise) == 0\n\n\ndef test_get_lang():\n sentence = \"This is a test sentence.\"\n expected_lang = \"en\"\n threshold_prob = 0.99\n\n lang = get_language(sentence)\n\n assert lang[0] == expected_lang\n assert lang[1] > threshold_prob\n\n lang = get_language(\"a\")\n assert lang[0] == \"unknown\"\n assert lang[1] == 0.0\n"
},
{
"path": "tests/common/test_models_builder.py",
"content": "import numpy as np\nimport pytest\n\nfrom onnxtr.io import Document\nfrom onnxtr.models import builder\n\nwords_per_page = 10\n\n\ndef test_documentbuilder():\n num_pages = 2\n\n # Don't resolve lines\n doc_builder = builder.DocumentBuilder(resolve_lines=False, resolve_blocks=False)\n pages = [np.zeros((100, 200, 3))] * num_pages\n boxes = np.random.rand(words_per_page, 6) # array format\n boxes[:2] *= boxes[2:4]\n objectness_scores = np.array([0.9] * words_per_page)\n # Arg consistency check\n with pytest.raises(ValueError):\n doc_builder(\n pages,\n [boxes, boxes],\n [objectness_scores, objectness_scores],\n [(\"hello\", 1.0)] * 3,\n [(100, 200), (100, 200)],\n [{\"value\": 0, \"confidence\": None}] * 3,\n )\n out = doc_builder(\n pages,\n [boxes, boxes],\n [objectness_scores, objectness_scores],\n [[(\"hello\", 1.0)] * words_per_page] * num_pages,\n [(100, 200), (100, 200)],\n [[{\"value\": 0, \"confidence\": None}] * words_per_page] * num_pages,\n )\n assert isinstance(out, Document)\n assert len(out.pages) == num_pages\n assert all(isinstance(page.page, np.ndarray) for page in out.pages) and all(\n page.page.shape == (100, 200, 3) for page in out.pages\n )\n # 1 Block & 1 line per page\n assert len(out.pages[0].blocks) == 1 and len(out.pages[0].blocks[0].lines) == 1\n assert len(out.pages[0].blocks[0].lines[0].words) == words_per_page\n\n # Resolve lines\n doc_builder = builder.DocumentBuilder(resolve_lines=True, resolve_blocks=True)\n out = doc_builder(\n pages,\n [boxes, boxes],\n [objectness_scores, objectness_scores],\n [[(\"hello\", 1.0)] * words_per_page] * num_pages,\n [(100, 200), (100, 200)],\n [[{\"value\": 0, \"confidence\": None}] * words_per_page] * num_pages,\n )\n\n # No detection\n boxes = np.zeros((0, 4))\n objectness_scores = np.zeros([0])\n out = doc_builder(\n pages, [boxes, boxes], [objectness_scores, objectness_scores], [[], []], [(100, 200), (100, 200)], [[]]\n )\n assert len(out.pages[0].blocks) == 0\n\n # Rotated boxes to export as straight boxes\n boxes = np.array([\n [[0.1, 0.1], [0.2, 0.2], [0.15, 0.25], [0.05, 0.15]],\n [[0.5, 0.5], [0.6, 0.6], [0.55, 0.65], [0.45, 0.55]],\n ])\n objectness_scores = np.array([0.99, 0.99])\n doc_builder_2 = builder.DocumentBuilder(resolve_blocks=False, resolve_lines=False, export_as_straight_boxes=True)\n out = doc_builder_2(\n [np.zeros((100, 100, 3))],\n [boxes],\n [objectness_scores],\n [[(\"hello\", 0.99), (\"word\", 0.99)]],\n [(100, 100)],\n [[{\"value\": 0, \"confidence\": None}] * 2],\n )\n assert out.pages[0].blocks[0].lines[0].words[-1].geometry == ((0.45, 0.5), (0.6, 0.65))\n assert out.pages[0].blocks[0].lines[0].words[-1].objectness_score == 0.99\n\n # Repr\n assert (\n repr(doc_builder) == \"DocumentBuilder(resolve_lines=True, \"\n \"resolve_blocks=True, paragraph_break=0.035, export_as_straight_boxes=False)\"\n )\n\n\n@pytest.mark.parametrize(\n \"input_boxes, sorted_idxs\",\n [\n [[[0, 0.5, 0.1, 0.6], [0, 0.3, 0.2, 0.4], [0, 0, 0.1, 0.1]], [2, 1, 0]], # vertical\n [[[0.7, 0.5, 0.85, 0.6], [0.2, 0.3, 0.4, 0.4], [0, 0, 0.1, 0.1]], [2, 1, 0]], # diagonal\n [[[0, 0.5, 0.1, 0.6], [0.15, 0.5, 0.25, 0.6], [0.5, 0.5, 0.6, 0.6]], [0, 1, 2]], # same line, 2p\n [[[0, 0.5, 0.1, 0.6], [0.2, 0.49, 0.35, 0.59], [0.8, 0.52, 0.9, 0.63]], [0, 1, 2]], # ~same line\n [[[0, 0.3, 0.4, 0.45], [0.5, 0.28, 0.75, 0.42], [0, 0.45, 0.1, 0.55]], [0, 1, 2]], # 2 lines\n [[[0, 0.3, 0.4, 0.35], [0.75, 0.28, 0.95, 0.42], [0, 0.45, 0.1, 0.55]], [0, 1, 2]], # 2 lines\n [\n [\n [[0.1, 0.1], [0.2, 0.2], [0.15, 0.25], [0.05, 0.15]],\n [[0.5, 0.5], [0.6, 0.6], [0.55, 0.65], [0.45, 0.55]],\n ],\n [0, 1],\n ], # rot\n ],\n)\ndef test_sort_boxes(input_boxes, sorted_idxs):\n doc_builder = builder.DocumentBuilder()\n assert doc_builder._sort_boxes(np.asarray(input_boxes))[0].tolist() == sorted_idxs\n\n\n@pytest.mark.parametrize(\n \"input_boxes, lines\",\n [\n [[[0, 0.5, 0.1, 0.6], [0, 0.3, 0.2, 0.4], [0, 0, 0.1, 0.1]], [[2], [1], [0]]], # vertical\n [[[0.7, 0.5, 0.85, 0.6], [0.2, 0.3, 0.4, 0.4], [0, 0, 0.1, 0.1]], [[2], [1], [0]]], # diagonal\n [[[0, 0.5, 0.14, 0.6], [0.15, 0.5, 0.25, 0.6], [0.5, 0.5, 0.6, 0.6]], [[0, 1], [2]]], # same line, 2p\n [[[0, 0.5, 0.18, 0.6], [0.2, 0.48, 0.35, 0.58], [0.8, 0.52, 0.9, 0.63]], [[0, 1], [2]]], # ~same line\n [[[0, 0.3, 0.48, 0.45], [0.5, 0.28, 0.75, 0.42], [0, 0.45, 0.1, 0.55]], [[0, 1], [2]]], # 2 lines\n [[[0, 0.3, 0.4, 0.35], [0.75, 0.28, 0.95, 0.42], [0, 0.45, 0.1, 0.55]], [[0], [1], [2]]], # 2 lines\n [\n [\n [[0.1, 0.1], [0.2, 0.2], [0.15, 0.25], [0.05, 0.15]],\n [[0.5, 0.5], [0.6, 0.6], [0.55, 0.65], [0.45, 0.55]],\n ],\n [[0], [1]],\n ], # rot\n ],\n)\ndef test_resolve_lines(input_boxes, lines):\n doc_builder = builder.DocumentBuilder()\n assert doc_builder._resolve_lines(np.asarray(input_boxes)) == lines\n"
},
{
"path": "tests/common/test_models_classification.py",
"content": "import cv2\nimport numpy as np\nimport pytest\n\nfrom onnxtr.models import classification, detection\nfrom onnxtr.models.classification.predictor import OrientationPredictor\nfrom onnxtr.models.engine import Engine\n\n\n@pytest.mark.parametrize(\n \"arch_name, input_shape\",\n [\n [\"mobilenet_v3_small_crop_orientation\", (256, 256, 3)],\n [\"mobilenet_v3_small_page_orientation\", (512, 512, 3)],\n ],\n)\ndef test_classification_models(arch_name, input_shape):\n batch_size = 8\n model = classification.__dict__[arch_name]()\n assert isinstance(model, Engine)\n input_tensor = np.random.rand(batch_size, *input_shape).astype(np.float32)\n out = model(input_tensor)\n assert isinstance(out, np.ndarray)\n assert out.shape == (8, 4)\n\n\n@pytest.mark.parametrize(\n \"arch_name\",\n [\n \"mobilenet_v3_small_crop_orientation\",\n \"mobilenet_v3_small_page_orientation\",\n ],\n)\ndef test_classification_zoo(arch_name):\n if \"crop\" in arch_name:\n batch_size = 16\n input_array = np.random.rand(batch_size, 3, 256, 256).astype(np.float32)\n # Model\n predictor = classification.zoo.crop_orientation_predictor(arch_name)\n\n with pytest.raises(ValueError):\n predictor = classification.zoo.crop_orientation_predictor(arch=\"wrong_model\")\n else:\n batch_size = 2\n input_array = np.random.rand(batch_size, 3, 512, 512).astype(np.float32)\n # Model\n predictor = classification.zoo.page_orientation_predictor(arch_name)\n\n with pytest.raises(ValueError):\n predictor = classification.zoo.page_orientation_predictor(arch=\"wrong_model\")\n # object check\n assert isinstance(predictor, OrientationPredictor)\n\n out = predictor(input_array)\n class_idxs, classes, confs = out[0], out[1], out[2]\n assert isinstance(class_idxs, list) and len(class_idxs) == batch_size\n assert isinstance(classes, list) and len(classes) == batch_size\n assert isinstance(confs, list) and len(confs) == batch_size\n assert all(isinstance(pred, int) for pred in class_idxs)\n assert all(isinstance(pred, int) for pred in classes) and all(pred in [0, 90, 180, -90] for pred in classes)\n assert all(isinstance(pred, float) for pred in confs)\n\n\n@pytest.mark.parametrize(\"quantized\", [False, True])\ndef test_crop_orientation_model(mock_text_box, quantized):\n text_box_0 = cv2.imread(mock_text_box)\n # rotates counter-clockwise\n text_box_270 = np.rot90(text_box_0, 1)\n text_box_180 = np.rot90(text_box_0, 2)\n text_box_90 = np.rot90(text_box_0, 3)\n classifier = classification.crop_orientation_predictor(\n \"mobilenet_v3_small_crop_orientation\", load_in_8_bit=quantized\n )\n assert classifier([text_box_0, text_box_270, text_box_180, text_box_90])[0] == [0, 1, 2, 3]\n # 270 degrees is equivalent to -90 degrees\n assert classifier([text_box_0, text_box_270, text_box_180, text_box_90])[1] == [0, -90, 180, 90]\n assert all(isinstance(pred, float) for pred in classifier([text_box_0, text_box_270, text_box_180, text_box_90])[2])\n\n # Test custom model loading\n classifier = classification.crop_orientation_predictor(\n classification.mobilenet_v3_small_crop_orientation(load_in_8_bit=quantized)\n )\n assert isinstance(classifier, OrientationPredictor)\n\n with pytest.raises(ValueError):\n _ = classification.crop_orientation_predictor(detection.db_resnet34())\n\n # Test with disabled predictor\n classifier = classification.crop_orientation_predictor(\"mobilenet_v3_small_crop_orientation\", disabled=True)\n assert classifier([text_box_0, text_box_270, text_box_180, text_box_90]) == [\n [0, 0, 0, 0],\n [0, 0, 0, 0],\n [1.0, 1.0, 1.0, 1.0],\n ]\n\n\n@pytest.mark.parametrize(\"quantized\", [False, True])\ndef test_page_orientation_model(mock_payslip, quantized):\n text_box_0 = cv2.imread(mock_payslip)\n # rotates counter-clockwise\n text_box_270 = np.rot90(text_box_0, 1)\n text_box_180 = np.rot90(text_box_0, 2)\n text_box_90 = np.rot90(text_box_0, 3)\n classifier = classification.crop_orientation_predictor(\n \"mobilenet_v3_small_page_orientation\", load_in_8_bit=quantized\n )\n assert classifier([text_box_0, text_box_270, text_box_180, text_box_90])[0] == [0, 1, 2, 3]\n # 270 degrees is equivalent to -90 degrees\n assert classifier([text_box_0, text_box_270, text_box_180, text_box_90])[1] == [0, -90, 180, 90]\n assert all(isinstance(pred, float) for pred in classifier([text_box_0, text_box_270, text_box_180, text_box_90])[2])\n\n # Test custom model loading\n classifier = classification.page_orientation_predictor(\n classification.mobilenet_v3_small_page_orientation(load_in_8_bit=quantized)\n )\n assert isinstance(classifier, OrientationPredictor)\n\n with pytest.raises(ValueError):\n _ = classification.page_orientation_predictor(detection.db_resnet34())\n\n # Test with disabled predictor\n classifier = classification.crop_orientation_predictor(\"mobilenet_v3_small_page_orientation\", disabled=True)\n assert classifier([text_box_0, text_box_270, text_box_180, text_box_90]) == [\n [0, 0, 0, 0],\n [0, 0, 0, 0],\n [1.0, 1.0, 1.0, 1.0],\n ]\n"
},
{
"path": "tests/common/test_models_detection.py",
"content": "import numpy as np\nimport pytest\n\nfrom onnxtr.models import detection\nfrom onnxtr.models.detection.postprocessor.base import GeneralDetectionPostProcessor\nfrom onnxtr.models.detection.predictor import DetectionPredictor\nfrom onnxtr.models.engine import Engine\n\n\ndef test_postprocessor():\n postprocessor = GeneralDetectionPostProcessor(assume_straight_pages=True)\n r_postprocessor = GeneralDetectionPostProcessor(assume_straight_pages=False)\n with pytest.raises(AssertionError):\n postprocessor(np.random.rand(2, 512, 512).astype(np.float32))\n mock_batch = np.random.rand(2, 512, 512, 1).astype(np.float32)\n out = postprocessor(mock_batch)\n r_out = r_postprocessor(mock_batch)\n # Batch composition\n assert isinstance(out, list)\n assert len(out) == 2\n assert all(isinstance(sample, list) and all(isinstance(v, np.ndarray) for v in sample) for sample in out)\n assert all(all(v.shape[1] == 5 for v in sample) for sample in out)\n assert all(all(v.shape[1] == 5 and v.shape[2] == 2 for v in sample) for sample in r_out)\n # Relative coords\n assert all(all(np.all(np.logical_and(v[:, :4] >= 0, v[:, :4] <= 1)) for v in sample) for sample in out)\n assert all(all(np.all(np.logical_and(v[:, :4] >= 0, v[:, :4] <= 1)) for v in sample) for sample in r_out)\n # Repr\n assert repr(postprocessor) == \"GeneralDetectionPostProcessor(bin_thresh=0.1, box_thresh=0.1)\"\n # Edge case when the expanded points of the polygon has two lists\n issue_points = np.array(\n [\n [869, 561],\n [923, 581],\n [925, 595],\n [915, 583],\n [889, 583],\n [905, 593],\n [882, 601],\n [901, 595],\n [904, 604],\n [876, 608],\n [915, 614],\n [911, 605],\n [925, 601],\n [930, 616],\n [911, 617],\n [900, 636],\n [931, 637],\n [904, 649],\n [932, 649],\n [932, 628],\n [918, 627],\n [934, 624],\n [935, 573],\n [909, 569],\n [934, 562],\n ],\n dtype=np.int32,\n )\n out = postprocessor.polygon_to_box(issue_points)\n r_out = r_postprocessor.polygon_to_box(issue_points)\n assert isinstance(out, tuple) and len(out) == 4\n assert isinstance(r_out, np.ndarray) and r_out.shape == (4, 2)\n\n\n@pytest.mark.parametrize(\"quantized\", [False, True])\n@pytest.mark.parametrize(\n \"arch_name, input_shape, output_size, out_prob\",\n [\n [\"db_resnet34\", (1024, 1024, 3), (1024, 1024, 1), True],\n [\"db_resnet50\", (1024, 1024, 3), (1024, 1024, 1), True],\n [\"db_mobilenet_v3_large\", (1024, 1024, 3), (1024, 1024, 1), True],\n [\"linknet_resnet18\", (1024, 1024, 3), (1024, 1024, 1), True],\n [\"linknet_resnet34\", (1024, 1024, 3), (1024, 1024, 1), True],\n [\"linknet_resnet50\", (1024, 1024, 3), (1024, 1024, 1), True],\n [\"fast_tiny\", (1024, 1024, 3), (1024, 1024, 1), True],\n [\"fast_small\", (1024, 1024, 3), (1024, 1024, 1), True],\n [\"fast_base\", (1024, 1024, 3), (1024, 1024, 1), True],\n ],\n)\ndef test_detection_models(arch_name, input_shape, output_size, out_prob, quantized):\n batch_size = 2\n model = detection.__dict__[arch_name](load_in_8_bit=quantized)\n assert isinstance(model, Engine)\n input_array = np.random.rand(batch_size, *input_shape).astype(np.float32)\n out = model(input_array, return_model_output=True)\n assert isinstance(out, dict)\n assert len(out) == 2\n # Check proba map\n assert out[\"out_map\"].shape == (batch_size, *output_size)\n assert out[\"out_map\"].dtype == np.float32\n if out_prob:\n assert np.all(out[\"out_map\"] >= 0) and np.all(out[\"out_map\"] <= 1)\n # Check boxes\n for boxes_list in out[\"preds\"]:\n for boxes in boxes_list:\n assert boxes.shape[1] == 5\n assert np.all(boxes[:, :2] < boxes[:, 2:4])\n assert np.all(boxes[:, :4] >= 0) and np.all(boxes[:, :4] <= 1)\n\n\n@pytest.mark.parametrize(\"quantized\", [False, True])\n@pytest.mark.parametrize(\n \"arch_name\",\n [\n \"db_resnet34\",\n \"db_resnet50\",\n \"db_mobilenet_v3_large\",\n \"linknet_resnet18\",\n \"linknet_resnet34\",\n \"linknet_resnet50\",\n \"fast_tiny\",\n \"fast_small\",\n \"fast_base\",\n ],\n)\ndef test_detection_zoo(arch_name, quantized):\n # Model\n predictor = detection.zoo.detection_predictor(\n arch_name, load_in_8_bit=quantized, preserve_aspect_ratio=False, symmetric_pad=False\n )\n # object check\n assert isinstance(predictor, DetectionPredictor)\n input_array = np.random.rand(2, 3, 1024, 1024).astype(np.float32)\n\n out, seq_maps = predictor(input_array, return_maps=True)\n assert isinstance(out, list)\n for box in out:\n assert isinstance(box, np.ndarray)\n assert box.shape[1] == 5\n assert np.all(box[:, :2] < box[:, 2:4])\n assert np.all(box[:, :4] >= 0) and np.all(box[:, :4] <= 1)\n assert all(isinstance(seq_map, np.ndarray) for seq_map in seq_maps)\n assert all(seq_map.shape[:2] == (1024, 1024) for seq_map in seq_maps)\n # check that all values in the seq_maps are between 0 and 1\n assert all((seq_map >= 0).all() and (seq_map <= 1).all() for seq_map in seq_maps)\n"
},
{
"path": "tests/common/test_models_detection_utils.py",
"content": "import numpy as np\nimport pytest\n\nfrom onnxtr.models.detection._utils import _remove_padding\n\n\n@pytest.mark.parametrize(\"pages\", [[np.zeros((1000, 1000))], [np.zeros((1000, 2000))], [np.zeros((2000, 1000))]])\n@pytest.mark.parametrize(\"preserve_aspect_ratio\", [True, False])\n@pytest.mark.parametrize(\"symmetric_pad\", [True, False])\n@pytest.mark.parametrize(\"assume_straight_pages\", [True, False])\ndef test_remove_padding(pages, preserve_aspect_ratio, symmetric_pad, assume_straight_pages):\n h, w = pages[0].shape\n # straight pages test cases\n if assume_straight_pages:\n loc_preds = [np.array([[0.7, 0.1, 0.7, 0.2]])]\n if h == w or not preserve_aspect_ratio:\n expected = loc_preds\n else:\n if symmetric_pad:\n if h > w:\n expected = [np.array([[0.9, 0.1, 0.9, 0.2]])]\n else:\n expected = [np.array([[0.7, 0.0, 0.7, 0.0]])]\n else:\n if h > w:\n expected = [np.array([[1.0, 0.1, 1.0, 0.2]])]\n else:\n expected = [np.array([[0.7, 0.2, 0.7, 0.4]])]\n # non-straight pages test cases\n else:\n loc_preds = [np.array([[[0.9, 0.1], [0.9, 0.2], [0.8, 0.2], [0.8, 0.2]]])]\n if h == w or not preserve_aspect_ratio:\n expected = loc_preds\n else:\n if symmetric_pad:\n if h > w:\n expected = [np.array([[[1.0, 0.1], [1.0, 0.2], [1.0, 0.2], [1.0, 0.2]]])]\n else:\n expected = [np.array([[[0.9, 0.0], [0.9, 0.0], [0.8, 0.0], [0.8, 0.0]]])]\n else:\n if h > w:\n expected = [np.array([[[1.0, 0.1], [1.0, 0.2], [1.0, 0.2], [1.0, 0.2]]])]\n else:\n expected = [np.array([[[0.9, 0.2], [0.9, 0.4], [0.8, 0.4], [0.8, 0.4]]])]\n\n result = _remove_padding(pages, loc_preds, preserve_aspect_ratio, symmetric_pad, assume_straight_pages)\n for res, exp in zip(result, expected):\n assert np.allclose(res, exp)\n"
},
{
"path": "tests/common/test_models_factory.py",
"content": "import json\nimport os\nimport tempfile\n\nimport pytest\n\nfrom onnxtr import models\nfrom onnxtr.models.factory import _save_model_and_config_for_hf_hub, from_hub, push_to_hf_hub\n\nAVAILABLE_ARCHS = {\n \"classification\": models.classification.zoo.ORIENTATION_ARCHS,\n \"detection\": models.detection.zoo.ARCHS,\n \"recognition\": models.recognition.zoo.ARCHS,\n}\n\n\ndef test_push_to_hf_hub():\n model = models.classification.mobilenet_v3_small_crop_orientation()\n with pytest.raises(ValueError):\n # run_config and/or arch must be specified\n push_to_hf_hub(model, model_name=\"test\", task=\"classification\")\n with pytest.raises(ValueError):\n # task must be one of classification, detection, recognition, obj_detection\n push_to_hf_hub(model, model_name=\"test\", task=\"invalid_task\", arch=\"mobilenet_v3_small\")\n with pytest.raises(ValueError):\n # arch not in available architectures for task\n push_to_hf_hub(model, model_name=\"test\", task=\"detection\", arch=\"crnn_mobilenet_v3_large\")\n\n\ndef test_models_huggingface_hub(tmpdir):\n with tempfile.TemporaryDirectory() as tmp_dir:\n for task_name, archs in AVAILABLE_ARCHS.items():\n for arch_name in archs:\n model = models.__dict__[task_name].__dict__[arch_name]()\n\n _save_model_and_config_for_hf_hub(model, arch=arch_name, task=task_name, save_dir=tmp_dir)\n\n assert hasattr(model, \"cfg\")\n assert len(os.listdir(tmp_dir)) == 2\n assert os.path.exists(tmp_dir + \"/model.onnx\")\n assert os.path.exists(tmp_dir + \"/config.json\")\n tmp_config = json.load(open(tmp_dir + \"/config.json\"))\n assert arch_name == tmp_config[\"arch\"]\n assert task_name == tmp_config[\"task\"]\n assert all(key in model.cfg.keys() for key in tmp_config.keys())\n\n # test from hub\n hub_model = from_hub(repo_id=\"Felix92/onnxtr-{}\".format(arch_name).replace(\"_\", \"-\"))\n assert isinstance(hub_model, type(model))\n"
},
{
"path": "tests/common/test_models_preprocessor.py",
"content": "import numpy as np\nimport pytest\n\nfrom onnxtr.models.preprocessor import PreProcessor\n\n\n@pytest.mark.parametrize(\n \"batch_size, output_size, input_tensor, expected_batches, expected_value\",\n [\n [2, (128, 128), np.full((3, 256, 128, 3), 255, dtype=np.uint8), 1, 0.5], # numpy uint8\n [2, (128, 128), np.ones((3, 256, 128, 3), dtype=np.float32), 1, 0.5], # numpy fp32\n [2, (128, 128), [np.full((256, 128, 3), 255, dtype=np.uint8)] * 3, 2, 0.5], # list of numpy uint8\n [2, (128, 128), [np.ones((256, 128, 3), dtype=np.float32)] * 3, 2, 0.5], # list of numpy fp32 list of tf fp32\n ],\n)\ndef test_preprocessor(batch_size, output_size, input_tensor, expected_batches, expected_value):\n processor = PreProcessor(output_size, batch_size)\n\n # Invalid input type\n with pytest.raises(TypeError):\n processor(42)\n # 4D check\n with pytest.raises(AssertionError):\n processor(np.full((256, 128, 3), 255, dtype=np.uint8))\n with pytest.raises(TypeError):\n processor(np.full((1, 256, 128, 3), 255, dtype=np.int32))\n # 3D check\n with pytest.raises(AssertionError):\n processor([np.full((3, 256, 128, 3), 255, dtype=np.uint8)])\n with pytest.raises(TypeError):\n processor([np.full((256, 128, 3), 255, dtype=np.int32)])\n\n out = processor(input_tensor)\n assert isinstance(out, list) and len(out) == expected_batches\n assert all(isinstance(b, np.ndarray) for b in out)\n assert all(b.dtype == np.float32 for b in out)\n assert all(b.shape[1:3] == output_size for b in out)\n assert all(np.all(b == expected_value) for b in out)\n assert len(repr(processor).split(\"\\n\")) == 4\n"
},
{
"path": "tests/common/test_models_recognition.py",
"content": "import numpy as np\nimport pytest\n\nfrom onnxtr.models import recognition\nfrom onnxtr.models.engine import Engine\nfrom onnxtr.models.recognition.core import RecognitionPostProcessor\nfrom onnxtr.models.recognition.predictor import RecognitionPredictor\nfrom onnxtr.models.recognition.predictor._utils import remap_preds, split_crops\nfrom onnxtr.utils.vocabs import VOCABS\n\n\ndef test_recognition_postprocessor():\n mock_vocab = VOCABS[\"french\"]\n post_processor = RecognitionPostProcessor(mock_vocab)\n assert post_processor.extra_repr() == f\"vocab_size={len(mock_vocab)}\"\n assert post_processor.vocab == mock_vocab\n assert post_processor._embedding == list(mock_vocab) + [\"\"]\n\n\n@pytest.mark.parametrize(\n \"crops, max_ratio, target_ratio, target_overlap_ratio, channels_last, num_crops\",\n [\n # No split required\n [[np.zeros((32, 128, 3), dtype=np.uint8)], 8, 4, 0.5, True, 1],\n [[np.zeros((3, 32, 128), dtype=np.uint8)], 8, 4, 0.5, False, 1],\n # Split required\n [[np.zeros((32, 1024, 3), dtype=np.uint8)], 8, 6, 0.5, True, 10],\n [[np.zeros((3, 32, 1024), dtype=np.uint8)], 8, 6, 0.5, False, 10],\n ],\n)\ndef test_split_crops(crops, max_ratio, target_ratio, target_overlap_ratio, channels_last, num_crops):\n new_crops, crop_map, should_remap = split_crops(crops, max_ratio, target_ratio, target_overlap_ratio, channels_last)\n assert len(new_crops) == num_crops\n assert len(crop_map) == len(crops)\n assert should_remap == (len(crops) != len(new_crops))\n\n\n@pytest.mark.parametrize(\n \"preds, crop_map, split_overlap_ratio, pred\",\n [\n # Nothing to remap\n ([(\"hello\", 0.5)], [0], 0.5, [(\"hello\", 0.5)]),\n # Merge\n ([(\"hellowo\", 0.5), (\"loworld\", 0.6)], [(0, 2, 0.5)], 0.5, [(\"helloworld\", 0.55)]),\n ],\n)\ndef test_remap_preds(preds, crop_map, split_overlap_ratio, pred):\n preds = remap_preds(preds, crop_map, split_overlap_ratio)\n assert len(preds) == len(pred)\n assert preds == pred\n assert all(isinstance(pred, tuple) for pred in preds)\n assert all(isinstance(pred[0], str) and isinstance(pred[1], float) for pred in preds)\n\n\n@pytest.mark.parametrize(\n \"inputs, max_ratio, target_ratio, target_overlap_ratio, expected_remap_required, expected_len, expected_shape, \"\n \"expected_crop_map, channels_last\",\n [\n # Don't split\n ([np.zeros((32, 32 * 4, 3))], 4, 4, 0.5, False, 1, (32, 128, 3), 0, True),\n # Split needed\n ([np.zeros((32, 32 * 4 + 1, 3))], 4, 4, 0.5, True, 2, (32, 128, 3), (0, 2, 0.9921875), True),\n # Larger max ratio prevents split\n ([np.zeros((32, 32 * 8, 3))], 8, 4, 0.5, False, 1, (32, 256, 3), 0, True),\n # Half-overlap, two crops\n ([np.zeros((32, 128 + 64, 3))], 4, 4, 0.5, True, 2, (32, 128, 3), (0, 2, 0.5), True),\n # Half-overlap, two crops, channels first\n ([np.zeros((3, 32, 128 + 64))], 4, 4, 0.5, True, 2, (3, 32, 128), (0, 2, 0.5), False),\n # Half-overlap with small max_ratio forces split\n ([np.zeros((32, 128 + 64, 3))], 2, 4, 0.5, True, 2, (32, 128, 3), (0, 2, 0.5), True),\n # > half last overlap ratio\n ([np.zeros((32, 128 + 32, 3))], 4, 4, 0.5, True, 2, (32, 128, 3), (0, 2, 0.75), True),\n # 3 crops, half last overlap\n ([np.zeros((32, 128 + 128, 3))], 4, 4, 0.5, True, 3, (32, 128, 3), (0, 3, 0.5), True),\n # 3 crops, > half last overlap\n ([np.zeros((32, 128 + 64 + 32, 3))], 4, 4, 0.5, True, 3, (32, 128, 3), (0, 3, 0.75), True),\n # Split into larger crops\n ([np.zeros((32, 192 * 2, 3))], 4, 6, 0.5, True, 3, (32, 192, 3), (0, 3, 0.5), True),\n # Test fallback for empty splits\n ([np.empty((1, 0, 3))], -1, 4, 0.5, False, 1, (1, 0, 3), (0), True),\n ],\n)\ndef test_split_crops_cases(\n inputs,\n max_ratio,\n target_ratio,\n target_overlap_ratio,\n expected_remap_required,\n expected_len,\n expected_shape,\n expected_crop_map,\n channels_last,\n):\n new_crops, crop_map, _remap_required = split_crops(\n inputs,\n max_ratio=max_ratio,\n target_ratio=target_ratio,\n split_overlap_ratio=target_overlap_ratio,\n channels_last=channels_last,\n )\n\n assert _remap_required == expected_remap_required\n assert len(new_crops) == expected_len\n assert len(crop_map) == 1\n\n if expected_remap_required:\n assert isinstance(crop_map[0], tuple)\n\n assert crop_map[0] == expected_crop_map\n\n for crop in new_crops:\n assert crop.shape == expected_shape\n\n\n@pytest.mark.parametrize(\n \"split_overlap_ratio\",\n [\n # lower bound\n 0.0,\n # upper bound\n 1.0,\n ],\n)\ndef test_invalid_split_overlap_ratio(split_overlap_ratio):\n with pytest.raises(ValueError):\n split_crops(\n [np.zeros((32, 32 * 4, 3))],\n max_ratio=4,\n target_ratio=4,\n split_overlap_ratio=split_overlap_ratio,\n )\n\n\n@pytest.mark.parametrize(\"quantized\", [False, True])\n@pytest.mark.parametrize(\n \"arch_name, input_shape\",\n [\n [\"crnn_vgg16_bn\", (32, 128, 3)],\n [\"crnn_mobilenet_v3_small\", (32, 128, 3)],\n [\"crnn_mobilenet_v3_large\", (32, 128, 3)],\n [\"sar_resnet31\", (32, 128, 3)],\n [\"master\", (32, 128, 3)],\n [\"vitstr_small\", (32, 128, 3)],\n [\"vitstr_base\", (32, 128, 3)],\n [\"parseq\", (32, 128, 3)],\n [\"viptr_tiny\", (32, 128, 3)],\n ],\n)\ndef test_recognition_models(arch_name, input_shape, quantized):\n mock_vocab = VOCABS[\"french\"]\n batch_size = 4\n model = recognition.__dict__[arch_name](load_in_8_bit=quantized)\n assert isinstance(model, Engine)\n input_array = np.random.rand(batch_size, *input_shape).astype(np.float32)\n\n out = model(input_array, return_model_output=True)\n assert isinstance(out, dict)\n assert len(out) == 2\n assert isinstance(out[\"preds\"], list)\n assert len(out[\"preds\"]) == batch_size\n assert all(isinstance(word, str) and isinstance(conf, float) and 0 <= conf <= 1 for word, conf in out[\"preds\"])\n\n assert isinstance(out[\"out_map\"], np.ndarray)\n assert out[\"out_map\"].shape[0] == 4\n\n # test model post processor\n post_processor = model.postprocessor\n decoded = post_processor(np.random.rand(2, len(mock_vocab), 30).astype(np.float32))\n assert isinstance(decoded, list)\n assert all(isinstance(word, str) and isinstance(conf, float) and 0 <= conf <= 1 for word, conf in decoded)\n assert len(decoded) == 2\n assert all(char in mock_vocab for word, _ in decoded for char in word)\n\n # Testing with a fixed batch size\n model = recognition.__dict__[arch_name]()\n model.fixed_batch_size = 1\n assert isinstance(model, Engine)\n input_array = np.random.rand(batch_size, *input_shape).astype(np.float32)\n\n out = model(input_array, return_model_output=True)\n assert isinstance(out, dict)\n assert len(out) == 2\n assert isinstance(out[\"preds\"], list)\n assert len(out[\"preds\"]) == batch_size\n assert all(isinstance(word, str) and isinstance(conf, float) and 0 <= conf <= 1 for word, conf in out[\"preds\"])\n\n assert isinstance(out[\"out_map\"], np.ndarray)\n assert out[\"out_map\"].shape[0] == 4\n\n\n@pytest.mark.parametrize(\"quantized\", [False, True])\n@pytest.mark.parametrize(\n \"input_shape\",\n [\n (128, 128, 3),\n (32, 1024, 3), # test case split wide crops\n ],\n)\n@pytest.mark.parametrize(\n \"arch_name\",\n [\n \"crnn_vgg16_bn\",\n \"crnn_mobilenet_v3_small\",\n \"crnn_mobilenet_v3_large\",\n \"sar_resnet31\",\n \"master\",\n \"vitstr_small\",\n \"vitstr_base\",\n \"parseq\",\n \"viptr_tiny\",\n ],\n)\ndef test_recognition_zoo(arch_name, input_shape, quantized):\n batch_size = 2\n # Model\n predictor = recognition.zoo.recognition_predictor(arch_name, load_in_8_bit=quantized)\n # object check\n assert isinstance(predictor, RecognitionPredictor)\n input_array = np.random.rand(batch_size, *input_shape).astype(np.float32)\n out = predictor(input_array)\n assert isinstance(out, list) and len(out) == batch_size\n assert all(isinstance(word, str) and isinstance(conf, float) for word, conf in out)\n\n with pytest.raises(ValueError):\n _ = recognition.zoo.recognition_predictor(arch=\"wrong_model\")\n"
},
{
"path": "tests/common/test_models_recognition_utils.py",
"content": "import pytest\n\nfrom onnxtr.models.recognition.utils import merge_multi_strings, merge_strings\n\n\n@pytest.mark.parametrize(\n \"a, b, overlap_ratio, merged\",\n [\n # Last character of first string and first of last string will be cropped when merging - indicated by X\n (\"abcX\", \"Xdef\", 0.5, \"abcdef\"),\n (\"abcdX\", \"Xdef\", 0.75, \"abcdef\"),\n (\"abcdeX\", \"Xdef\", 0.9, \"abcdef\"),\n (\"abcdefX\", \"Xdef\", 0.9, \"abcdef\"),\n # Long repetition - four of seven characters in the second string are in the estimated overlap\n # X-chars will be cropped during merge, because they might be cut off during splitting of corresponding image\n (\"abccccX\", \"Xcccccc\", 4 / 7, \"abcccccccc\"),\n (\"abc\", \"\", 0.5, \"abc\"),\n (\"\", \"abc\", 0.5, \"abc\"),\n (\"a\", \"b\", 0.5, \"ab\"),\n # No overlap of input strings after crop\n (\"abcdX\", \"Xefghi\", 0.33, \"abcdefghi\"),\n # No overlap of input strings after crop with shorter inputs\n (\"bcdX\", \"Xefgh\", 0.4, \"bcdefgh\"),\n # No overlap of input strings after crop with even shorter inputs\n (\"cdX\", \"Xefg\", 0.5, \"cdefg\"),\n # Full overlap of input strings\n (\"abcdX\", \"Xbcde\", 1.0, \"abcde\"),\n # One repetition within inputs\n (\"ababX\", \"Xabde\", 0.8, \"ababde\"),\n # Multiple repetitions within inputs\n (\"ababX\", \"Xabab\", 0.8, \"ababab\"),\n # Multiple repetitions within inputs with shorter input strings\n (\"abaX\", \"Xbab\", 1.0, \"abab\"),\n # Longer multiple repetitions within inputs with half overlap\n (\"cabababX\", \"Xabababc\", 0.5, \"cabababababc\"),\n # Longer multiple repetitions within inputs with full overlap\n (\"ababaX\", \"Xbabab\", 1.0, \"ababab\"),\n # One different letter in overlap\n (\"one_differon\", \"ferent_letter\", 0.5, \"one_differont_letter\"),\n # First string empty after crop\n (\"-\", \"test\", 0.9, \"-test\"),\n # Second string empty after crop\n (\"test\", \"-\", 0.9, \"test-\"),\n ],\n)\ndef test_merge_strings(a, b, overlap_ratio, merged):\n assert merged == merge_strings(a, b, overlap_ratio)\n\n\n@pytest.mark.parametrize(\n \"seq_list, overlap_ratio, last_overlap_ratio, merged\",\n [\n # One character at each conjunction point will be cropped when merging - indicated by X\n ([\"abcX\", \"Xdef\"], 0.5, 0.5, \"abcdef\"),\n ([\"abcdX\", \"XdefX\", \"XefghX\", \"Xijk\"], 0.5, 0.5, \"abcdefghijk\"),\n ([\"abcdX\", \"XdefX\", \"XefghiX\", \"Xaijk\"], 0.5, 0.8, \"abcdefghijk\"),\n ([\"aaaa\", \"aaab\", \"aabc\"], 0.8, 0.3, \"aaaabc\"),\n # Handle empty input\n ([], 0.5, 0.4, \"\"),\n ],\n)\ndef test_merge_multi_strings(seq_list, overlap_ratio, last_overlap_ratio, merged):\n assert merged == merge_multi_strings(seq_list, overlap_ratio, last_overlap_ratio)\n"
},
{
"path": "tests/common/test_models_zoo.py",
"content": "import numpy as np\nimport pytest\n\nfrom onnxtr import models\nfrom onnxtr.io import Document, DocumentFile\nfrom onnxtr.models import detection, recognition\nfrom onnxtr.models.classification import mobilenet_v3_small_crop_orientation, mobilenet_v3_small_page_orientation\nfrom onnxtr.models.classification.zoo import crop_orientation_predictor, page_orientation_predictor\nfrom onnxtr.models.detection.predictor import DetectionPredictor\nfrom onnxtr.models.detection.zoo import ARCHS as DET_ARCHS\nfrom onnxtr.models.detection.zoo import detection_predictor\nfrom onnxtr.models.predictor import OCRPredictor\nfrom onnxtr.models.preprocessor import PreProcessor\nfrom onnxtr.models.recognition.predictor import RecognitionPredictor\nfrom onnxtr.models.recognition.zoo import ARCHS as RECO_ARCHS\nfrom onnxtr.models.recognition.zoo import recognition_predictor\nfrom onnxtr.models.zoo import ocr_predictor\nfrom onnxtr.utils.repr import NestedObject\n\n\n# Create a dummy callback\nclass _DummyCallback:\n def __call__(self, loc_preds):\n return loc_preds\n\n\n@pytest.mark.parametrize(\n \"assume_straight_pages, straighten_pages, disable_page_orientation, disable_crop_orientation\",\n [\n [True, False, False, False],\n [False, False, True, True],\n [True, True, False, False],\n [False, True, True, True],\n [True, False, True, False],\n ],\n)\ndef test_ocrpredictor(\n mock_pdf, assume_straight_pages, straighten_pages, disable_page_orientation, disable_crop_orientation\n):\n det_bsize = 4\n det_predictor = DetectionPredictor(\n PreProcessor(output_size=(1024, 1024), batch_size=det_bsize),\n detection.db_mobilenet_v3_large(assume_straight_pages=assume_straight_pages),\n )\n\n reco_bsize = 16\n reco_predictor = RecognitionPredictor(\n PreProcessor(output_size=(32, 128), batch_size=reco_bsize, preserve_aspect_ratio=True),\n recognition.crnn_vgg16_bn(),\n )\n\n doc = DocumentFile.from_pdf(mock_pdf)\n\n predictor = OCRPredictor(\n det_predictor,\n reco_predictor,\n assume_straight_pages=assume_straight_pages,\n straighten_pages=straighten_pages,\n detect_orientation=True,\n detect_language=True,\n resolve_lines=True,\n resolve_blocks=True,\n disable_page_orientation=disable_page_orientation,\n disable_crop_orientation=disable_crop_orientation,\n )\n\n assert (\n predictor._page_orientation_disabled if disable_page_orientation else not predictor._page_orientation_disabled\n )\n assert (\n predictor._crop_orientation_disabled if disable_crop_orientation else not predictor._crop_orientation_disabled\n )\n\n if assume_straight_pages:\n assert predictor.crop_orientation_predictor is None\n if predictor.detect_orientation or predictor.straighten_pages:\n assert isinstance(predictor.page_orientation_predictor, NestedObject)\n else:\n assert predictor.page_orientation_predictor is None\n else:\n assert isinstance(predictor.crop_orientation_predictor, NestedObject)\n assert isinstance(predictor.page_orientation_predictor, NestedObject)\n\n out = predictor(doc)\n assert isinstance(out, Document)\n assert len(out.pages) == 2\n # Dimension check\n with pytest.raises(ValueError):\n input_page = (255 * np.random.rand(1, 256, 512, 3)).astype(np.uint8)\n _ = predictor([input_page])\n\n assert out.pages[0].orientation[\"value\"] in range(-2, 3)\n assert isinstance(out.pages[0].language[\"value\"], str)\n assert isinstance(out.render(), str)\n assert isinstance(out.pages[0].render(), str)\n assert isinstance(out.export(), dict)\n assert isinstance(out.pages[0].export(), dict)\n\n with pytest.raises(ValueError):\n _ = ocr_predictor(\"unknown_arch\")\n\n # Test with custom orientation models\n custom_crop_orientation_model = mobilenet_v3_small_crop_orientation()\n custom_page_orientation_model = mobilenet_v3_small_page_orientation()\n\n if assume_straight_pages:\n if predictor.detect_orientation or predictor.straighten_pages:\n # Overwrite the default orientation models\n predictor.crop_orientation_predictor = crop_orientation_predictor(custom_crop_orientation_model)\n predictor.page_orientation_predictor = page_orientation_predictor(custom_page_orientation_model)\n else:\n # Overwrite the default orientation models\n predictor.crop_orientation_predictor = crop_orientation_predictor(custom_crop_orientation_model)\n predictor.page_orientation_predictor = page_orientation_predictor(custom_page_orientation_model)\n\n out = predictor(doc)\n orientation = 0\n assert out.pages[0].orientation[\"value\"] == orientation\n\n\ndef test_trained_ocr_predictor(mock_payslip):\n doc = DocumentFile.from_images(mock_payslip)\n\n det_predictor = detection_predictor(\n \"db_resnet50\",\n batch_size=2,\n assume_straight_pages=True,\n symmetric_pad=True,\n preserve_aspect_ratio=False,\n )\n reco_predictor = recognition_predictor(\"crnn_vgg16_bn\", batch_size=128)\n\n predictor = OCRPredictor(\n det_predictor,\n reco_predictor,\n assume_straight_pages=True,\n straighten_pages=True,\n preserve_aspect_ratio=False,\n resolve_lines=True,\n resolve_blocks=True,\n )\n # test hooks\n predictor.add_hook(_DummyCallback())\n\n out = predictor(doc)\n\n assert out.pages[0].blocks[0].lines[0].words[0].value == \"Mr.\"\n geometry_mr = np.array([[0.1083984375, 0.0634765625], [0.1494140625, 0.0859375]])\n assert np.allclose(np.array(out.pages[0].blocks[0].lines[0].words[0].geometry), geometry_mr, rtol=0.05)\n\n assert out.pages[0].blocks[1].lines[0].words[-1].value == \"revised\"\n geometry_revised = np.array([[0.7548828125, 0.126953125], [0.8388671875, 0.1484375]])\n assert np.allclose(np.array(out.pages[0].blocks[1].lines[0].words[-1].geometry), geometry_revised, rtol=0.05)\n\n det_predictor = detection_predictor(\n \"db_resnet50\",\n batch_size=2,\n assume_straight_pages=True,\n preserve_aspect_ratio=True,\n symmetric_pad=True,\n )\n\n predictor = OCRPredictor(\n det_predictor,\n reco_predictor,\n assume_straight_pages=True,\n straighten_pages=True,\n preserve_aspect_ratio=True,\n symmetric_pad=True,\n resolve_lines=True,\n resolve_blocks=True,\n )\n\n out = predictor(doc)\n\n assert \"Mr\" in out.pages[0].blocks[0].lines[0].words[0].value\n\n # test list archs\n archs = predictor.list_archs()\n assert isinstance(archs, dict)\n assert archs[\"recognition_archs\"] == RECO_ARCHS\n assert archs[\"detection_archs\"] == DET_ARCHS\n\n\ndef _test_predictor(predictor):\n # Output checks\n assert isinstance(predictor, OCRPredictor)\n\n doc = [np.zeros((1024, 1024, 3), dtype=np.uint8)]\n out = predictor(doc)\n # Document\n assert isinstance(out, Document)\n\n # The input doc has 1 page\n assert len(out.pages) == 1\n # Dimension check\n with pytest.raises(ValueError):\n input_page = (255 * np.random.rand(1, 256, 512, 3)).astype(np.uint8)\n _ = predictor([input_page])\n\n\n@pytest.mark.parametrize(\"quantized\", [False, True])\n@pytest.mark.parametrize(\n \"det_arch, reco_arch\",\n [[det_arch, reco_arch] for det_arch, reco_arch in zip(detection.zoo.ARCHS, recognition.zoo.ARCHS)],\n)\ndef test_zoo_models(det_arch, reco_arch, quantized):\n # Model\n predictor = models.ocr_predictor(det_arch, reco_arch, load_in_8_bit=quantized)\n _test_predictor(predictor)\n\n # passing model instance directly\n det_model = detection.__dict__[det_arch]()\n reco_model = recognition.__dict__[reco_arch]()\n predictor = models.ocr_predictor(det_model, reco_model)\n _test_predictor(predictor)\n\n # passing recognition model as detection model\n with pytest.raises(ValueError):\n models.ocr_predictor(det_arch=reco_model)\n\n # passing detection model as recognition model\n with pytest.raises(ValueError):\n models.ocr_predictor(reco_arch=det_model)\n"
},
{
"path": "tests/common/test_transforms.py",
"content": "import numpy as np\nimport pytest\n\nfrom onnxtr.transforms import Normalize, Resize\n\n\ndef test_resize():\n output_size = (32, 32)\n transfo = Resize(output_size)\n input_t = np.ones((64, 64, 3), dtype=np.float32)\n out = transfo(input_t)\n\n assert np.all(out == 255)\n assert out.shape[:2] == output_size\n assert repr(transfo) == f\"Resize(output_size={output_size}, interpolation='2')\"\n\n transfo = Resize(output_size, preserve_aspect_ratio=True)\n input_t = np.ones((32, 64, 3), dtype=np.float32)\n out = transfo(input_t)\n\n assert out.shape[:2] == output_size\n assert not np.all(out == 255)\n # Asymetric padding\n assert np.all(out[-1] == 0) and np.all(out[0] == 255)\n\n # Symetric padding\n transfo = Resize(output_size, preserve_aspect_ratio=True, symmetric_pad=True)\n assert repr(transfo) == (\n f\"Resize(output_size={output_size}, interpolation='2', preserve_aspect_ratio=True, symmetric_pad=True)\"\n )\n out = transfo(input_t)\n assert out.shape[:2] == output_size\n # symetric padding\n assert np.all(out[-1] == 0) and np.all(out[0] == 0)\n\n # Inverse aspect ratio\n input_t = np.ones((64, 32, 3), dtype=np.float32)\n out = transfo(input_t)\n\n assert not np.all(out == 1)\n assert out.shape[:2] == output_size\n\n # Same aspect ratio\n output_size = (32, 128)\n transfo = Resize(output_size, preserve_aspect_ratio=True)\n out = transfo(np.ones((16, 64, 3), dtype=np.float32))\n assert out.shape[:2] == output_size\n\n\n@pytest.mark.parametrize(\n \"input_shape\",\n [\n [8, 32, 32, 3],\n [32, 32, 3],\n [32, 3],\n ],\n)\ndef test_normalize(input_shape):\n mean, std = [0.5, 0.5, 0.5], [0.5, 0.5, 0.5]\n transfo = Normalize(mean, std)\n input_t = np.ones(input_shape, dtype=np.float32)\n\n out = transfo(input_t)\n\n assert np.all(out == 1)\n assert repr(transfo) == f\"Normalize(mean={mean}, std={std})\"\n\n with pytest.raises(AssertionError):\n Normalize(mean=\"32\")\n\n with pytest.raises(AssertionError):\n Normalize(std=\"32\")\n"
},
{
"path": "tests/common/test_utils_data.py",
"content": "import os\nimport tempfile\nfrom pathlib import PosixPath\nfrom unittest.mock import patch\n\nimport pytest\n\nfrom onnxtr.utils.data import _urlretrieve, download_from_url\n\n\ndef test__urlretrieve():\n with tempfile.TemporaryDirectory() as temp_dir:\n file_path = os.path.join(temp_dir, \"crnn_mobilenet_v3_small-bded4d49.onnx\")\n _urlretrieve(\n \"https://github.com/felixdittrich92/OnnxTR/releases/download/v0.0.1/crnn_mobilenet_v3_small-bded4d49.onnx\",\n file_path,\n )\n assert os.path.exists(file_path), f\"File {file_path} does not exist.\"\n\n\n@patch(\"onnxtr.utils.data._urlretrieve\")\n@patch(\"pathlib.Path.mkdir\")\n@patch.dict(os.environ, {\"HOME\": \"/\"}, clear=True)\ndef test_download_from_url(mkdir_mock, urlretrieve_mock):\n download_from_url(\"test_url\")\n urlretrieve_mock.assert_called_with(\"test_url\", PosixPath(\"/.cache/onnxtr/test_url\"))\n\n\n@patch.dict(os.environ, {\"ONNXTR_CACHE_DIR\": \"/test\"}, clear=True)\n@patch(\"onnxtr.utils.data._urlretrieve\")\n@patch(\"pathlib.Path.mkdir\")\ndef test_download_from_url_customizing_cache_dir(mkdir_mock, urlretrieve_mock):\n download_from_url(\"test_url\")\n urlretrieve_mock.assert_called_with(\"test_url\", PosixPath(\"/test/test_url\"))\n\n\n@patch.dict(os.environ, {\"HOME\": \"/\"}, clear=True)\n@patch(\"pathlib.Path.mkdir\", side_effect=OSError)\n@patch(\"logging.error\")\ndef test_download_from_url_error_creating_directory(logging_mock, mkdir_mock):\n with pytest.raises(OSError):\n download_from_url(\"test_url\")\n logging_mock.assert_called_with(\n \"Failed creating cache direcotry at /.cache/onnxtr.\"\n \" You can change default cache directory using 'ONNXTR_CACHE_DIR' environment variable if needed.\"\n )\n\n\n@patch.dict(os.environ, {\"HOME\": \"/\", \"ONNXTR_CACHE_DIR\": \"/test\"}, clear=True)\n@patch(\"pathlib.Path.mkdir\", side_effect=OSError)\n@patch(\"logging.error\")\ndef test_download_from_url_error_creating_directory_with_env_var(logging_mock, mkdir_mock):\n with pytest.raises(OSError):\n download_from_url(\"test_url\")\n logging_mock.assert_called_with(\n \"Failed creating cache direcotry at /test using path from 'ONNXTR_CACHE_DIR' environment variable.\"\n )\n"
},
{
"path": "tests/common/test_utils_fonts.py",
"content": "from PIL.ImageFont import FreeTypeFont, ImageFont\n\nfrom onnxtr.utils.fonts import get_font\n\n\ndef test_get_font():\n # Attempts to load recommended OS font\n font = get_font()\n\n assert isinstance(font, (ImageFont, FreeTypeFont))\n"
},
{
"path": "tests/common/test_utils_geometry.py",
"content": "from copy import deepcopy\nfrom math import hypot\n\nimport numpy as np\nimport pytest\n\nfrom onnxtr.io import DocumentFile\nfrom onnxtr.utils import geometry\n\n\ndef test_bbox_to_polygon():\n assert geometry.bbox_to_polygon(((0, 0), (1, 1))) == ((0, 0), (1, 0), (0, 1), (1, 1))\n\n\ndef test_polygon_to_bbox():\n assert geometry.polygon_to_bbox(((0, 0), (1, 0), (0, 1), (1, 1))) == ((0, 0), (1, 1))\n\n\ndef test_order_points():\n # bbox format (xmin, ymin, xmax, ymax)\n bbox = np.array([1, 2, 5, 6])\n expected_bbox = np.array([\n [1, 2], # top-left\n [5, 2], # top-right\n [5, 6], # bottom-right\n [1, 6], # bottom-left\n ])\n out_bbox = geometry.order_points(bbox)\n assert np.all(out_bbox == expected_bbox)\n\n # quadrangle (unordered)\n quad = np.array([\n [5, 6], # br\n [1, 2], # tl\n [1, 6], # bl\n [5, 2], # tr\n ])\n expected_quad = expected_bbox\n out_quad = geometry.order_points(quad)\n assert np.all(out_quad == expected_quad)\n\n # already ordered quad\n ordered_quad = expected_bbox.copy()\n out_ordered = geometry.order_points(ordered_quad)\n assert np.all(out_ordered == expected_bbox)\n\n # float inputs\n quad_float = quad.astype(np.float32)\n out_float = geometry.order_points(quad_float)\n assert out_float.dtype == quad_float.dtype\n assert np.allclose(out_float, expected_quad)\n\n with pytest.raises(ValueError):\n geometry.order_points(np.array([1, 2, 3])) # wrong shape\n\n with pytest.raises(ValueError):\n geometry.order_points(np.zeros((5, 2))) # too many points\n\n\ndef test_detach_scores():\n # box test\n boxes = np.array([[0.1, 0.1, 0.2, 0.2, 0.9], [0.15, 0.15, 0.2, 0.2, 0.8]])\n pred = geometry.detach_scores([boxes])\n target1 = np.array([[0.1, 0.1, 0.2, 0.2], [0.15, 0.15, 0.2, 0.2]])\n target2 = np.array([0.9, 0.8])\n assert np.all(pred[0] == target1) and np.all(pred[1] == target2)\n # polygon test\n boxes = np.array([\n [[0.1, 0.1], [0.2, 0.2], [0.15, 0.25], [0.05, 0.15], [0.0, 0.9]],\n [[0.15, 0.15], [0.2, 0.2], [0.15, 0.25], [0.05, 0.15], [0.0, 0.8]],\n ])\n pred = geometry.detach_scores([boxes])\n target1 = np.array([\n [[0.1, 0.1], [0.2, 0.2], [0.15, 0.25], [0.05, 0.15]],\n [[0.15, 0.15], [0.2, 0.2], [0.15, 0.25], [0.05, 0.15]],\n ])\n target2 = np.array([0.9, 0.8])\n assert np.all(pred[0] == target1) and np.all(pred[1] == target2)\n\n\ndef test_resolve_enclosing_bbox():\n assert geometry.resolve_enclosing_bbox([((0, 0.5), (1, 0)), ((0.5, 0), (1, 0.25))]) == ((0, 0), (1, 0.5))\n pred = geometry.resolve_enclosing_bbox(np.array([[0.1, 0.1, 0.2, 0.2], [0.15, 0.15, 0.2, 0.2]]))\n assert pred.all() == np.array([0.1, 0.1, 0.2, 0.2]).all()\n\n\ndef test_resolve_enclosing_rbbox():\n box1 = np.asarray([[0.1, 0.1], [0.2, 0.2], [0.15, 0.25], [0.05, 0.15]])\n box2 = np.asarray([[0.5, 0.5], [0.6, 0.6], [0.55, 0.65], [0.45, 0.55]])\n\n pred = geometry.resolve_enclosing_rbbox([box1, box2])\n expected_raw = np.asarray([[0.05, 0.15], [0.1, 0.1], [0.6, 0.6], [0.55, 0.65]])\n target = geometry.order_points(expected_raw)\n assert np.allclose(pred, target, atol=1e-3)\n\n\ndef test_remap_boxes():\n pred = geometry.remap_boxes(\n np.asarray([[[0.25, 0.25], [0.25, 0.75], [0.75, 0.25], [0.75, 0.75]]]), (10, 10), (20, 20)\n )\n target = np.asarray([[[0.375, 0.375], [0.375, 0.625], [0.625, 0.375], [0.625, 0.625]]])\n assert np.all(pred == target)\n\n pred = geometry.remap_boxes(\n np.asarray([[[0.25, 0.25], [0.25, 0.75], [0.75, 0.25], [0.75, 0.75]]]), (10, 10), (20, 10)\n )\n target = np.asarray([[[0.25, 0.375], [0.25, 0.625], [0.75, 0.375], [0.75, 0.625]]])\n assert np.all(pred == target)\n\n with pytest.raises(ValueError):\n geometry.remap_boxes(\n np.asarray([[[0.25, 0.25], [0.25, 0.75], [0.75, 0.25], [0.75, 0.75]]]), (80, 40, 150), (160, 40)\n )\n\n with pytest.raises(ValueError):\n geometry.remap_boxes(np.asarray([[[0.25, 0.25], [0.25, 0.75], [0.75, 0.25], [0.75, 0.75]]]), (80, 40), (160,))\n\n orig_dimension = (100, 100)\n dest_dimensions = (200, 100)\n # Unpack dimensions\n height_o, width_o = orig_dimension\n height_d, width_d = dest_dimensions\n\n orig_box = np.asarray([[[0.25, 0.25], [0.25, 0.25], [0.75, 0.75], [0.75, 0.75]]])\n\n pred = geometry.remap_boxes(orig_box, orig_dimension, dest_dimensions)\n\n # Switch to absolute coords\n orig = np.stack((orig_box[:, :, 0] * width_o, orig_box[:, :, 1] * height_o), axis=2)[0]\n dest = np.stack((pred[:, :, 0] * width_d, pred[:, :, 1] * height_d), axis=2)[0]\n\n len_orig = hypot(orig[0][0] - orig[2][0], orig[0][1] - orig[2][1])\n len_dest = hypot(dest[0][0] - dest[2][0], dest[0][1] - dest[2][1])\n assert len_orig == len_dest\n\n alpha_orig = np.rad2deg(np.arctan((orig[0][1] - orig[2][1]) / (orig[0][0] - orig[2][0])))\n alpha_dest = np.rad2deg(np.arctan((dest[0][1] - dest[2][1]) / (dest[0][0] - dest[2][0])))\n assert alpha_orig == alpha_dest\n\n\ndef test_rotate_boxes():\n boxes = np.array([[0.1, 0.1, 0.8, 0.3, 0.5]])\n rboxes = np.array([[0.1, 0.1], [0.8, 0.1], [0.8, 0.3], [0.1, 0.3]])\n # Angle = 0\n rotated = geometry.rotate_boxes(boxes, angle=0.0, orig_shape=(1, 1))\n assert np.all(rotated == rboxes)\n # Angle < 1:\n rotated = geometry.rotate_boxes(boxes, angle=0.5, orig_shape=(1, 1))\n assert np.all(rotated == rboxes)\n # Angle = 30\n rotated = geometry.rotate_boxes(boxes, angle=30, orig_shape=(1, 1))\n assert rotated.shape == (1, 4, 2)\n\n boxes = np.array([[0.0, 0.0, 0.6, 0.2, 0.5]])\n # Angle = -90:\n rotated = geometry.rotate_boxes(boxes, angle=-90, orig_shape=(1, 1), min_angle=0)\n assert np.allclose(rotated, np.array([[[1, 0.0], [1, 0.6], [0.8, 0.6], [0.8, 0.0]]]))\n # Angle = 90\n rotated = geometry.rotate_boxes(boxes, angle=+90, orig_shape=(1, 1), min_angle=0)\n assert np.allclose(rotated, np.array([[[0, 1.0], [0, 0.4], [0.2, 0.4], [0.2, 1.0]]]))\n\n\n@pytest.fixture\ndef sample_geoms():\n return np.array([\n [[10, 10], [20, 10], [20, 20], [10, 20]],\n [\n [\n 30,\n 30,\n ],\n [40, 30],\n [40, 40],\n [30, 40],\n ],\n ])\n\n\ndef test_rotate_abs_geoms(sample_geoms):\n img_shape = (100, 100)\n angle = 45.0\n expanded_polys = geometry.rotate_abs_geoms(sample_geoms, angle, img_shape)\n\n # Check if the output has the correct shape\n assert expanded_polys.shape == sample_geoms.shape\n\n\ndef test_rotate_image():\n img = np.ones((32, 64, 3), dtype=np.float32)\n rotated = geometry.rotate_image(img, 30.0)\n assert rotated.shape[:-1] == (32, 64)\n assert rotated[0, 0, 0] == 0\n assert rotated[0, :, 0].sum() > 1\n\n # Expand\n rotated = geometry.rotate_image(img, 30.0, expand=True)\n assert rotated.shape[:-1] == (60, 120)\n assert rotated[0, :, 0].sum() <= 1\n\n # Expand\n rotated = geometry.rotate_image(img, 30.0, expand=True, preserve_origin_shape=True)\n assert rotated.shape[:-1] == (32, 64)\n assert rotated[0, :, 0].sum() <= 1\n\n # Expand with 90° rotation\n rotated = geometry.rotate_image(img, 90.0, expand=True)\n assert rotated.shape[:-1] == (64, 128)\n assert rotated[0, :, 0].sum() <= 1\n\n\ndef test_remove_image_padding():\n img = np.ones((32, 64, 3), dtype=np.float32)\n padded = np.pad(img, ((10, 10), (20, 20), (0, 0)))\n cropped = geometry.remove_image_padding(padded)\n assert np.all(cropped == img)\n\n # No padding\n cropped = geometry.remove_image_padding(img)\n assert np.all(cropped == img)\n\n\n@pytest.mark.parametrize(\n \"abs_geoms, img_size, rel_geoms\",\n [\n # Full image (boxes)\n [np.array([[0, 0, 32, 32]]), (32, 32), np.array([[0, 0, 1, 1]], dtype=np.float32)],\n # Full image (polygons)\n [\n np.array([[[0, 0], [32, 0], [32, 32], [0, 32]]]),\n (32, 32),\n np.array([[[0, 0], [1, 0], [1, 1], [0, 1]]], dtype=np.float32),\n ],\n # Quarter image (boxes)\n [np.array([[0, 0, 16, 16]]), (32, 32), np.array([[0, 0, 0.5, 0.5]], dtype=np.float32)],\n # Quarter image (polygons)\n [\n np.array([[[0, 0], [16, 0], [16, 16], [0, 16]]]),\n (32, 32),\n np.array([[[0, 0], [0.5, 0], [0.5, 0.5], [0, 0.5]]], dtype=np.float32),\n ],\n ],\n)\ndef test_convert_to_relative_coords(abs_geoms, img_size, rel_geoms):\n assert np.all(geometry.convert_to_relative_coords(abs_geoms, img_size) == rel_geoms)\n\n # Wrong format\n with pytest.raises(ValueError):\n geometry.convert_to_relative_coords(np.zeros((3, 5)), (32, 32))\n\n\ndef test_estimate_page_angle():\n straight_polys = np.array([\n [[0.3, 0.3], [0.4, 0.3], [0.4, 0.4], [0.3, 0.4]],\n [[0.4, 0.4], [0.5, 0.4], [0.5, 0.5], [0.4, 0.5]],\n [[0.5, 0.5], [0.6, 0.5], [0.6, 0.6], [0.5, 0.6]],\n ])\n rotated_polys = geometry.rotate_boxes(straight_polys, angle=20, orig_shape=(512, 512))\n angle = geometry.estimate_page_angle(rotated_polys)\n assert np.isclose(angle, 20)\n # Test divide by zero / NaN\n invalid_poly = np.array([[[0.5, 0.5], [0.5, 0.5], [0.5, 0.5], [0.5, 0.5]]])\n angle = geometry.estimate_page_angle(invalid_poly)\n assert angle == 0.0\n\n\ndef test_extract_crops(mock_pdf):\n doc_img = DocumentFile.from_pdf(mock_pdf)[0]\n num_crops = 2\n rel_boxes = np.array(\n [[idx / num_crops, idx / num_crops, (idx + 1) / num_crops, (idx + 1) / num_crops] for idx in range(num_crops)],\n dtype=np.float32,\n )\n abs_boxes = np.array(\n [\n [\n int(idx * doc_img.shape[1] / num_crops),\n int(idx * doc_img.shape[0]) / num_crops,\n int((idx + 1) * doc_img.shape[1] / num_crops),\n int((idx + 1) * doc_img.shape[0] / num_crops),\n ]\n for idx in range(num_crops)\n ],\n dtype=np.float32,\n )\n\n with pytest.raises(AssertionError):\n geometry.extract_crops(doc_img, np.zeros((1, 5)))\n\n for boxes in (rel_boxes, abs_boxes):\n croped_imgs = geometry.extract_crops(doc_img, boxes)\n # Number of crops\n assert len(croped_imgs) == num_crops\n # Data type and shape\n assert all(isinstance(crop, np.ndarray) for crop in croped_imgs)\n assert all(crop.ndim == 3 for crop in croped_imgs)\n\n # Identity\n assert np.all(\n doc_img == geometry.extract_crops(doc_img, np.array([[0, 0, 1, 1]], dtype=np.float32), channels_last=True)[0]\n )\n torch_img = np.transpose(doc_img, axes=(-1, 0, 1))\n assert np.all(\n torch_img\n == np.transpose(\n geometry.extract_crops(doc_img, np.array([[0, 0, 1, 1]], dtype=np.float32), channels_last=False)[0],\n axes=(-1, 0, 1),\n )\n )\n\n # No box\n assert geometry.extract_crops(doc_img, np.zeros((0, 4))) == []\n\n\n@pytest.mark.parametrize(\"assume_horizontal\", [True, False])\ndef test_extract_rcrops(mock_pdf, assume_horizontal):\n doc_img = DocumentFile.from_pdf(mock_pdf)[0]\n num_crops = 2\n rel_boxes = np.array(\n [\n [\n [idx / num_crops, idx / num_crops],\n [idx / num_crops + 0.1, idx / num_crops],\n [idx / num_crops + 0.1, idx / num_crops + 0.1],\n [idx / num_crops, idx / num_crops],\n ]\n for idx in range(num_crops)\n ],\n dtype=np.float32,\n )\n abs_boxes = deepcopy(rel_boxes)\n abs_boxes[:, :, 0] *= doc_img.shape[1]\n abs_boxes[:, :, 1] *= doc_img.shape[0]\n abs_boxes = abs_boxes.astype(np.int64)\n\n with pytest.raises(AssertionError):\n geometry.extract_rcrops(doc_img, np.zeros((1, 8)), assume_horizontal=assume_horizontal)\n for boxes in (rel_boxes, abs_boxes):\n croped_imgs = geometry.extract_rcrops(doc_img, boxes, assume_horizontal=assume_horizontal)\n # Number of crops\n assert len(croped_imgs) == num_crops\n # Data type and shape\n assert all(isinstance(crop, np.ndarray) for crop in croped_imgs)\n assert all(crop.ndim == 3 for crop in croped_imgs)\n\n # No box\n assert geometry.extract_rcrops(doc_img, np.zeros((0, 4, 2)), assume_horizontal=assume_horizontal) == []\n\n\n@pytest.mark.parametrize(\n \"format,input_shape,expected_shape\",\n [\n (\"BCHW\", (32, 3, 64, 64), (32, 3, 64, 64)),\n (\"BCHW\", (32, 64, 64, 3), (32, 3, 64, 64)),\n (\"BHWC\", (32, 64, 64, 3), (32, 64, 64, 3)),\n (\"BHWC\", (32, 3, 64, 64), (32, 64, 64, 3)),\n (\"XYZ\", (32, 3, 64, 64), (32, 3, 64, 64)),\n (\"CHW\", (3, 64, 64), (3, 64, 64)),\n (\"CHW\", (64, 64, 3), (3, 64, 64)),\n (\"HWC\", (64, 64, 3), (64, 64, 3)),\n (\"HWC\", (3, 64, 64), (64, 64, 3)),\n ],\n)\ndef test_shape_translate(format, input_shape, expected_shape):\n sample_data = np.random.rand(*input_shape).astype(np.float32)\n output_data = geometry.shape_translate(sample_data, format)\n\n # Assert that the output data has the expected shape\n assert output_data.shape == expected_shape\n"
},
{
"path": "tests/common/test_utils_multithreading.py",
"content": "import os\nfrom multiprocessing.pool import ThreadPool\nfrom unittest.mock import patch\n\nimport pytest\n\nfrom onnxtr.utils.multithreading import multithread_exec\n\n\n@pytest.mark.parametrize(\n \"input_seq, func, output_seq\",\n [\n [[1, 2, 3], lambda x: 2 * x, [2, 4, 6]],\n [[1, 2, 3], lambda x: x**2, [1, 4, 9]],\n [\n [\"this is\", \"show me\", \"I know\"],\n lambda x: x + \" the way\",\n [\"this is the way\", \"show me the way\", \"I know the way\"],\n ],\n ],\n)\ndef test_multithread_exec(input_seq, func, output_seq):\n assert list(multithread_exec(func, input_seq)) == output_seq\n assert list(multithread_exec(func, input_seq, 0)) == output_seq\n\n\n@patch.dict(os.environ, {\"ONNXTR_MULTIPROCESSING_DISABLE\": \"TRUE\"}, clear=True)\ndef test_multithread_exec_multiprocessing_disable():\n with patch.object(ThreadPool, \"map\") as mock_tp_map:\n multithread_exec(lambda x: x, [1, 2])\n assert not mock_tp_map.called\n"
},
{
"path": "tests/common/test_utils_reconstitution.py",
"content": "import numpy as np\nfrom test_io_elements import _mock_pages\n\nfrom onnxtr.utils import reconstitution\n\n\ndef test_synthesize_page():\n pages = _mock_pages()\n # Test without probability rendering\n render_no_proba = reconstitution.synthesize_page(pages[0].export(), draw_proba=False)\n assert isinstance(render_no_proba, np.ndarray)\n assert render_no_proba.shape == (*pages[0].dimensions, 3)\n\n # Test with probability rendering\n render_with_proba = reconstitution.synthesize_page(pages[0].export(), draw_proba=True)\n assert isinstance(render_with_proba, np.ndarray)\n assert render_with_proba.shape == (*pages[0].dimensions, 3)\n\n # Test with only one line\n pages_one_line = pages[0].export()\n pages_one_line[\"blocks\"][0][\"lines\"] = [pages_one_line[\"blocks\"][0][\"lines\"][0]]\n render_one_line = reconstitution.synthesize_page(pages_one_line, draw_proba=True)\n assert isinstance(render_one_line, np.ndarray)\n assert render_one_line.shape == (*pages[0].dimensions, 3)\n\n # Test with polygons\n pages_poly = pages[0].export()\n pages_poly[\"blocks\"][0][\"lines\"][0][\"geometry\"] = [(0, 0), (0, 1), (1, 1), (1, 0)]\n render_poly = reconstitution.synthesize_page(pages_poly, draw_proba=True)\n assert isinstance(render_poly, np.ndarray)\n assert render_poly.shape == (*pages[0].dimensions, 3)\n"
},
{
"path": "tests/common/test_utils_visualization.py",
"content": "import numpy as np\nimport pytest\nfrom test_io_elements import _mock_pages\n\nfrom onnxtr.utils import visualization\n\n\ndef test_visualize_page():\n pages = _mock_pages()\n image = np.ones((300, 200, 3))\n visualization.visualize_page(pages[0].export(), image, words_only=False)\n visualization.visualize_page(pages[0].export(), image, words_only=True, interactive=False)\n visualization.visualize_page(\n pages[0].export(), image, words_only=True, interactive=False, preserve_aspect_ratio=True\n )\n # geometry checks\n with pytest.raises(ValueError):\n visualization.create_obj_patch([1, 2], (100, 100))\n\n with pytest.raises(ValueError):\n visualization.create_obj_patch((1, 2), (100, 100))\n\n with pytest.raises(ValueError):\n visualization.create_obj_patch((1, 2, 3, 4, 5), (100, 100))\n # polygon patch\n pages = _mock_pages(polygons=True)\n image = np.ones((300, 200, 3))\n visualization.visualize_page(pages[0].export(), image, words_only=False)\n visualization.visualize_page(pages[0].export(), image, words_only=True, interactive=False)\n visualization.visualize_page(\n pages[0].export(), image, words_only=True, interactive=False, preserve_aspect_ratio=True\n )\n\n\ndef test_draw_boxes():\n image = np.ones((256, 256, 3), dtype=np.float32)\n boxes = [\n [0.1, 0.1, 0.2, 0.2],\n [0.15, 0.15, 0.19, 0.2], # to suppress\n [0.5, 0.5, 0.6, 0.55],\n [0.55, 0.5, 0.7, 0.55], # to suppress\n ]\n visualization.draw_boxes(boxes=np.array(boxes), image=image, block=False)\n"
},
{
"path": "tests/common/test_utils_vocabs.py",
"content": "from collections import Counter\n\nfrom onnxtr.utils import VOCABS\n\n\ndef test_vocabs_duplicates():\n for key, vocab in VOCABS.items():\n assert isinstance(vocab, str)\n\n duplicates = [char for char, count in Counter(vocab).items() if count > 1]\n assert not duplicates, f\"Duplicate characters in {key} vocab: {duplicates}\"\n"
},
{
"path": "tests/conftest.py",
"content": "from io import BytesIO\n\nimport cv2\nimport pytest\nimport requests\nfrom PIL import Image, ImageDraw\n\nfrom onnxtr.io import reader\nfrom onnxtr.utils import geometry\nfrom onnxtr.utils.fonts import get_font\n\n\ndef synthesize_text_img(\n text: str,\n font_size: int = 32,\n font_family=None,\n background_color=None,\n text_color=None,\n) -> Image.Image:\n background_color = (0, 0, 0) if background_color is None else background_color\n text_color = (255, 255, 255) if text_color is None else text_color\n\n font = get_font(font_family, font_size)\n left, top, right, bottom = font.getbbox(text)\n text_w, text_h = right - left, bottom - top\n h, w = int(round(1.3 * text_h)), int(round(1.1 * text_w))\n # If single letter, make the image square, otherwise expand to meet the text size\n img_size = (h, w) if len(text) > 1 else (max(h, w), max(h, w))\n\n img = Image.new(\"RGB\", img_size[::-1], color=background_color)\n d = ImageDraw.Draw(img)\n\n # Offset so that the text is centered\n text_pos = (int(round((img_size[1] - text_w) / 2)), int(round((img_size[0] - text_h) / 2)))\n # Draw the text\n d.text(text_pos, text, font=font, fill=text_color)\n return img\n\n\n@pytest.fixture(scope=\"session\")\ndef mock_vocab():\n return \"3K}7eé;5àÎYho]QwV6qU~W\\\"XnbBvcADfËmy.9ÔpÛ*{CôïE%M4#ÈR:g@T$x?0î£|za1ù8,OG€P-kçHëÀÂ2É/ûIJ'j(LNÙFut[)èZs+&°Sd=Ï!<â_Ç>rêi`l\" # noqa\n\n\n@pytest.fixture(scope=\"session\")\ndef mock_pdf(tmpdir_factory):\n # Page 1\n text_img = synthesize_text_img(\"I am a jedi!\", background_color=(255, 255, 255), text_color=(0, 0, 0))\n page = Image.new(text_img.mode, (1240, 1754), (255, 255, 255))\n page.paste(text_img, (50, 100))\n\n # Page 2\n text_img = synthesize_text_img(\"No, I am your father.\", background_color=(255, 255, 255), text_color=(0, 0, 0))\n _page = Image.new(text_img.mode, (1240, 1754), (255, 255, 255))\n _page.paste(text_img, (40, 300))\n\n # Save the PDF\n fn = tmpdir_factory.mktemp(\"data\").join(\"mock_pdf_file.pdf\")\n page.save(str(fn), \"PDF\", save_all=True, append_images=[_page])\n\n return str(fn)\n\n\n@pytest.fixture(scope=\"session\")\ndef mock_payslip(tmpdir_factory):\n url = \"https://3.bp.blogspot.com/-Es0oHTCrVEk/UnYA-iW9rYI/AAAAAAAAAFI/hWExrXFbo9U/s1600/003.jpg\"\n file = BytesIO(requests.get(url).content)\n folder = tmpdir_factory.mktemp(\"data\")\n fn = str(folder.join(\"mock_payslip.jpeg\"))\n with open(fn, \"wb\") as f:\n f.write(file.getbuffer())\n return fn\n\n\n@pytest.fixture(scope=\"session\")\ndef mock_tilted_payslip(mock_payslip, tmpdir_factory):\n image = reader.read_img_as_numpy(mock_payslip)\n image = geometry.rotate_image(image, 30, expand=True)\n tmp_path = str(tmpdir_factory.mktemp(\"data\").join(\"mock_tilted_payslip.jpg\"))\n cv2.imwrite(tmp_path, image)\n return tmp_path\n\n\n@pytest.fixture(scope=\"session\")\ndef mock_text_box_stream():\n url = \"https://doctr-static.mindee.com/models?id=v0.5.1/word-crop.png&src=0\"\n return requests.get(url).content\n\n\n@pytest.fixture(scope=\"session\")\ndef mock_text_box(mock_text_box_stream, tmpdir_factory):\n file = BytesIO(mock_text_box_stream)\n fn = tmpdir_factory.mktemp(\"data\").join(\"mock_text_box_file.png\")\n with open(fn, \"wb\") as f:\n f.write(file.getbuffer())\n return str(fn)\n\n\n@pytest.fixture(scope=\"session\")\ndef mock_artefact_image_stream():\n url = \"https://github.com/mindee/doctr/releases/download/v0.8.1/artefact_dummy.jpg\"\n return requests.get(url).content\n"
}
]