chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,62 @@
|
||||
---
|
||||
name: Breaking bug report
|
||||
about: Create a report about a breaking bug
|
||||
title: "[BUG: Breaking]"
|
||||
labels: 'bug: breaking'
|
||||
assignees: ''
|
||||
|
||||
---
|
||||
|
||||
## 🧨 Describe the Bug
|
||||
|
||||
A clear and concise description of the breaking issue (e.g., crash, OOM, exception, etc).
|
||||
|
||||
## 📄 Input Document
|
||||
|
||||
Attach the PDF or input file that triggered the error.
|
||||
|
||||
## 📤 Output Trace / Stack Trace
|
||||
|
||||
Paste the **complete** stack trace or error output, if available.
|
||||
|
||||
<details>
|
||||
<summary>Click to expand</summary>
|
||||
|
||||
```
|
||||
Paste stack trace here
|
||||
```
|
||||
|
||||
</details>
|
||||
|
||||
## ⚙️ Environment
|
||||
|
||||
Please fill in all relevant details:
|
||||
|
||||
- **Marker version**:
|
||||
- **Surya version**:
|
||||
- **Python version**:
|
||||
- **PyTorch version**:
|
||||
- **Transformers version**:
|
||||
- **Operating System** (incl. container info if relevant):
|
||||
|
||||
## ✅ Expected Behavior
|
||||
|
||||
What did you expect Marker to do?
|
||||
|
||||
## 📟 Command or Code Used
|
||||
|
||||
Paste the **exact bash command** or **Python code** you used to run Marker:
|
||||
|
||||
<details>
|
||||
<summary>Click to expand</summary>
|
||||
|
||||
```bash
|
||||
# or Python code block
|
||||
your_command_here --with-flags
|
||||
```
|
||||
|
||||
</details>
|
||||
|
||||
## 📎 Additional Context
|
||||
|
||||
Any other context that might help us debug this (e.g., CLI options, working directory, runtime settings).
|
||||
@@ -0,0 +1,24 @@
|
||||
---
|
||||
name: Feature request
|
||||
about: Suggest an idea for this project
|
||||
title: "[FEAT]"
|
||||
labels: enhancement
|
||||
assignees: ''
|
||||
|
||||
---
|
||||
|
||||
## ✨ Is your feature request related to a problem?
|
||||
|
||||
A clear and concise description of what the problem is.
|
||||
|
||||
## 💡 Describe the Solution You'd Like
|
||||
|
||||
A concise description of what you want to happen or how you envision it working.
|
||||
|
||||
## 📋 Alternatives Considered
|
||||
|
||||
Any alternative solutions or workarounds you've tried.
|
||||
|
||||
## 🧩 Additional Context
|
||||
|
||||
Any additional context, references, or related issues.
|
||||
@@ -0,0 +1,57 @@
|
||||
---
|
||||
name: Output bug report
|
||||
about: Create a report about poor output quality
|
||||
title: "[BUG: Output]"
|
||||
labels: 'bug: output'
|
||||
assignees: ''
|
||||
|
||||
---
|
||||
|
||||
## 📝 Describe the Output Issue
|
||||
|
||||
A clear and concise description of the incorrect or unexpected output.
|
||||
|
||||
## 📄 Input Document
|
||||
|
||||
Attach the PDF or input file used.
|
||||
|
||||
## 📤 Current Output
|
||||
|
||||
Paste the Markdown or HTML that Marker generated:
|
||||
|
||||
````markdown
|
||||
Paste output here
|
||||
`````
|
||||
|
||||
## ✅ Expected Output
|
||||
|
||||
Describe or paste what you expected Marker to generate.
|
||||
|
||||
## ⚙️ Environment
|
||||
|
||||
Please fill in all relevant details:
|
||||
|
||||
* **Marker version**:
|
||||
* **Surya version**:
|
||||
* **Python version**:
|
||||
* **PyTorch version**:
|
||||
* **Transformers version**:
|
||||
* **Operating System**:
|
||||
|
||||
## 📟 Command or Code Used
|
||||
|
||||
Paste the **exact bash command** or **Python code** you used to run Marker:
|
||||
|
||||
<details>
|
||||
<summary>Click to expand</summary>
|
||||
|
||||
```bash
|
||||
# or Python code block
|
||||
your_command_here --with-flags
|
||||
```
|
||||
|
||||
</details>
|
||||
|
||||
## 📎 Additional Context
|
||||
|
||||
Any other relevant info, configs, or assumptions.
|
||||
@@ -0,0 +1,30 @@
|
||||
name: Unit tests
|
||||
|
||||
on: [push]
|
||||
|
||||
jobs:
|
||||
build:
|
||||
runs-on: ${{ matrix.os }}
|
||||
strategy:
|
||||
matrix:
|
||||
os: [t4_gpu, ubuntu-latest, windows-latest]
|
||||
fail-fast: false
|
||||
env:
|
||||
# T4 can't run bf16 in vllm; size for the 16GB card. (No-op on the
|
||||
# CPU runners, which skip the VLM-backed tests.)
|
||||
VLLM_DTYPE: float16
|
||||
VLLM_GPU_TYPE: t4
|
||||
SURYA_INFERENCE_STARTUP_TIMEOUT: "1200"
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v3
|
||||
with:
|
||||
version: latest
|
||||
enable-cache: true
|
||||
- name: Set up Python 3.11
|
||||
run: uv python install 3.11
|
||||
- name: Install dependencies
|
||||
run: uv sync --frozen --group dev
|
||||
- name: Run tests
|
||||
run: uv run pytest
|
||||
@@ -0,0 +1,32 @@
|
||||
name: "Surya CLA Assistant"
|
||||
on:
|
||||
issue_comment:
|
||||
types: [created]
|
||||
pull_request_target:
|
||||
types: [opened,closed,synchronize]
|
||||
|
||||
# explicitly configure permissions, in case your GITHUB_TOKEN workflow permissions are set to read-only in repository settings
|
||||
permissions:
|
||||
actions: write
|
||||
contents: write
|
||||
pull-requests: write
|
||||
statuses: write
|
||||
|
||||
jobs:
|
||||
CLAAssistant:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: "Surya CLA Assistant"
|
||||
if: (github.event.comment.body == 'recheck' || github.event.comment.body == 'I have read the CLA Document and I hereby sign the CLA') || github.event_name == 'pull_request_target'
|
||||
uses: contributor-assistant/github-action@v2.3.0
|
||||
env:
|
||||
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
# the below token should have repo scope and must be manually added by you in the repository's secret
|
||||
# This token is required only if you have configured to store the signatures in a remote repository/organization
|
||||
PERSONAL_ACCESS_TOKEN: ${{ secrets.PERSONAL_ACCESS_TOKEN }}
|
||||
with:
|
||||
path-to-signatures: 'signatures/version1/cla.json'
|
||||
path-to-document: 'https://github.com/datalab-to/surya/blob/master/CLA.md'
|
||||
# branch should not be protected
|
||||
branch: 'master'
|
||||
allowlist: VikParuchuri
|
||||
@@ -0,0 +1,23 @@
|
||||
name: Python package
|
||||
on:
|
||||
push:
|
||||
tags:
|
||||
- "v*.*.*"
|
||||
jobs:
|
||||
build:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v3
|
||||
with:
|
||||
version: latest
|
||||
enable-cache: true
|
||||
- name: Set up Python 3.11
|
||||
run: uv python install 3.11
|
||||
- name: Build package
|
||||
run: uv build
|
||||
- name: Publish package
|
||||
env:
|
||||
UV_PUBLISH_TOKEN: ${{ secrets.PYPI_TOKEN }}
|
||||
run: uv publish
|
||||
@@ -0,0 +1,40 @@
|
||||
name: Test CLI scripts
|
||||
|
||||
on: [push]
|
||||
|
||||
jobs:
|
||||
build:
|
||||
runs-on: t4_gpu
|
||||
env:
|
||||
# T4 (Turing, compute 7.5) can't run bf16 in vllm; size vllm for the 16GB
|
||||
# card and give the cold start (image pull + model download) headroom.
|
||||
VLLM_DTYPE: float16
|
||||
VLLM_GPU_TYPE: t4
|
||||
SURYA_INFERENCE_STARTUP_TIMEOUT: "1200"
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v3
|
||||
with:
|
||||
version: latest
|
||||
enable-cache: true
|
||||
- name: Set up Python 3.11
|
||||
run: uv python install 3.11
|
||||
- name: Install dependencies
|
||||
run: uv sync --frozen --group dev
|
||||
- name: Download benchmark data
|
||||
run: |
|
||||
wget -O benchmark_data.zip "https://drive.google.com/uc?export=download&id=1NHrdYatR1rtqs2gPVfdvO0BAvocH8CJi"
|
||||
unzip -o benchmark_data.zip
|
||||
- name: Test detection
|
||||
run: uv run surya_detect benchmark_data/pdfs/switch_trans.pdf --page_range 0
|
||||
# Spawn the vllm server once and reuse it across the OCR/layout/table
|
||||
# steps (--keep_server) instead of paying a cold start three times.
|
||||
- name: Test OCR
|
||||
run: uv run surya_ocr benchmark_data/pdfs/switch_trans.pdf --page_range 0 --keep_server
|
||||
- name: Test layout
|
||||
run: uv run surya_layout benchmark_data/pdfs/switch_trans.pdf --page_range 0 --keep_server
|
||||
- name: Test table
|
||||
run: uv run surya_table benchmark_data/pdfs/switch_trans.pdf --page_range 0 --keep_server
|
||||
- name: Test detection folder
|
||||
run: uv run surya_detect benchmark_data/pdfs --page_range 0
|
||||
+172
@@ -0,0 +1,172 @@
|
||||
private.py
|
||||
.DS_Store
|
||||
local.env
|
||||
experiments
|
||||
test_data
|
||||
training
|
||||
wandb
|
||||
notebooks
|
||||
results
|
||||
data
|
||||
slices
|
||||
|
||||
# Byte-compiled / optimized / DLL files
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
*$py.class
|
||||
|
||||
# C extensions
|
||||
*.so
|
||||
|
||||
# Distribution / packaging
|
||||
.Python
|
||||
build/
|
||||
develop-eggs/
|
||||
dist/
|
||||
downloads/
|
||||
eggs/
|
||||
.eggs/
|
||||
lib/
|
||||
lib64/
|
||||
parts/
|
||||
sdist/
|
||||
var/
|
||||
wheels/
|
||||
share/python-wheels/
|
||||
*.egg-info/
|
||||
.installed.cfg
|
||||
*.egg
|
||||
MANIFEST
|
||||
|
||||
# PyInstaller
|
||||
# Usually these files are written by a python script from a template
|
||||
# before PyInstaller builds the exe, so as to inject date/other infos into it.
|
||||
*.manifest
|
||||
*.spec
|
||||
|
||||
# Installer logs
|
||||
pip-log.txt
|
||||
pip-delete-this-directory.txt
|
||||
|
||||
# Unit test / coverage reports
|
||||
htmlcov/
|
||||
.tox/
|
||||
.nox/
|
||||
.coverage
|
||||
.coverage.*
|
||||
.cache
|
||||
nosetests.xml
|
||||
coverage.xml
|
||||
*.cover
|
||||
*.py,cover
|
||||
.hypothesis/
|
||||
.pytest_cache/
|
||||
cover/
|
||||
|
||||
# Translations
|
||||
*.mo
|
||||
*.pot
|
||||
|
||||
# Django stuff:
|
||||
*.log
|
||||
local_settings.py
|
||||
db.sqlite3
|
||||
db.sqlite3-journal
|
||||
|
||||
# Flask stuff:
|
||||
instance/
|
||||
.webassets-cache
|
||||
|
||||
# Scrapy stuff:
|
||||
.scrapy
|
||||
|
||||
# Sphinx documentation
|
||||
docs/_build/
|
||||
|
||||
# PyBuilder
|
||||
.pybuilder/
|
||||
target/
|
||||
|
||||
# Jupyter Notebook
|
||||
.ipynb_checkpoints
|
||||
|
||||
# IPython
|
||||
profile_default/
|
||||
ipython_config.py
|
||||
|
||||
# pyenv
|
||||
# For a library or package, you might want to ignore these files since the code is
|
||||
# intended to run in multiple environments; otherwise, check them in:
|
||||
# .python-version
|
||||
|
||||
# pipenv
|
||||
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
|
||||
# However, in case of collaboration, if having platform-specific dependencies or dependencies
|
||||
# having no cross-platform support, pipenv may install dependencies that don't work, or not
|
||||
# install all needed dependencies.
|
||||
#Pipfile.lock
|
||||
|
||||
# poetry
|
||||
# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
|
||||
# This is especially recommended for binary packages to ensure reproducibility, and is more
|
||||
# commonly ignored for libraries.
|
||||
# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
|
||||
#poetry.lock
|
||||
|
||||
# pdm
|
||||
# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
|
||||
#pdm.lock
|
||||
# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
|
||||
# in version control.
|
||||
# https://pdm.fming.dev/#use-with-ide
|
||||
.pdm.toml
|
||||
|
||||
# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
|
||||
__pypackages__/
|
||||
|
||||
# Celery stuff
|
||||
celerybeat-schedule
|
||||
celerybeat.pid
|
||||
|
||||
# SageMath parsed files
|
||||
*.sage.py
|
||||
|
||||
# Environments
|
||||
.env
|
||||
.venv
|
||||
env/
|
||||
venv/
|
||||
ENV/
|
||||
env.bak/
|
||||
venv.bak/
|
||||
|
||||
# Spyder project settings
|
||||
.spyderproject
|
||||
.spyproject
|
||||
|
||||
# Rope project settings
|
||||
.ropeproject
|
||||
|
||||
# mkdocs documentation
|
||||
/site
|
||||
|
||||
# mypy
|
||||
.mypy_cache/
|
||||
.dmypy.json
|
||||
dmypy.json
|
||||
|
||||
# Pyre type checker
|
||||
.pyre/
|
||||
|
||||
# pytype static type analyzer
|
||||
.pytype/
|
||||
|
||||
# Cython debug symbols
|
||||
cython_debug/
|
||||
|
||||
# PyCharm
|
||||
# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
|
||||
# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
|
||||
# and can be added to the global gitignore or merged into this file. For a more nuclear
|
||||
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
|
||||
.idea/
|
||||
@@ -0,0 +1,12 @@
|
||||
repos:
|
||||
- repo: https://github.com/astral-sh/ruff-pre-commit
|
||||
# Ruff version.
|
||||
rev: v0.9.10
|
||||
hooks:
|
||||
# Run the linter.
|
||||
- id: ruff
|
||||
types_or: [ python, pyi ]
|
||||
args: [ --fix ]
|
||||
# Run the formatter.
|
||||
- id: ruff-format
|
||||
types_or: [ python, pyi ]
|
||||
@@ -0,0 +1,11 @@
|
||||
cff-version: 1.2.0
|
||||
message: "If you use this software, please cite it using the following metadata."
|
||||
title: "Surya: A lightweight framework for analyzing documents and PDFs at scale"
|
||||
authors:
|
||||
- family-names: Paruchuri
|
||||
given-names: Vikas
|
||||
- name: Datalab Team
|
||||
date-released: 2025-05-13
|
||||
url: https://github.com/datalab-to/surya
|
||||
version: 1.0.0
|
||||
repository-code: https://github.com/datalab-to/surya
|
||||
@@ -0,0 +1,24 @@
|
||||
Surya Contributor Agreement
|
||||
|
||||
This Surya Contributor Agreement ("SCA") applies to any contribution that you make to any product or project managed by us (the "project"), and sets out the intellectual property rights you grant to us in the contributed materials. The term "us" shall mean Endless Labs, Inc. The term "you" shall mean the person or entity identified below.
|
||||
|
||||
If you agree to be bound by these terms, sign by writing "I have read the CLA document and I hereby sign the CLA" in response to the CLA bot Github comment. Read this agreement carefully before signing. These terms and conditions constitute a binding legal agreement.
|
||||
|
||||
1. The term 'contribution' or 'contributed materials' means any source code, object code, patch, tool, sample, graphic, specification, manual, documentation, or any other material posted or submitted by you to the project.
|
||||
2. With respect to any worldwide copyrights, or copyright applications and registrations, in your contribution:
|
||||
- you hereby assign to us joint ownership, and to the extent that such assignment is or becomes invalid, ineffective or unenforceable, you hereby grant to us a perpetual, irrevocable, non-exclusive, worldwide, no-charge, royalty free, unrestricted license to exercise all rights under those copyrights. This includes, at our option, the right to sublicense these same rights to third parties through multiple levels of sublicensees or other licensing arrangements, including dual-license structures for commercial customers;
|
||||
- you agree that each of us can do all things in relation to your contribution as if each of us were the sole owners, and if one of us makes a derivative work of your contribution, the one who makes the derivative work (or has it made will be the sole owner of that derivative work;
|
||||
- you agree that you will not assert any moral rights in your contribution against us, our licensees or transferees;
|
||||
- you agree that we may register a copyright in your contribution and exercise all ownership rights associated with it; and
|
||||
- you agree that neither of us has any duty to consult with, obtain the consent of, pay or render an accounting to the other for any use or distribution of vour contribution.
|
||||
3. With respect to any patents you own, or that you can license without payment to any third party, you hereby grant to us a perpetual, irrevocable, non-exclusive, worldwide, no-charge, royalty-free license to:
|
||||
- make, have made, use, sell, offer to sell, import, and otherwise transfer your contribution in whole or in part, alone or in combination with or included in any product, work or materials arising out of the project to which your contribution was submitted, and
|
||||
- at our option, to sublicense these same rights to third parties through multiple levels of sublicensees or other licensing arrangements.
|
||||
If you or your affiliates institute patent litigation against any entity (including a cross-claim or counterclaim in a lawsuit) alleging that the contribution or any project it was submitted to constitutes direct or contributory patent infringement, then any patent licenses granted to you under this agreement for that contribution shall terminate as of the date such litigation is filed.
|
||||
4. Except as set out above, you keep all right, title, and interest in your contribution. The rights that you grant to us under these terms are effective on the date you first submitted a contribution to us, even if your submission took place before the date you sign these terms. Any contribution we make available under any license will also be made available under a suitable FSF (Free Software Foundation) or OSI (Open Source Initiative) approved license.
|
||||
5. You covenant, represent, warrant and agree that:
|
||||
- each contribution that you submit is and shall be an original work of authorship and you can legally grant the rights set out in this SCA;
|
||||
- to the best of your knowledge, each contribution will not violate any third party's copyrights, trademarks, patents, or other intellectual property rights; and
|
||||
- each contribution shall be in compliance with U.S. export control laws and other applicable export and import laws.
|
||||
You agree to notify us if you become aware of any circumstance which would make any of the foregoing representations inaccurate in any respect. Endless Labs, Inc. may publicly disclose your participation in the project, including the fact that you have signed the SCA.
|
||||
6. This SCA is governed by the laws of the State of California and applicable U.S. Federal law. Any choice of law rules will not apply.
|
||||
@@ -0,0 +1,51 @@
|
||||
Apache License
|
||||
Version 2.0, January 2004
|
||||
http://www.apache.org/licenses/
|
||||
|
||||
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
|
||||
|
||||
1. Definitions.
|
||||
|
||||
"License" shall mean the terms and conditions for use, reproduction, and distribution as defined by Sections 1 through 9 of this document.
|
||||
|
||||
"Licensor" shall mean the copyright owner or entity authorized by the copyright owner that is granting the License.
|
||||
|
||||
"Legal Entity" shall mean the union of the acting entity and all other entities that control, are controlled by, or are under common control with that entity. For the purposes of this definition, "control" means (i) the power, direct or indirect, to cause the direction or management of such entity, whether by contract or otherwise, or (ii) ownership of fifty percent (50%) or more of the outstanding shares, or (iii) beneficial ownership of such entity.
|
||||
|
||||
"You" (or "Your") shall mean an individual or Legal Entity exercising permissions granted by this License.
|
||||
|
||||
"Source" form shall mean the preferred form for making modifications, including but not limited to software source code, documentation source, and configuration files.
|
||||
|
||||
"Object" form shall mean any form resulting from mechanical transformation or translation of a Source form, including but not limited to compiled object code, generated documentation, and conversions to other media types.
|
||||
|
||||
"Work" shall mean the work of authorship, whether in Source or Object form, made available under the License, as indicated by a copyright notice that is included in or attached to the work (an example is provided in the Appendix below).
|
||||
|
||||
"Derivative Works" shall mean any work, whether in Source or Object form, that is based on (or derived from) the Work and for which the editorial revisions, annotations, elaborations, or other modifications represent, as a whole, an original work of authorship. For the purposes of this License, Derivative Works shall not include works that remain separable from, or merely link (or bind by name) to the interfaces of, the Work and Derivative Works thereof.
|
||||
|
||||
"Contribution" shall mean any work of authorship, including the original version of the Work and any modifications or additions to that Work or Derivative Works thereof, that is intentionally submitted to Licensor for inclusion in the Work by the copyright owner or by an individual or Legal Entity authorized to submit on behalf of the copyright owner. For the purposes of this definition, "submitted" means any form of electronic, verbal, or written communication sent to the Licensor or its representatives, including but not limited to communication on electronic mailing lists, source code control systems, and issue tracking systems that are managed by, or on behalf of, the Licensor for the purpose of discussing and improving the Work, but excluding communication that is conspicuously marked or otherwise designated in writing by the copyright owner as "Not a Contribution."
|
||||
|
||||
"Contributor" shall mean Licensor and any individual or Legal Entity on behalf of whom a Contribution has been received by Licensor and subsequently incorporated within the Work.
|
||||
|
||||
2. Grant of Copyright License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable copyright license to reproduce, prepare Derivative Works of, publicly display, publicly perform, sublicense, and distribute the Work and such Derivative Works in Source or Object form.
|
||||
|
||||
3. Grant of Patent License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable (except as stated in this section) patent license to make, have made, use, offer to sell, sell, import, and otherwise transfer the Work, where such license applies only to those patent claims licensable by such Contributor that are necessarily infringed by their Contribution(s) alone or by combination of their Contribution(s) with the Work to which such Contribution(s) was submitted. If You institute patent litigation against any entity (including a cross-claim or counterclaim in a lawsuit) alleging that the Work or a Contribution incorporated within the Work constitutes direct or contributory patent infringement, then any patent licenses granted to You under this License for that Work shall terminate as of the date such litigation is filed.
|
||||
|
||||
4. Redistribution. You may reproduce and distribute copies of the Work or Derivative Works thereof in any medium, with or without modifications, and in Source or Object form, provided that You meet the following conditions:
|
||||
|
||||
You must give any other recipients of the Work or Derivative Works a copy of this License; and
|
||||
You must cause any modified files to carry prominent notices stating that You changed the files; and
|
||||
You must retain, in the Source form of any Derivative Works that You distribute, all copyright, patent, trademark, and attribution notices from the Source form of the Work, excluding those notices that do not pertain to any part of the Derivative Works; and
|
||||
If the Work includes a "NOTICE" text file as part of its distribution, then any Derivative Works that You distribute must include a readable copy of the attribution notices contained within such NOTICE file, excluding those notices that do not pertain to any part of the Derivative Works, in at least one of the following places: within a NOTICE text file distributed as part of the Derivative Works; within the Source form or documentation, if provided along with the Derivative Works; or, within a display generated by the Derivative Works, if and wherever such third-party notices normally appear. The contents of the NOTICE file are for informational purposes only and do not modify the License. You may add Your own attribution notices within Derivative Works that You distribute, alongside or as an addendum to the NOTICE text from the Work, provided that such additional attribution notices cannot be construed as modifying the License.
|
||||
You may add Your own copyright statement to Your modifications and may provide additional or different license terms and conditions for use, reproduction, or distribution of Your modifications, or for any such Derivative Works as a whole, provided Your use, reproduction, and distribution of the Work otherwise complies with the conditions stated in this License.
|
||||
|
||||
5. Submission of Contributions. Unless You explicitly state otherwise, any Contribution intentionally submitted for inclusion in the Work by You to the Licensor shall be under the terms and conditions of this License, without any additional terms or conditions. Notwithstanding the above, nothing herein shall supersede or modify the terms of any separate license agreement you may have executed with Licensor regarding such Contributions.
|
||||
|
||||
6. Trademarks. This License does not grant permission to use the trade names, trademarks, service marks, or product names of the Licensor, except as required for reasonable and customary use in describing the origin of the Work and reproducing the content of the NOTICE file.
|
||||
|
||||
7. Disclaimer of Warranty. Unless required by applicable law or agreed to in writing, Licensor provides the Work (and each Contributor provides its Contributions) on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied, including, without limitation, any warranties or conditions of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A PARTICULAR PURPOSE. You are solely responsible for determining the appropriateness of using or redistributing the Work and assume any risks associated with Your exercise of permissions under this License.
|
||||
|
||||
8. Limitation of Liability. In no event and under no legal theory, whether in tort (including negligence), contract, or otherwise, unless required by applicable law (such as deliberate and grossly negligent acts) or agreed to in writing, shall any Contributor be liable to You for damages, including any direct, indirect, special, incidental, or consequential damages of any character arising as a result of this License or out of the use or inability to use the Work (including but not limited to damages for loss of goodwill, work stoppage, computer failure or malfunction, or any and all other commercial damages or losses), even if such Contributor has been advised of the possibility of such damages.
|
||||
|
||||
9. Accepting Warranty or Additional Liability. While redistributing the Work or Derivative Works thereof, You may choose to offer, and charge a fee for, acceptance of support, warranty, indemnity, or other liability obligations and/or rights consistent with this License. However, in accepting such obligations, You may act only on Your own behalf and on Your sole responsibility, not on behalf of any other Contributor, and only if You agree to indemnify, defend, and hold each Contributor harmless for any liability incurred by, or claims asserted against, such Contributor by reason of your accepting any such warranty or additional liability.
|
||||
|
||||
END OF TERMS AND CONDITIONS
|
||||
@@ -0,0 +1,59 @@
|
||||
AI PUBS OPEN RAIL-M LICENSE (MODIFIED)
|
||||
|
||||
Version 0.1, March 2, 2023 (Modified)
|
||||
http://licenses.ai/
|
||||
|
||||
PLEASE READ THESE TERMS CAREFULLY BEFORE USING THE MODEL OR A DERIVATIVE WORKS OF THE MODEL MADE AVAILABLE IN CONNECTION WITH THESE TERMS. BY DOWNLOADING, REPRODUCING, DISTRIBUTING OR USING THE MODEL OR A DERIVATIVE WORK OF THE MODEL IN ANY MANNER, YOU (“YOU”) AGREE TO BE BOUND BY THESE TERMS (THE “AGREEMENT”) TO THE EXCLUSION OF ALL OTHER TERMS. YOU REPRESENT AND WARRANT THAT YOU HAVE THE AUTHORITY TO ENTER INTO THIS AGREEMENT; IF YOU ARE ENTERING INTO THIS AGREEMENT ON BEHALF OF AN ORGANIZATION OR ENTITY, REFERENCES TO AND “YOU” IN THIS AGREEMENT, REFER TO THAT ORGANIZATION OR ENTITY. IF YOU DO NOT AGREE TO ALL OF THE FOLLOWING, YOU MAY NOT DOWNLOAD, REPRODUCE, DISTRIBUTE OR USE THE MODEL OR A DERIVATIVE WORK OF THE MODEL IN ANY MANNER.
|
||||
Section I: PREAMBLE
|
||||
This OpenRAIL-M License, as modified, is generally applicable to any machine-learning Model.
|
||||
The “Open” nomenclature indicates that the licensed Model is be freely accessible to downstream and other users. The “RAIL” nomenclature indicates that there are use restrictions prohibiting the use of the Model. These restrictions are intended to avoid potential misuse. This License specifies that the use restrictions in the original License must apply to such derivatives.
|
||||
NOW THEREFORE, You and Licensor agree as follows:
|
||||
1. Definitions
|
||||
(a) “Complementary Material” means the applicable source code and scripts used to define, run, load, benchmark or evaluate the Model, and used to prepare data for training or evaluation, if any. This includes any accompanying documentation, tutorials, examples, and any related information, if any. Complementary Material is not licensed under this License.
|
||||
(b) "Contribution" means any work, including the original version of the Model and any modifications or additions to that Model or Derivatives of the Model thereof, that is intentionally submitted to Licensor for inclusion in the Model by the rights owner or by an individual or legal entity authorized to submit on behalf of the rights owner. For the purposes of this definition, “submitted” means any form of electronic, verbal, or written communication sent to the Licensor or its representatives, including but not limited to communication on electronic mailing lists, source code control systems, and issue tracking systems that are managed by, or on behalf of, the Licensor for the purpose of discussing and improving the Model, but excluding communication that is conspicuously marked or otherwise designated in writing by the rights owner as "Not a Contribution."
|
||||
(c) "Contributor" means Licensor and any individual or legal entity on behalf of whom a Contribution has been received by Licensor and subsequently incorporated within the Model.
|
||||
(d) “Data” means a collection of information and/or content extracted from the dataset used with the Model, including to train, pretrain, or otherwise evaluate the Model. The Data is not licensed under this License.
|
||||
(e) “Derivatives of the Model” means all modifications to the Model, works based on the Model, or any other model which is created or initialized by transfer of patterns of the weights, parameters, activations or output of the Model, to the other model, in order to cause the other model to perform similarly to the Model, including - but not limited to - distillation methods entailing the use of intermediate data representations or methods based on the generation of synthetic data by the Model for training the other model.
|
||||
(f) “Distribution” means any transmission, reproduction, publication, distribution, or other sharing of the Model or Derivatives of the Model to a third party, including providing the Model as a hosted service made available by electronic or other remote means, including but not limited to API-based or web access.
|
||||
(g) “Harm” includes but is not limited to physical, mental, psychological, financial and reputational damage, pain, or loss
|
||||
(h) "License" means the terms and conditions for use, reproduction, and Distribution as defined in this document.
|
||||
(i) “Licensor” means the rights owner or entity authorized by the rights owner that is granting the License, including the persons or entities that may have rights in the Model and/or distributing the Model.
|
||||
(j) “Model” means any accompanying machine-learning based assemblies (including checkpoints), consisting of learnt weights, parameters (including optimizer states), corresponding to the model architecture as embodied in the Complementary Material, that have been trained or tuned, in whole or in part on the Data, using the Complementary Material.
|
||||
(k) “Output” means the results of operating a Model as embodied in informational content resulting therefrom.
|
||||
(l) “Third Parties” means individuals or legal entities that are not under common control with Licensor or You.
|
||||
(m) "You" (or "Your") means an individual or legal entity exercising permissions granted by this License and/or making use of the Model for whichever purpose and in any field of use, including usage of the Model in an end-use application, including but not limited to a chatbot, translator, or image generator.
|
||||
Section II: INTELLECTUAL PROPERTY RIGHTS
|
||||
Both copyright and patent grants may apply to the Model and Derivatives of the Model. The Model and Derivatives of the Model are subject to additional terms as described in Section III, which shall govern the use of the Model and Derivatives of the Model even in the event Section II is held unenforceable.
|
||||
2. Grant of Copyright License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable copyright license to reproduce, prepare, publicly display, publicly perform, sublicense, and distribute the Model and Derivatives of the Model.
|
||||
3. Grant of Patent License. Subject to the terms and conditions of this License and where and as applicable, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable (except as stated in this paragraph) patent license to make, have made, use, offer to sell, sell, import, and otherwise transfer the Model and/or Derivatives of the Model where such license applies only to those patent claims licensable by such Contributor that are necessarily infringed by their Contribution(s) alone or by combination of their Contribution(s) with the Model or Derivatives of the Model to which such Contribution(s) was submitted. If You institute patent litigation against any entity (including a cross-claim or counterclaim in a lawsuit) alleging that the Model or Derivative of the Model and/or a Contribution incorporated within the Model or Derivative of the Model constitutes direct or contributory patent infringement, then any patent licenses granted to You under this License for the Model and/or Derivative of the Model shall terminate as of the date such litigation is asserted or filed.
|
||||
Section III: CONDITIONS OF USAGE, DISTRIBUTION AND REDISTRIBUTION
|
||||
4. Distribution and Redistribution. You may host the Model or Derivatives of the Model for remote access by Third Parties, including but not limited to software-as-a-service, reproduce, or Distribute copies of the Model or Derivatives of the Model thereof in any medium, with or without modifications, provided that You meet the conditions in this Section III:
|
||||
(a) Use-based restrictions in paragraph 5 MUST be included as an enforceable provision by You in any type of legal agreement (for example, a license) governing the use and/or distribution of the Model or Derivatives of the Model, and You shall give notice to subsequent users You Distribute to, that the Model and Derivatives of the Model are subject to paragraph 5;
|
||||
(b) You must give any Third Party recipients of the Model or Derivatives of the Model a copy of this License;
|
||||
(c) You must cause any modified files to carry prominent notices stating that You changed the files; and
|
||||
(d) You must retain all copyright, patent, trademark, and attribution notices excluding those notices that do not pertain to any part of the Model or Derivatives of the Model.
|
||||
You may add Your own copyright statement to Your modifications and may provide additional or different license terms and conditions, consistent with paragraph 4.a., for use, reproduction, or Distribution of Your modifications, or for any such Derivatives of the Model as a whole, provided Your use, reproduction, and Distribution of the Model otherwise complies with the conditions stated in this License.
|
||||
5. Use-based restrictions. The restrictions set forth in Attachment A are considered Use-based restrictions. Accordingly, You cannot use the Model or the Derivatives of the Model in violation of such restrictions. You may use the Model subject to this License, including only for lawful purposes and in accordance with the License. Use may include creating any content with, fine-tuning, updating, running, training, evaluating and/or re-parametrizing the Model. You shall require all of Your users who use the Model or a Derivative of the Model to comply with the terms of this paragraph 5.
|
||||
6. The Output You Generate. Except as set forth herein, Licensor claims no rights in the Output You generate using the Model. You are solely responsible for the Output you generate and its subsequent uses. No use of the Output can contravene any provision as stated in the License.
|
||||
7. Attribution. In connection with any Output, or use of Distribution of any Model or Derivatives of the Model, You agree to give appropriate credit and attribution to Licensor, provide a link to the original Model or Derivatives of the Model, provide a copy of this License, and identify any changes You have made to the Model or Derivatives of the Model (collectively, the “Attribution”). The Attribution must not suggest endorsement by any Licensor.
|
||||
8. Share-a-Like. As a condition to the license and authorizations herein, You agree to apply this License (to the exclusion of all others) to any and all copies of the Model, Derivatives of the Model, any changes or improvements to the Model or Derivatives of the Model, and to the Output and any derivatives, changes or improvements to or of the Output.
|
||||
Section IV: OTHER PROVISIONS
|
||||
9. Updates and Runtime Restrictions. To the maximum extent permitted by law, Licensor reserves the right to restrict (remotely or otherwise) usage of the Model in violation of this License, update the Model through electronic means, or cause modification to the Output resulting from updates to the Model based.
|
||||
10. Trademarks and related. Nothing in this License permits You to make use of Licensors’ trademarks, trade names, logos or to otherwise suggest endorsement or misrepresent the relationship between the parties; and any rights not expressly granted herein are reserved by the Licensors.
|
||||
11. Disclaimer of Warranty. Unless required by applicable law or agreed to in writing, Licensor provides the Model (and each Contributor provides its Contributions) on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied, including, without limitation, any warranties or conditions of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A PARTICULAR PURPOSE. You are solely responsible for determining the appropriateness of using or redistributing the Model and Derivatives of the Model, and assume any risks associated with Your exercise of permissions under this License.
|
||||
12. Limitation of Liability. In no event and under no legal theory, whether in tort (including negligence), contract, or otherwise, unless required by applicable law (such as deliberate and grossly negligent acts) or agreed to in writing, shall any Contributor be liable to You for damages, including any direct, indirect, special, incidental, or consequential damages of any character arising as a result of this License or out of the use or inability to use the Model (including but not limited to damages for loss of goodwill, work stoppage, computer failure or malfunction, or any and all other commercial damages or losses), even if such Contributor has been advised of the possibility of such damages.
|
||||
13. Accepting Warranty or Additional Liability. While Distributing the Model or Derivatives of the Model, You may choose to charge a fee in exchange for support, warranty, indemnity, or other obligations and/or rights consistent with this License. However, in accepting such obligations, You may act only on Your own behalf and on Your sole responsibility, not on behalf of any other Contributor or Licensor, and only if You agree to indemnify, defend, and hold each Contributor and the Licensor harmless for any liability incurred by, or claims asserted against, such Contributor or Licensor by reason of your accepting any such warranty or additional liability.
|
||||
14. If any provision of this License is held to be invalid, illegal or unenforceable, the remaining provisions shall be unaffected thereby and remain valid as if such provision had not been set forth herein.
|
||||
END OF TERMS AND CONDITIONS
|
||||
|
||||
Attachment A
|
||||
USE RESTRICTIONS
|
||||
As conditions to the Licenses set forth in this Agreement, You agree not to use, reproduce, modify, create or Distribute the Model, Derivatives of the Model, or Output (collectively, “Use”) in any of the following ways:
|
||||
1. Legal:
|
||||
(a) In any way that violates any applicable national, federal, state, local or international law or regulation; or
|
||||
(b) to directly or indirectly infringe or misappropriate any third party intellectual property rights (including those of Licensor or any Contributor)
|
||||
2. Commercial:
|
||||
(a) for any purpose if You (your employer, or the entity you are affiliated with) generated more than two million US Dollars ($5,000,000) in gross revenue in the prior year, except where Your Use is limited to personal use or research purposes;
|
||||
(b) for any purpose if You (your employer, or the entity you are affiliated with) has raised more than two million US dollars ($5,000,000) in total equity or debt funding from any source, except where Your Use is limited to personal use or research purposes; or
|
||||
(c) for any purpose if You (your employer, or the entity you are affiliated with) provides or otherwise makes available any product or service that competes with any product or service offered by or made available by Licensor or any of its affiliates.
|
||||
Commercial and broader use licenses may be available from Licensor at the following URL: https://www.datalab.to/
|
||||
@@ -0,0 +1,504 @@
|
||||
<p align="center">
|
||||
<img src="static/datalab-logo.png" alt="Datalab Logo" width="150"/>
|
||||
</p>
|
||||
<h1 align="center">Datalab</h1>
|
||||
<p align="center">
|
||||
<strong>State of the Art models for Document Intelligence</strong>
|
||||
</p>
|
||||
<p align="center">
|
||||
<a href="https://www.apache.org/licenses/LICENSE-2.0"><img src="https://img.shields.io/badge/Code%20License-Apache--2.0-green.svg" alt="Code License"></a>
|
||||
<a href="https://www.datalab.to/pricing"><img src="https://img.shields.io/badge/Model%20License-OpenRAIL--M-blue.svg" alt="Model License"></a>
|
||||
<a href="https://discord.gg/KuZwXNGnfH"><img src="https://img.shields.io/badge/Discord-Join%20us-5865F2?logo=discord&logoColor=white" alt="Discord"></a>
|
||||
</p>
|
||||
<p align="center">
|
||||
<a href="https://www.datalab.to"><img src="https://img.shields.io/badge/Homepage-datalab.to-blue" alt="Homepage"></a>
|
||||
<a href="https://documentation.datalab.to"><img src="https://img.shields.io/badge/Docs-Read%20the%20docs-blue" alt="Docs"></a>
|
||||
<a href="https://www.datalab.to/playground"><img src="https://img.shields.io/badge/Datalab Playground-Try%20it-orange" alt="Datalab Playground"></a>
|
||||
</p>
|
||||
|
||||
<hr/>
|
||||
|
||||
# Surya
|
||||
|
||||
Surya is a 650M param OCR model with these features:
|
||||
|
||||
- Accuracy - scores 83.3% on [olmOCR-bench](https://huggingface.co/datasets/allenai/olmOCR-bench) (top under 3B params)
|
||||
- Speed - throughput of 5 pages/s on an RTX 5090
|
||||
- Multilingual - scores 87.2% on an internal benchmark set of 91 languages (more [here](#multilingual))
|
||||
- Layout analysis (table, image, header, etc.) with reading order
|
||||
- Table recognition (rows + columns)
|
||||
|
||||
We also ship smaller models for line-level text detection and ocr error detection. It works on a range of documents (see [usage](#usage) and [benchmarks](#benchmarks)).
|
||||
|
||||
## Try Datalab's Managed Platform
|
||||
|
||||
Our managed platform runs both Surya, and variants of our highest accuracy model, [Chandra](https://github.com/datalab-to/chandra).
|
||||
|
||||
Get started with **$5 in free credits** — [sign up](https://www.datalab.to/?utm_source=gh-surya) (takes under 30 seconds) or try our free [public playground](https://www.datalab.to/playground?utm_source=gh-surya).
|
||||
|
||||
## Model Information
|
||||
|
||||
<img src="static/images/olmocr_size_chart.png" width="700"/>
|
||||
|
||||
|
||||
| Detection | OCR |
|
||||
|:----------------------------------------------------------------:|:-----------------------------------------------------------------------:|
|
||||
| <img src="static/images/excerpt.png" width="280"/> | <img src="static/images/excerpt_text.png" width="280"/> |
|
||||
|
||||
| Layout | Table Recognition |
|
||||
|:------------------------------------------------------------------:|:-------------------------------------------------------------:|
|
||||
| <img src="static/images/excerpt_layout.png" width="280"/> | <img src="static/images/scanned_tablerec.png" width="280"/> |
|
||||
|
||||
|
||||
Surya is named for the [Hindu sun god](https://en.wikipedia.org/wiki/Surya), who has universal vision.
|
||||
|
||||
## Examples
|
||||
|
||||
Each row links to five annotated views of the same page: text-line detection, OCR, layout, reading order, and (when present) table recognition.
|
||||
|
||||
| Name | Detection | OCR | Layout | Order | Table Rec |
|
||||
|-------------------|:-----------------------------------:|------------------------------------------:|---------------------------------------------:|------------------------------------------------:|------------------------------------------------:|
|
||||
| Newspaper | [Image](static/images/newspaper.png) | [Image](static/images/newspaper_text.png) | [Image](static/images/newspaper_layout.png) | [Image](static/images/newspaper_reading.png) | |
|
||||
| Textbook | [Image](static/images/textbook.png) | [Image](static/images/textbook_text.png) | [Image](static/images/textbook_layout.png) | [Image](static/images/textbook_reading.png) | |
|
||||
| Tax Form | [Image](static/images/form.png) | [Image](static/images/form_text.png) | [Image](static/images/form_layout.png) | [Image](static/images/form_reading.png) | [Image](static/images/form_tablerec.png) |
|
||||
| Handwritten Notes | [Image](static/images/handwritten.png) | [Image](static/images/handwritten_text.png) | [Image](static/images/handwritten_layout.png) | [Image](static/images/handwritten_reading.png) | [Image](static/images/handwritten_tablerec.png) |
|
||||
| Corporate Doc | [Image](static/images/corporate.png) | [Image](static/images/corporate_text.png) | [Image](static/images/corporate_layout.png) | [Image](static/images/corporate_reading.png) | [Image](static/images/corporate_tablerec.png) |
|
||||
|
||||
# Commercial usage
|
||||
|
||||
The Surya code is licensed under Apache 2.0. The model weights use a modified AI Pubs Open Rail-M license (free for research, personal use, and startups under $5M funding/revenue). For broader commercial licensing of the model weights, visit our pricing page [here](https://www.datalab.to/pricing?utm_source=gh-surya).
|
||||
|
||||
# Installation
|
||||
|
||||
Install with:
|
||||
|
||||
```shell
|
||||
pip install surya-ocr
|
||||
```
|
||||
|
||||
## Inference backend prerequisites
|
||||
|
||||
Surya auto-spawns the server on first use, and you need `vllm` (NVIDIA GPU) or `llama.cpp` (CPU / Apple Silicon):
|
||||
|
||||
- **NVIDIA GPU:** [Docker](https://docs.docker.com/get-docker/) plus the [NVIDIA Container Toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html).
|
||||
- **CPU / Apple Silicon:** the `llama-server` binary from llama.cpp:
|
||||
```shell
|
||||
brew install llama.cpp # macOS
|
||||
# or grab a release from https://github.com/ggml-org/llama.cpp/releases
|
||||
```
|
||||
|
||||
## Upgrading from Surya v1
|
||||
|
||||
If you have v1 code, you can migrate to this:
|
||||
|
||||
```python
|
||||
# v2
|
||||
from surya.inference import SuryaInferenceManager
|
||||
from surya.recognition import RecognitionPredictor
|
||||
|
||||
manager = SuryaInferenceManager() # auto-spawns vllm or llama-server
|
||||
rec = RecognitionPredictor(manager)
|
||||
predictions = rec([image])
|
||||
```
|
||||
|
||||
What's different:
|
||||
- `SuryaInferenceManager` replaces `FoundationPredictor`. Same manager instance is shared across `LayoutPredictor`, `RecognitionPredictor`, `TableRecPredictor`.
|
||||
- Output schemas changed: see the per-section JSON tables below. Highlights — `text_lines` → `blocks` (with `html`); layout dropped `top_k`, added `count`; table_rec dropped `is_header` / `colspan` / `rowspan` from cells.
|
||||
|
||||
# Usage
|
||||
|
||||
Surya 2 runs layout, OCR, and table recognition through a single VLM. The inference manager will spawn one for you on first use; you can also point it at an existing server via `SURYA_INFERENCE_URL=http://host:port/v1`.
|
||||
|
||||
- Inspect the settings in `surya/settings.py`. You can override any setting via env var (e.g. `SURYA_INFERENCE_BACKEND=vllm`).
|
||||
- Text detection and OCR errors are separate models.
|
||||
|
||||
### Server lifecycle (`--keep_server`)
|
||||
|
||||
By default each command spawns the VLM server on startup and shuts it down on
|
||||
exit — so running several commands in a row pays the startup (and, on GPU, the
|
||||
model-load) cost every time. Pass `--keep_server` to leave the server running
|
||||
so later commands attach to it instead of re-spawning:
|
||||
|
||||
```shell
|
||||
surya_ocr DATA_PATH --keep_server # spawns the server and leaves it up
|
||||
surya_layout DATA_PATH # attaches to the running server
|
||||
surya_table DATA_PATH # ...and so on, no re-spawn
|
||||
```
|
||||
|
||||
`--keep_server` works on every command. Stop the server when you're done
|
||||
(`docker stop` the `surya-vllm-*` container, or kill the `llama-server`
|
||||
process), or set `SURYA_INFERENCE_KEEP_ALIVE=1` to make keep-alive the default.
|
||||
|
||||
## Interactive App
|
||||
|
||||
I've included a streamlit app that lets you interactively try Surya on images or PDF files. Run it with:
|
||||
|
||||
```shell
|
||||
pip install streamlit pdftext
|
||||
surya_gui
|
||||
```
|
||||
|
||||
## OCR (text recognition)
|
||||
|
||||
This command will write out a json file with the detected text and bboxes:
|
||||
|
||||
```shell
|
||||
surya_ocr DATA_PATH
|
||||
```
|
||||
|
||||
- `DATA_PATH` can be an image, pdf, or folder of images/pdfs
|
||||
- `--images` will save images of the pages and detected blocks (optional)
|
||||
- `--output_dir` specifies the directory to save results to instead of the default
|
||||
- `--page_range` specifies the page range to process in the PDF, specified as a single number, a comma separated list, a range, or comma separated ranges - example: `0,5-10,20`.
|
||||
- `--keep_server` leaves the inference server running after the command exits so later commands reuse it (see [Server lifecycle](#server-lifecycle---keep_server)). Available on every command.
|
||||
|
||||
The `results.json` file contains a dict keyed by input filename (no extension). Each value is a list of page dicts. Each page dict contains:
|
||||
|
||||
- `blocks` - per-block OCR results in reading order
|
||||
- `label` - canonicalized layout label (e.g. `Text`, `SectionHeader`, `Table`, `Equation`, `Picture`, `Form`, `PageHeader`, ...). See `surya/layout/label.py:LAYOUT_PRED_RELABEL` for the full canonical-name set.
|
||||
- `raw_label` - original label emitted by the model, before canonicalization
|
||||
- `reading_order` - 0-indexed position in layout output
|
||||
- `html` - block content as HTML (math wrapped in `<math>...</math>`, tables as `<table>...</table>`, etc.). `""` if the block was skipped
|
||||
- `polygon` - 4-corner polygon in `[[x0,y0],[x1,y0],[x1,y1],[x0,y1]]` order
|
||||
- `bbox` - axis-aligned `[x0, y0, x1, y1]` derived from the polygon
|
||||
- `confidence` - mean per-token probability across the block's decode (0-1)
|
||||
- `skipped` - true if the block was a visual label (e.g. Picture) and not OCR'd
|
||||
- `error` - true if the block OCR call failed
|
||||
- `image_bbox` - `[0, 0, width, height]` for the page image
|
||||
|
||||
**Performance tips**
|
||||
|
||||
- Throughput is governed by the inference backend. With `vllm`, raise `--max-num-seqs` / `--max-num-batched-tokens` (or `SURYA_INFERENCE_PARALLEL` on the client side) to keep more pages in flight. With `llama.cpp`, set `SURYA_INFERENCE_PARALLEL` to match `--parallel` on `llama-server`.
|
||||
- DPI can also impact throughput significantly - you can adjust the DPI settings to make the right throughput/accuracy tradeoff for your usecase. Try going from 192 to 96 for improved throughput.
|
||||
- MTP can also impact latency/throughput - you can adjust the vllm mtp config in settings.
|
||||
|
||||
### From python
|
||||
|
||||
```python
|
||||
from PIL import Image
|
||||
from surya.inference import SuryaInferenceManager
|
||||
from surya.recognition import RecognitionPredictor
|
||||
|
||||
manager = SuryaInferenceManager()
|
||||
recognition_predictor = RecognitionPredictor(manager)
|
||||
|
||||
# Default: full-page OCR. One VLM call per page. Returns one PageOCRResult per
|
||||
# image: `.blocks` (each with label, html, polygon, bbox, confidence, ...) and
|
||||
# `.image_bbox` — the same schema as block mode.
|
||||
predictions = recognition_predictor([Image.open(IMAGE_PATH)])
|
||||
|
||||
# Block mode: pre-run layout, then per-block OCR. Same return schema as above.
|
||||
# Auto-selected when `layout_results` is passed.
|
||||
from surya.layout import LayoutPredictor
|
||||
layout = LayoutPredictor(manager)
|
||||
layouts = layout([Image.open(IMAGE_PATH)])
|
||||
predictions = recognition_predictor([Image.open(IMAGE_PATH)], layouts)
|
||||
```
|
||||
|
||||
|
||||
## Text line detection
|
||||
|
||||
This command will write out a json file with the detected bboxes.
|
||||
|
||||
```shell
|
||||
surya_detect DATA_PATH
|
||||
```
|
||||
|
||||
- `DATA_PATH` can be an image, pdf, or folder of images/pdfs
|
||||
- `--images` will save images of the pages and detected text lines (optional)
|
||||
- `--output_dir` specifies the directory to save results to instead of the default
|
||||
- `--page_range` specifies the page range to process in the PDF, specified as a single number, a comma separated list, a range, or comma separated ranges - example: `0,5-10,20`.
|
||||
|
||||
The `results.json` file will contain a json dictionary where the keys are the input filenames without extensions. Each value will be a list of dictionaries, one per page of the input document. Each page dictionary contains:
|
||||
|
||||
- `bboxes` - detected bounding boxes for text
|
||||
- `bbox` - the axis-aligned rectangle for the text line in (x1, y1, x2, y2) format. (x1, y1) is the top left corner, and (x2, y2) is the bottom right corner.
|
||||
- `polygon` - the polygon for the text line in (x1, y1), (x2, y2), (x3, y3), (x4, y4) format. The points are in clockwise order from the top left.
|
||||
- `confidence` - the confidence of the model in the detected text (0-1)
|
||||
- `vertical_lines` - vertical lines detected in the document
|
||||
- `bbox` - the axis-aligned line coordinates.
|
||||
- `page` - the page number in the file
|
||||
- `image_bbox` - the bbox for the image in (x1, y1, x2, y2) format. (x1, y1) is the top left corner, and (x2, y2) is the bottom right corner. All line bboxes will be contained within this bbox.
|
||||
|
||||
**Performance tips**
|
||||
|
||||
Detection is a torch model. `DETECTOR_BATCH_SIZE` defaults to an auto-picked value at runtime; override the env var to control VRAM usage on GPU and raise it on larger cards.
|
||||
|
||||
### From python
|
||||
|
||||
```python
|
||||
from PIL import Image
|
||||
from surya.detection import DetectionPredictor
|
||||
|
||||
det_predictor = DetectionPredictor()
|
||||
predictions = det_predictor([Image.open(IMAGE_PATH)])
|
||||
```
|
||||
|
||||
## Layout and reading order
|
||||
|
||||
This command will write out a json file with the detected layout and reading order.
|
||||
|
||||
```shell
|
||||
surya_layout DATA_PATH
|
||||
```
|
||||
|
||||
- `DATA_PATH` can be an image, pdf, or folder of images/pdfs
|
||||
- `--images` will save images of the pages and detected text lines (optional)
|
||||
- `--output_dir` specifies the directory to save results to instead of the default
|
||||
- `--page_range` specifies the page range to process in the PDF, specified as a single number, a comma separated list, a range, or comma separated ranges - example: `0,5-10,20`.
|
||||
|
||||
The `results.json` file contains a dict keyed by input filename (no extension). Each value is a list of page dicts. Each page dict contains:
|
||||
|
||||
- `bboxes` - layout boxes in reading order
|
||||
- `polygon` - 4-corner polygon `[[x0,y0],[x1,y0],[x1,y1],[x0,y1]]`
|
||||
- `bbox` - axis-aligned `[x0, y0, x1, y1]` derived from the polygon
|
||||
- `label` - canonicalized label. One of `Caption`, `Footnote`, `Equation`, `ListGroup`, `PageHeader`, `PageFooter`, `Picture`, `SectionHeader`, `Table`, `Text`, `Figure`, `Code`, `Form`, `TableOfContents`, `ChemicalBlock`, `Diagram`, `Bibliography`, `BlankPage`
|
||||
- `raw_label` - original label emitted by the model
|
||||
- `position` - 0-indexed reading order
|
||||
- `count` - model's token estimate for OCR'ing this block (rounded to multiples of 50; used to size the per-block decode budget)
|
||||
- `confidence` - mean per-token probability across the layout decode (0-1)
|
||||
- `image_bbox` - `[0, 0, width, height]`
|
||||
- `raw` - raw JSON the layout model emitted, for debugging
|
||||
- `error` - true if the layout call failed
|
||||
|
||||
**Performance tips**
|
||||
|
||||
Layout runs through the shared inference backend. Throughput tuning is the same as OCR — see Performance tips above.
|
||||
|
||||
### From python
|
||||
|
||||
```python
|
||||
from PIL import Image
|
||||
from surya.inference import SuryaInferenceManager
|
||||
from surya.layout import LayoutPredictor
|
||||
|
||||
layout_predictor = LayoutPredictor(SuryaInferenceManager())
|
||||
layout_predictions = layout_predictor([Image.open(IMAGE_PATH)])
|
||||
```
|
||||
|
||||
## Table Recognition
|
||||
|
||||
This command will write out a json file with the detected table cells and row/column ids, along with row/column bounding boxes. If you want to get cell positions and text, along with nice formatting, check out the [marker](https://github.com/datalab-to/marker) repo. You can use the `TableConverter` to detect and extract tables in images and PDFs. It supports output in json (with bboxes), markdown, and html.
|
||||
|
||||
```shell
|
||||
surya_table DATA_PATH
|
||||
```
|
||||
|
||||
- `DATA_PATH` can be an image, pdf, or folder of images/pdfs
|
||||
- `--images` will save annotated row + column overlays alongside the json (optional)
|
||||
- `--output_dir` specifies the directory to save results to instead of the default
|
||||
- `--page_range` specifies the page range to process in the PDF, specified as a single number, a comma separated list, a range, or comma separated ranges - example: `0,5-10,20`.
|
||||
- `--skip_table_detection` tells table recognition not to detect tables first. Use this if your image is already cropped to a table.
|
||||
|
||||
The `results.json` file contains a dict keyed by input filename (no extension). Each value is a list of per-table dicts. Each table dict contains:
|
||||
|
||||
- `rows` - detected table rows in reading order
|
||||
- `polygon` / `bbox` - row geometry (same convention as everywhere else)
|
||||
- `row_id` - 0-indexed row id
|
||||
- `cols` - detected table columns
|
||||
- `polygon` / `bbox` - column geometry
|
||||
- `col_id` - 0-indexed column id
|
||||
- `cells` - geometric row × column intersections (simple mode)
|
||||
- `polygon` / `bbox` - cell geometry
|
||||
- `row_id`, `col_id`, `cell_id`
|
||||
- `html` - full `<table>...</table>` HTML (only populated when `predict_full` is used; handles spanning cells / header rows). `null` in simple mode.
|
||||
- `mode` - `"simple"` or `"full"`
|
||||
- `image_bbox` - the table crop bbox
|
||||
- `error` - true if the table_rec call failed
|
||||
- `raw` - raw model output, for debugging
|
||||
|
||||
**Performance tips**
|
||||
|
||||
Table recognition routes through the shared VLM. Throughput tuning is the same as OCR.
|
||||
|
||||
### From python
|
||||
|
||||
```python
|
||||
from PIL import Image
|
||||
from surya.inference import SuryaInferenceManager
|
||||
from surya.table_rec import TableRecPredictor
|
||||
|
||||
table_rec_predictor = TableRecPredictor(SuryaInferenceManager())
|
||||
|
||||
# Default: rows + columns only, cells derived from intersections.
|
||||
table_predictions = table_rec_predictor([Image.open(IMAGE_PATH)])
|
||||
|
||||
# Or full HTML output (better for spanning cells / headers):
|
||||
# table_predictions = table_rec_predictor.predict_full([image])
|
||||
```
|
||||
|
||||
## Math / equations
|
||||
|
||||
Surya 2 handles math inline as part of full-page OCR — recognized equations
|
||||
come back inside `<math>...</math>` tags in the same HTML output as
|
||||
surrounding prose, in KaTeX-compatible LaTeX. No separate LaTeX OCR pass.
|
||||
|
||||
# Inference Backends
|
||||
|
||||
Layout / OCR / table_rec all share one VLM, served either by `vllm` (GPU) or `llama.cpp` (CPU / Apple Silicon). The `SuryaInferenceManager` will spawn one automatically; you can also point at a pre-running server:
|
||||
|
||||
```bash
|
||||
# Attach to an existing vllm
|
||||
export SURYA_INFERENCE_BACKEND=vllm
|
||||
export SURYA_INFERENCE_URL=http://localhost:8000/v1
|
||||
```
|
||||
|
||||
| Setting | Default | Notes |
|
||||
|-----------------------------------|-----------------------------------|--------------------------------------------------------|
|
||||
| `SURYA_INFERENCE_BACKEND` | auto (vllm if NVIDIA, else llamacpp) | `vllm` \| `llamacpp` \| unset (auto) |
|
||||
| `SURYA_INFERENCE_URL` | (auto-spawn) | Attach to a running OpenAI-compatible server |
|
||||
| `SURYA_INFERENCE_PARALLEL` | 8 | Client-side concurrency to the backend |
|
||||
| `SURYA_INFERENCE_KEEP_ALIVE` | false | Leave the spawned server up after exit (cf. `--keep_server`) |
|
||||
| `SURYA_GUIDED_LAYOUT` | true | JSON-schema-constrained layout decode |
|
||||
|
||||
# Limitations
|
||||
|
||||
- This is specialized for document OCR. Performance on photos or natural scenes is not the goal.
|
||||
- Layout / OCR / table_rec all need a running inference backend (vllm or llama.cpp). Detection runs purely on torch and works without it.
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
If OCR isn't working properly:
|
||||
|
||||
- Try increasing resolution of the image so the text is bigger. If the resolution is already very high, try decreasing it to no more than a `2048px` width.
|
||||
- Preprocessing the image (binarizing, deskewing, etc) can help with very old/blurry images.
|
||||
- You can adjust `DETECTOR_BLANK_THRESHOLD` and `DETECTOR_TEXT_THRESHOLD` if you don't get good results. `DETECTOR_BLANK_THRESHOLD` controls the space between lines - any prediction below this number will be considered blank space. `DETECTOR_TEXT_THRESHOLD` controls how text is joined - any number above this is considered text. `DETECTOR_TEXT_THRESHOLD` should always be higher than `DETECTOR_BLANK_THRESHOLD`, and both should be in the 0-1 range. Looking at the heatmap from the debug output of the detector can tell you how to adjust these (if you see faint things that look like boxes, lower the thresholds, and if you see bboxes being joined together, raise the thresholds).
|
||||
|
||||
# Manual install
|
||||
|
||||
If you want to develop surya, you can install it manually with [uv](https://docs.astral.sh/uv/):
|
||||
|
||||
```bash
|
||||
git clone https://github.com/datalab-to/surya.git
|
||||
cd surya
|
||||
uv sync --group dev # installs runtime + dev deps
|
||||
uv run surya_ocr ... # or `source .venv/bin/activate` to enter the venv
|
||||
```
|
||||
|
||||
# Benchmarks
|
||||
|
||||
Surya 2 is a single VLM that handles layout analysis, OCR (full-page or
|
||||
per-block), and table recognition in one model. We evaluate end-to-end on
|
||||
[olmOCR-bench](https://huggingface.co/datasets/allenai/olmOCR-bench) — the
|
||||
standard quality benchmark for document parsers.
|
||||
|
||||
## olmOCR-bench
|
||||
|
||||
Pareto-optimal on the size-vs-score frontier, and best in class under 3B params.
|
||||
|
||||
| Model | Params | Score |
|
||||
|-----------------------------|----------:|---------:|
|
||||
| Infinity-Parser2-Pro | 35.1B | 87.6 |
|
||||
| Chandra OCR 2 (Datalab) | 4.0B | 85.9 |
|
||||
| dots.mocr | 3.0B | 83.9 |
|
||||
| **Surya OCR 2** (Datalab) | **0.65B** | **83.3** |
|
||||
| LightOnOCR 2-1B \* | 1.0B | 83.2 |
|
||||
| Chandra OCR 1 (Datalab) | 9.0B | 83.1 |
|
||||
| olmOCR (anchored) | 8.3B | 77.4 |
|
||||
| GOT OCR | 0.6B | 48.3 |
|
||||
|
||||
\* **LightOnOCR 2-1B** uses a different benchmark methodology than the other entries (see their [release notes](https://huggingface.co/lightonai/LightOnOCR-2-1B)); the score is included for context but is not directly comparable.
|
||||
|
||||
Comparison scores from the [olmOCR-bench dataset card](https://huggingface.co/datasets/allenai/olmOCR-bench).
|
||||
|
||||
Surya 2, per-source pass rate on the `default` preset (8,413 tests total):
|
||||
|
||||
| ArXiv | Base | Hdr/Ftr | TinyTxt | MultCol | OldScan | OldMath | Tables |
|
||||
|------:|-----:|--------:|--------:|--------:|--------:|--------:|-------:|
|
||||
| 88.3 | 99.7 | 92.5 | 93.7 | 82.4 | 41.8 | 81.4 | 86.6 |
|
||||
|
||||
## Multilingual
|
||||
|
||||
We also evaluate Surya 2 against a 91-language internal benchmark covering
|
||||
text accuracy, layout, tables, math, and reading order in documents drawn
|
||||
from each language.
|
||||
|
||||
**Overall pass rate: 87.2% across 91 languages.** 38 of the
|
||||
91 languages score ≥ 90%; 76 score ≥ 80%.
|
||||
|
||||
Top 15 widely-spoken languages:
|
||||
|
||||
| Code | Language | Score |
|
||||
|------|-------------|------:|
|
||||
| `ar` | Arabic | 72.7% |
|
||||
| `bn` | Bengali | 82.7% |
|
||||
| `zh` | Chinese | 82.5% |
|
||||
| `en` | English | 92.3% |
|
||||
| `fr` | French | 89.3% |
|
||||
| `de` | German | 89.7% |
|
||||
| `hi` | Hindi | 82.2% |
|
||||
| `it` | Italian | 93.0% |
|
||||
| `ja` | Japanese | 86.2% |
|
||||
| `ko` | Korean | 86.7% |
|
||||
| `fa` | Persian | 82.3% |
|
||||
| `pt` | Portuguese | 86.1% |
|
||||
| `ru` | Russian | 88.8% |
|
||||
| `es` | Spanish | 90.7% |
|
||||
| `vi` | Vietnamese | 73.2% |
|
||||
|
||||
See [static/docs/multilingual.md](static/docs/multilingual.md) for the full 91-language table.
|
||||
|
||||
## Throughput
|
||||
|
||||
Full-page OCR, 96 DPI input (~2,400 output tokens/page average), measured
|
||||
client-side against a running inference server.
|
||||
|
||||
### RTX 5090 (vllm)
|
||||
|
||||
`vllm/vllm-openai:v0.20.1`, single RTX 5090 (32 GB).
|
||||
|
||||
| Concurrency | Pages/s | Tokens/s | p50 (ms) | p95 (ms) | avg tok/page |
|
||||
|------------:|--------:|----------:|---:|---:|---:|
|
||||
| 128 | 5.35 | 12,884 | 18,915 | 42,538 | 2,410 |
|
||||
|
||||
### Apple Silicon (llama.cpp / Metal)
|
||||
|
||||
`llama-server` with Metal backend.
|
||||
|
||||
| `--parallel` | Pages/s | Tokens/s | p50 (ms) | p95 (ms) | avg tok/page | Power |
|
||||
|-------------:|---------:|---------:|---:|---:|---:|---:|
|
||||
| 8 | 0.108 | 254 | 59,313 | 129,173 | 2,360 | ~30 W |
|
||||
|
||||
## Reproducing
|
||||
|
||||
We score Surya 2 on olmOCR-bench by serving the model with `vllm` (or
|
||||
`llama.cpp`) and running the olmOCR-bench harness from
|
||||
[allenai/olmocr](https://github.com/allenai/olmocr), with some adjustments applied to account for our output HTML format.
|
||||
|
||||
# Training
|
||||
|
||||
Layout, OCR, and table recognition all share a single vision-language model
|
||||
(Qwen3.5-style architecture, ~650M params). It's trained on diverse document
|
||||
images to emit either a layout JSON or a full-page HTML output, depending on
|
||||
prompt. Text-line detection is a separate small torch model — a modified
|
||||
EfficientViT segformer trained from scratch on document line annotations.
|
||||
|
||||
If you want help finetuning Surya on your own data, or to use our managed
|
||||
training stack, reach us at hi@datalab.to.
|
||||
|
||||
# Thanks
|
||||
|
||||
This work would not have been possible without amazing open source AI work:
|
||||
|
||||
- [Qwen3-VL](https://huggingface.co/Qwen) from Alibaba
|
||||
- [vllm](https://github.com/vllm-project/vllm) and [llama.cpp](https://github.com/ggerganov/llama.cpp) for inference
|
||||
- [Segformer](https://arxiv.org/pdf/2105.15203.pdf) from NVIDIA
|
||||
- [EfficientViT](https://github.com/mit-han-lab/efficientvit) from MIT
|
||||
- [timm](https://github.com/huggingface/pytorch-image-models) from Ross Wightman
|
||||
- [transformers](https://github.com/huggingface/transformers) from huggingface
|
||||
- [CRAFT](https://github.com/clovaai/CRAFT-pytorch), a great scene text detection model
|
||||
|
||||
Thank you to everyone who makes open source AI possible.
|
||||
|
||||
# Citation
|
||||
|
||||
If you use surya (or the associated models) in your work or research, please consider citing us using the following BibTeX entry:
|
||||
|
||||
```bibtex
|
||||
@misc{paruchuri2025surya,
|
||||
author = {Vikas Paruchuri and Datalab Team},
|
||||
title = {Surya: A lightweight document OCR and analysis toolkit},
|
||||
year = {2025},
|
||||
howpublished = {\url{https://github.com/datalab-to/surya}},
|
||||
note = {GitHub repository},
|
||||
}
|
||||
@@ -0,0 +1,7 @@
|
||||
# WeHub 来源说明
|
||||
|
||||
- 原始项目：`datalab-to/surya`
|
||||
- 原始仓库：https://github.com/datalab-to/surya
|
||||
- 导入方式：上游默认分支的最新快照
|
||||
- 原作者、版权和许可证信息以原始仓库及本仓库 LICENSE 为准
|
||||
- 本文件仅用于记录来源，不代表 WeHub 是原项目作者
|
||||
@@ -0,0 +1,68 @@
|
||||
[project]
|
||||
name = "surya-ocr"
|
||||
version = "0.21.1"
|
||||
description = "OCR, layout, reading order, and table recognition in 90+ languages."
|
||||
readme = "README.md"
|
||||
license = { text = "Apache-2.0" }
|
||||
authors = [
|
||||
{ name = "Vik Paruchuri", email = "vik@datalab.to" },
|
||||
]
|
||||
requires-python = ">=3.10,<4"
|
||||
keywords = ["ocr", "pdf", "text detection", "text recognition", "tables"]
|
||||
dependencies = [
|
||||
"transformers>=5.12.1",
|
||||
"torch>=2.7.0,<3",
|
||||
"torchvision>=0.20.0,<1",
|
||||
"pydantic>=2.5.3,<3",
|
||||
"pydantic-settings>=2.1.0,<3",
|
||||
"python-dotenv>=1.0.0,<2",
|
||||
"pillow>=10.2.0,<11",
|
||||
"pypdfium2==5.9.0",
|
||||
"filetype>=1.2.0,<2",
|
||||
"click>=8.1.8,<9",
|
||||
"platformdirs>=4.3.6,<5",
|
||||
"opencv-python-headless==4.11.0.86",
|
||||
"openai>=1.55.0,<3",
|
||||
"httpx>=0.27.0,<1",
|
||||
"huggingface-hub>=1.5.0,<2",
|
||||
"filelock>=3.16.0,<4",
|
||||
"beautifulsoup4>=4.12.0,<5",
|
||||
]
|
||||
|
||||
[project.urls]
|
||||
Repository = "https://github.com/datalab-to/surya"
|
||||
|
||||
[project.scripts]
|
||||
surya_detect = "surya.scripts.detect_text:detect_text_cli"
|
||||
surya_ocr = "surya.scripts.ocr_text:ocr_text_cli"
|
||||
surya_layout = "surya.scripts.detect_layout:detect_layout_cli"
|
||||
surya_gui = "surya.scripts.run_streamlit_app:streamlit_app_cli"
|
||||
surya_table = "surya.scripts.table_recognition:table_recognition_cli"
|
||||
surya_screenshot = "surya.scripts.screenshot_app:main"
|
||||
|
||||
[dependency-groups]
|
||||
dev = [
|
||||
"pre-commit>=4.2.0",
|
||||
"jupyter>=1.0.0",
|
||||
"pymupdf>=1.23.8",
|
||||
"datasets>=2.16.1",
|
||||
"streamlit>=1.31.0",
|
||||
"pytest>=8.3.4",
|
||||
"pdftext>=0.5.1",
|
||||
"tabulate>=0.9.0",
|
||||
"flask>=3.0.0",
|
||||
]
|
||||
|
||||
[build-system]
|
||||
requires = ["hatchling"]
|
||||
build-backend = "hatchling.build"
|
||||
|
||||
[tool.hatch.build.targets.wheel]
|
||||
packages = ["surya"]
|
||||
|
||||
[tool.ruff]
|
||||
# rf-detr is vendored byte-for-byte from upstream (Roboflow); don't lint/format it.
|
||||
# force-exclude makes the exclude apply even when pre-commit passes explicit paths.
|
||||
extend-exclude = ["surya/common/rfdetr"]
|
||||
force-exclude = true
|
||||
|
||||
@@ -0,0 +1,7 @@
|
||||
[pytest]
|
||||
testpaths=tests
|
||||
pythonpath=.
|
||||
filterwarnings =
|
||||
ignore::UserWarning
|
||||
ignore::PendingDeprecationWarning
|
||||
ignore::DeprecationWarning
|
||||
@@ -0,0 +1,244 @@
|
||||
{
|
||||
"signedContributors": [
|
||||
{
|
||||
"name": "rishiraj",
|
||||
"id": 44090649,
|
||||
"comment_id": 2170578748,
|
||||
"created_at": "2024-06-15T19:31:20Z",
|
||||
"repoId": 741297064,
|
||||
"pullRequestNo": 135
|
||||
},
|
||||
{
|
||||
"name": "mmacvicar",
|
||||
"id": 59354,
|
||||
"comment_id": 2236493182,
|
||||
"created_at": "2024-07-18T13:17:43Z",
|
||||
"repoId": 741297064,
|
||||
"pullRequestNo": 152
|
||||
},
|
||||
{
|
||||
"name": "jimexist",
|
||||
"id": 622789,
|
||||
"comment_id": 2255151376,
|
||||
"created_at": "2024-07-29T07:23:55Z",
|
||||
"repoId": 741297064,
|
||||
"pullRequestNo": 160
|
||||
},
|
||||
{
|
||||
"name": "michaeldriscoll-avant",
|
||||
"id": 85255083,
|
||||
"comment_id": 2259143427,
|
||||
"created_at": "2024-07-30T20:21:33Z",
|
||||
"repoId": 741297064,
|
||||
"pullRequestNo": 161
|
||||
},
|
||||
{
|
||||
"name": "EdoardoPona",
|
||||
"id": 29152472,
|
||||
"comment_id": 2271115922,
|
||||
"created_at": "2024-08-06T11:58:00Z",
|
||||
"repoId": 741297064,
|
||||
"pullRequestNo": 167
|
||||
},
|
||||
{
|
||||
"name": "hidenori-endo",
|
||||
"id": 15546605,
|
||||
"comment_id": 2307217499,
|
||||
"created_at": "2024-08-23T14:31:17Z",
|
||||
"repoId": 741297064,
|
||||
"pullRequestNo": 182
|
||||
},
|
||||
{
|
||||
"name": "dobosevych",
|
||||
"id": 12053536,
|
||||
"comment_id": 2430376828,
|
||||
"created_at": "2024-10-22T21:48:34Z",
|
||||
"repoId": 741297064,
|
||||
"pullRequestNo": 220
|
||||
},
|
||||
{
|
||||
"name": "iammosespaulr",
|
||||
"id": 28682735,
|
||||
"comment_id": 2447941238,
|
||||
"created_at": "2024-10-30T17:55:23Z",
|
||||
"repoId": 741297064,
|
||||
"pullRequestNo": 235
|
||||
},
|
||||
{
|
||||
"name": "ArthurMor4is",
|
||||
"id": 42987302,
|
||||
"comment_id": 2515315717,
|
||||
"created_at": "2024-12-03T18:37:45Z",
|
||||
"repoId": 741297064,
|
||||
"pullRequestNo": 255
|
||||
},
|
||||
{
|
||||
"name": "tarun-menta",
|
||||
"id": 66506307,
|
||||
"comment_id": 2543457960,
|
||||
"created_at": "2024-12-15T05:43:33Z",
|
||||
"repoId": 741297064,
|
||||
"pullRequestNo": 261
|
||||
},
|
||||
{
|
||||
"name": "jonaskahn",
|
||||
"id": 4338500,
|
||||
"comment_id": 2556622097,
|
||||
"created_at": "2024-12-20T09:36:20Z",
|
||||
"repoId": 741297064,
|
||||
"pullRequestNo": 269
|
||||
},
|
||||
{
|
||||
"name": "kumsumit",
|
||||
"id": 95072784,
|
||||
"comment_id": 2574534622,
|
||||
"created_at": "2025-01-07T07:05:59Z",
|
||||
"repoId": 741297064,
|
||||
"pullRequestNo": 276
|
||||
},
|
||||
{
|
||||
"name": "kevinhu",
|
||||
"id": 6051736,
|
||||
"comment_id": 2614135351,
|
||||
"created_at": "2025-01-25T23:34:12Z",
|
||||
"repoId": 741297064,
|
||||
"pullRequestNo": 291
|
||||
},
|
||||
{
|
||||
"name": "zanussbaum",
|
||||
"id": 33707069,
|
||||
"comment_id": 3008673416,
|
||||
"created_at": "2025-06-26T14:20:46Z",
|
||||
"repoId": 741297064,
|
||||
"pullRequestNo": 403
|
||||
},
|
||||
{
|
||||
"name": "mebriki",
|
||||
"id": 35892987,
|
||||
"comment_id": 3154706976,
|
||||
"created_at": "2025-08-05T10:54:27Z",
|
||||
"repoId": 741297064,
|
||||
"pullRequestNo": 418
|
||||
},
|
||||
{
|
||||
"name": "starikovplusplus",
|
||||
"id": 56602036,
|
||||
"comment_id": 3168958011,
|
||||
"created_at": "2025-08-08T18:29:50Z",
|
||||
"repoId": 741297064,
|
||||
"pullRequestNo": 423
|
||||
},
|
||||
{
|
||||
"name": "sandy0kwon",
|
||||
"id": 78377296,
|
||||
"comment_id": 3207932260,
|
||||
"created_at": "2025-08-20T20:07:15Z",
|
||||
"repoId": 741297064,
|
||||
"pullRequestNo": 434
|
||||
},
|
||||
{
|
||||
"name": "n0kovo",
|
||||
"id": 16690056,
|
||||
"comment_id": 3208251881,
|
||||
"created_at": "2025-08-20T22:22:06Z",
|
||||
"repoId": 741297064,
|
||||
"pullRequestNo": 435
|
||||
},
|
||||
{
|
||||
"name": "davidxifeng",
|
||||
"id": 158052,
|
||||
"comment_id": 3249594859,
|
||||
"created_at": "2025-09-03T14:52:16Z",
|
||||
"repoId": 741297064,
|
||||
"pullRequestNo": 445
|
||||
},
|
||||
{
|
||||
"name": "u-ashish",
|
||||
"id": 14264791,
|
||||
"comment_id": 3258734182,
|
||||
"created_at": "2025-09-05T15:16:48Z",
|
||||
"repoId": 741297064,
|
||||
"pullRequestNo": 447
|
||||
},
|
||||
{
|
||||
"name": "Mohking1",
|
||||
"id": 63689545,
|
||||
"comment_id": 3314908963,
|
||||
"created_at": "2025-09-20T11:21:42Z",
|
||||
"repoId": 741297064,
|
||||
"pullRequestNo": 462
|
||||
},
|
||||
{
|
||||
"name": "wkpark",
|
||||
"id": 232347,
|
||||
"comment_id": 3330009557,
|
||||
"created_at": "2025-09-24T17:42:55Z",
|
||||
"repoId": 741297064,
|
||||
"pullRequestNo": 464
|
||||
},
|
||||
{
|
||||
"name": "coval3nte",
|
||||
"id": 65908512,
|
||||
"comment_id": 3848768229,
|
||||
"created_at": "2026-02-04T17:28:32Z",
|
||||
"repoId": 741297064,
|
||||
"pullRequestNo": 483
|
||||
},
|
||||
{
|
||||
"name": "bailey-coding",
|
||||
"id": 29517254,
|
||||
"comment_id": 3955014177,
|
||||
"created_at": "2026-02-24T22:09:52Z",
|
||||
"repoId": 741297064,
|
||||
"pullRequestNo": 487
|
||||
},
|
||||
{
|
||||
"name": "Br1an67",
|
||||
"id": 29810238,
|
||||
"comment_id": 3979412700,
|
||||
"created_at": "2026-03-01T07:32:18Z",
|
||||
"repoId": 741297064,
|
||||
"pullRequestNo": 489
|
||||
},
|
||||
{
|
||||
"name": "Jkker",
|
||||
"id": 23471930,
|
||||
"comment_id": 4184890988,
|
||||
"created_at": "2026-04-03T19:31:54Z",
|
||||
"repoId": 741297064,
|
||||
"pullRequestNo": 493
|
||||
},
|
||||
{
|
||||
"name": "YizukiAme",
|
||||
"id": 104178195,
|
||||
"comment_id": 4224654618,
|
||||
"created_at": "2026-04-10T14:55:26Z",
|
||||
"repoId": 741297064,
|
||||
"pullRequestNo": 498
|
||||
},
|
||||
{
|
||||
"name": "milahu",
|
||||
"id": 12958815,
|
||||
"comment_id": 4386642879,
|
||||
"created_at": "2026-05-06T09:17:07Z",
|
||||
"repoId": 741297064,
|
||||
"pullRequestNo": 509
|
||||
},
|
||||
{
|
||||
"name": "iRonin",
|
||||
"id": 27929,
|
||||
"comment_id": 4604690173,
|
||||
"created_at": "2026-06-02T16:49:40Z",
|
||||
"repoId": 741297064,
|
||||
"pullRequestNo": 512
|
||||
},
|
||||
{
|
||||
"name": "gkriegspeedbay",
|
||||
"id": 269034776,
|
||||
"comment_id": 4697896223,
|
||||
"created_at": "2026-06-13T07:41:35Z",
|
||||
"repoId": 741297064,
|
||||
"pullRequestNo": 519
|
||||
}
|
||||
]
|
||||
}
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 6.0 KiB |
@@ -0,0 +1,101 @@
|
||||
# Surya 2 — full 91-language results
|
||||
|
||||
Overall pass rate: **87.2%** across 91 languages (32,055 tests). Each language
|
||||
has between ~30 and ~750 tests.
|
||||
|
||||
Sorted alphabetically by language name. See [README](../../README.md#multilingual)
|
||||
for the curated subset.
|
||||
|
||||
| Code | Language | Score |
|
||||
|------|-------------------|------:|
|
||||
| `af` | Afrikaans | 87.7% |
|
||||
| `sq` | Albanian | 88.9% |
|
||||
| `am` | Amharic | 74.6% |
|
||||
| `ar` | Arabic | 72.7% |
|
||||
| `hy` | Armenian | 90.1% |
|
||||
| `as` | Assamese | 86.4% |
|
||||
| `az` | Azerbaijani | 91.5% |
|
||||
| `eu` | Basque | 85.5% |
|
||||
| `be` | Belarusian | 98.5% |
|
||||
| `bn` | Bengali | 82.7% |
|
||||
| `bs` | Bosnian | 92.5% |
|
||||
| `br` | Breton | 93.6% |
|
||||
| `bg` | Bulgarian | 95.4% |
|
||||
| `my` | Burmese | 88.2% |
|
||||
| `ca` | Catalan | 86.4% |
|
||||
| `zh` | Chinese | 82.5% |
|
||||
| `hr` | Croatian | 92.5% |
|
||||
| `cs` | Czech | 85.8% |
|
||||
| `da` | Danish | 84.5% |
|
||||
| `nl` | Dutch | 86.5% |
|
||||
| `en` | English | 92.3% |
|
||||
| `eo` | Esperanto | 83.1% |
|
||||
| `et` | Estonian | 76.8% |
|
||||
| `fi` | Finnish | 87.5% |
|
||||
| `fr` | French | 89.3% |
|
||||
| `gl` | Galician | 85.5% |
|
||||
| `ka` | Georgian | 91.4% |
|
||||
| `de` | German | 89.7% |
|
||||
| `el` | Greek | 80.7% |
|
||||
| `gu` | Gujarati | 83.4% |
|
||||
| `ha` | Hausa | 89.6% |
|
||||
| `he` | Hebrew | 90.9% |
|
||||
| `hi` | Hindi | 82.2% |
|
||||
| `hu` | Hungarian | 90.6% |
|
||||
| `is` | Icelandic | 89.5% |
|
||||
| `id` | Indonesian | 90.3% |
|
||||
| `ga` | Irish | 92.8% |
|
||||
| `it` | Italian | 93.0% |
|
||||
| `ja` | Japanese | 86.2% |
|
||||
| `jv` | Javanese | 91.1% |
|
||||
| `kn` | Kannada | 79.2% |
|
||||
| `kk` | Kazakh | 90.3% |
|
||||
| `km` | Khmer | 75.0% |
|
||||
| `ko` | Korean | 86.7% |
|
||||
| `ku` | Kurdish | 93.9% |
|
||||
| `ky` | Kyrgyz | 92.3% |
|
||||
| `lo` | Lao | 72.6% |
|
||||
| `la` | Latin | 86.1% |
|
||||
| `lv` | Latvian | 90.3% |
|
||||
| `lt` | Lithuanian | 85.4% |
|
||||
| `mk` | Macedonian | 95.3% |
|
||||
| `mg` | Malagasy | 95.0% |
|
||||
| `ms` | Malay | 91.2% |
|
||||
| `ml` | Malayalam | 84.7% |
|
||||
| `mr` | Marathi | 85.9% |
|
||||
| `mn` | Mongolian | 94.3% |
|
||||
| `ne` | Nepali | 84.9% |
|
||||
| `no` | Norwegian | 93.6% |
|
||||
| `or` | Oriya | 60.0% |
|
||||
| `ps` | Pashto | 72.0% |
|
||||
| `fa` | Persian | 82.3% |
|
||||
| `pl` | Polish | 91.4% |
|
||||
| `pt` | Portuguese | 86.1% |
|
||||
| `pa` | Punjabi | 76.5% |
|
||||
| `ro` | Romanian | 86.9% |
|
||||
| `ru` | Russian | 88.8% |
|
||||
| `sa` | Sanskrit | 78.8% |
|
||||
| `gd` | Scottish Gaelic | 92.5% |
|
||||
| `sr` | Serbian | 94.1% |
|
||||
| `sd` | Sindhi | 87.3% |
|
||||
| `si` | Sinhala | 85.4% |
|
||||
| `sk` | Slovak | 90.4% |
|
||||
| `sl` | Slovenian | 91.4% |
|
||||
| `so` | Somali | 97.9% |
|
||||
| `es` | Spanish | 90.7% |
|
||||
| `su` | Sundanese | 92.6% |
|
||||
| `sw` | Swahili | 93.5% |
|
||||
| `sv` | Swedish | 91.4% |
|
||||
| `ta` | Tamil | 89.9% |
|
||||
| `te` | Telugu | 79.2% |
|
||||
| `th` | Thai | 76.4% |
|
||||
| `tr` | Turkish | 85.4% |
|
||||
| `uk` | Ukrainian | 92.1% |
|
||||
| `ur` | Urdu | 68.7% |
|
||||
| `ug` | Uyghur | 70.2% |
|
||||
| `uz` | Uzbek | 88.9% |
|
||||
| `vi` | Vietnamese | 73.2% |
|
||||
| `cy` | Welsh | 95.1% |
|
||||
| `fy` | Western Frisian | 90.9% |
|
||||
| `xh` | Xhosa | 90.3% |
|
||||
| `yi` | Yiddish | 82.5% |
|
||||
@@ -0,0 +1,2 @@
|
||||
*
|
||||
!.gitignore
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 164 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 164 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 164 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 159 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 162 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 331 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 337 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 342 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 530 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 501 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 496 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 507 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 500 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 322 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 174 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 180 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 181 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 167 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 292 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 348 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 5.4 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 5.3 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 5.4 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 1.8 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 80 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 337 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 189 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 190 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 196 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 220 KiB |
@@ -0,0 +1,10 @@
|
||||
import os
|
||||
|
||||
# The rf-detr fast detector's DINOv2 pos-embed interpolation uses an antialiased-bicubic
|
||||
# op that has no MPS kernel. PyTorch reads PYTORCH_ENABLE_MPS_FALLBACK at `import torch`,
|
||||
# so set it here — before torch is imported — to let that one op fall back to CPU while
|
||||
# the rest of the model runs on MPS. This is only effective when surya is imported before
|
||||
# torch; when torch is already loaded, `_pick_device` (surya.common.rfdetr_torch) probes
|
||||
# the op at runtime and falls back to CPU device selection instead. setdefault so we never
|
||||
# clobber a value the user set deliberately.
|
||||
os.environ.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1")
|
||||
@@ -0,0 +1,3 @@
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,64 @@
|
||||
"""Pixel-content heuristics for detecting blank or near-uniform image regions.
|
||||
|
||||
Used by both the layout predictor (drop hallucinated layout blocks over empty
|
||||
space) and the recognition predictor (drop hallucinated text blocks from
|
||||
full-page OCR, decide whether an empty full-page output is a correct blank-page
|
||||
read or a failure).
|
||||
|
||||
Two signals, combined:
|
||||
* near-white fraction — most pixels have every RGB channel above a threshold
|
||||
* pixel-value standard deviation — the region is essentially one color
|
||||
(catches uniform-color fills that the white check misses)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
|
||||
# Per-channel value at/above which a pixel is considered "near-white".
|
||||
# Tolerates the small noise typical of PDF renders at 96 DPI.
|
||||
BLANK_WHITE_THRESHOLD = 245
|
||||
# Fraction of pixels that must be near-white for a region to count as blank.
|
||||
BLANK_PIXEL_FRACTION = 0.99
|
||||
# Pixel-value std below which a region is "essentially one color" regardless
|
||||
# of what that color is (catches solid-fill rectangles, dark banners, etc.).
|
||||
UNIFORM_COLOR_STD = 8.0
|
||||
|
||||
|
||||
def near_white_fraction(
|
||||
image: Image.Image, white_threshold: int = BLANK_WHITE_THRESHOLD
|
||||
) -> float:
|
||||
"""Fraction of pixels where every RGB channel ≥ ``white_threshold``."""
|
||||
arr = np.asarray(image.convert("RGB"))
|
||||
if arr.size == 0:
|
||||
return 0.0
|
||||
return float(np.all(arr >= white_threshold, axis=-1).mean())
|
||||
|
||||
|
||||
def is_blank_region(
|
||||
image: Image.Image,
|
||||
*,
|
||||
white_threshold: int = BLANK_WHITE_THRESHOLD,
|
||||
blank_pixel_fraction: float = BLANK_PIXEL_FRACTION,
|
||||
uniform_color_std: float = UNIFORM_COLOR_STD,
|
||||
) -> bool:
|
||||
"""True iff the image is essentially blank — either mostly near-white or
|
||||
near-uniform color. Use this on a per-block crop or a whole page.
|
||||
|
||||
Returns False for empty (0-pixel) crops so callers don't accidentally
|
||||
treat a degenerate bbox as blank.
|
||||
"""
|
||||
arr = np.asarray(image.convert("RGB"))
|
||||
if arr.size == 0:
|
||||
return False
|
||||
if np.all(arr >= white_threshold, axis=-1).mean() > blank_pixel_fraction:
|
||||
return True
|
||||
# Per-channel std — a uniform solid color (e.g., red banner with RGB=(200,50,50))
|
||||
# has each channel constant across pixels, but mixing channels inflates the
|
||||
# aggregate std. Check each channel independently.
|
||||
per_channel_std = arr.reshape(-1, arr.shape[-1]).std(axis=0)
|
||||
if float(per_channel_std.max()) < uniform_color_std:
|
||||
return True
|
||||
return False
|
||||
@@ -0,0 +1,163 @@
|
||||
from typing import Optional, Iterable, List, Union, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def verify_out_features_out_indices(
|
||||
out_features: Optional[Iterable[str]],
|
||||
out_indices: Optional[Iterable[int]],
|
||||
stage_names: Optional[Iterable[str]],
|
||||
):
|
||||
"""
|
||||
Verify that out_indices and out_features are valid for the given stage_names.
|
||||
"""
|
||||
if stage_names is None:
|
||||
raise ValueError("Stage_names must be set for transformers backbones")
|
||||
|
||||
if out_features is not None:
|
||||
if not isinstance(out_features, (list,)):
|
||||
raise ValueError(f"out_features must be a list got {type(out_features)}")
|
||||
if any(feat not in stage_names for feat in out_features):
|
||||
raise ValueError(
|
||||
f"out_features must be a subset of stage_names: {stage_names} got {out_features}"
|
||||
)
|
||||
if len(out_features) != len(set(out_features)):
|
||||
raise ValueError(
|
||||
f"out_features must not contain any duplicates, got {out_features}"
|
||||
)
|
||||
if out_features != (
|
||||
sorted_feats := [feat for feat in stage_names if feat in out_features]
|
||||
):
|
||||
raise ValueError(
|
||||
f"out_features must be in the same order as stage_names, expected {sorted_feats} got {out_features}"
|
||||
)
|
||||
|
||||
if out_indices is not None:
|
||||
if not isinstance(out_indices, (list, tuple)):
|
||||
raise ValueError(
|
||||
f"out_indices must be a list or tuple, got {type(out_indices)}"
|
||||
)
|
||||
# Convert negative indices to their positive equivalent: [-1,] -> [len(stage_names) - 1,]
|
||||
positive_indices = tuple(
|
||||
idx % len(stage_names) if idx < 0 else idx for idx in out_indices
|
||||
)
|
||||
if any(idx for idx in positive_indices if idx not in range(len(stage_names))):
|
||||
raise ValueError(
|
||||
f"out_indices must be valid indices for stage_names {stage_names}, got {out_indices}"
|
||||
)
|
||||
if len(positive_indices) != len(set(positive_indices)):
|
||||
msg = f"out_indices must not contain any duplicates, got {out_indices}"
|
||||
msg += (
|
||||
f"(equivalent to {positive_indices}))"
|
||||
if positive_indices != out_indices
|
||||
else ""
|
||||
)
|
||||
raise ValueError(msg)
|
||||
if positive_indices != tuple(sorted(positive_indices)):
|
||||
sorted_negative = tuple(
|
||||
idx
|
||||
for _, idx in sorted(
|
||||
zip(positive_indices, out_indices), key=lambda x: x[0]
|
||||
)
|
||||
)
|
||||
raise ValueError(
|
||||
f"out_indices must be in the same order as stage_names, expected {sorted_negative} got {out_indices}"
|
||||
)
|
||||
|
||||
if out_features is not None and out_indices is not None:
|
||||
if len(out_features) != len(out_indices):
|
||||
raise ValueError(
|
||||
"out_features and out_indices should have the same length if both are set"
|
||||
)
|
||||
if out_features != [stage_names[idx] for idx in out_indices]:
|
||||
raise ValueError(
|
||||
"out_features and out_indices should correspond to the same stages if both are set"
|
||||
)
|
||||
|
||||
|
||||
def _align_output_features_output_indices(
|
||||
out_features: Optional[List[str]],
|
||||
out_indices: Optional[Union[List[int], Tuple[int]]],
|
||||
stage_names: List[str],
|
||||
):
|
||||
"""
|
||||
Finds the corresponding `out_features` and `out_indices` for the given `stage_names`.
|
||||
|
||||
The logic is as follows:
|
||||
- `out_features` not set, `out_indices` set: `out_features` is set to the `out_features` corresponding to the
|
||||
`out_indices`.
|
||||
- `out_indices` not set, `out_features` set: `out_indices` is set to the `out_indices` corresponding to the
|
||||
`out_features`.
|
||||
- `out_indices` and `out_features` not set: `out_indices` and `out_features` are set to the last stage.
|
||||
- `out_indices` and `out_features` set: input `out_indices` and `out_features` are returned.
|
||||
|
||||
Args:
|
||||
out_features (`List[str]`): The names of the features for the backbone to output.
|
||||
out_indices (`List[int]` or `Tuple[int]`): The indices of the features for the backbone to output.
|
||||
stage_names (`List[str]`): The names of the stages of the backbone.
|
||||
"""
|
||||
if out_indices is None and out_features is None:
|
||||
out_indices = [len(stage_names) - 1]
|
||||
out_features = [stage_names[-1]]
|
||||
elif out_indices is None and out_features is not None:
|
||||
out_indices = [stage_names.index(layer) for layer in out_features]
|
||||
elif out_features is None and out_indices is not None:
|
||||
out_features = [stage_names[idx] for idx in out_indices]
|
||||
return out_features, out_indices
|
||||
|
||||
|
||||
def get_aligned_output_features_output_indices(
|
||||
out_features: Optional[List[str]],
|
||||
out_indices: Optional[Union[List[int], Tuple[int]]],
|
||||
stage_names: List[str],
|
||||
) -> Tuple[List[str], List[int]]:
|
||||
"""
|
||||
Get the `out_features` and `out_indices` so that they are aligned.
|
||||
|
||||
The logic is as follows:
|
||||
- `out_features` not set, `out_indices` set: `out_features` is set to the `out_features` corresponding to the
|
||||
`out_indices`.
|
||||
- `out_indices` not set, `out_features` set: `out_indices` is set to the `out_indices` corresponding to the
|
||||
`out_features`.
|
||||
- `out_indices` and `out_features` not set: `out_indices` and `out_features` are set to the last stage.
|
||||
- `out_indices` and `out_features` set: they are verified to be aligned.
|
||||
|
||||
Args:
|
||||
out_features (`List[str]`): The names of the features for the backbone to output.
|
||||
out_indices (`List[int]` or `Tuple[int]`): The indices of the features for the backbone to output.
|
||||
stage_names (`List[str]`): The names of the stages of the backbone.
|
||||
"""
|
||||
# First verify that the out_features and out_indices are valid
|
||||
verify_out_features_out_indices(
|
||||
out_features=out_features, out_indices=out_indices, stage_names=stage_names
|
||||
)
|
||||
output_features, output_indices = _align_output_features_output_indices(
|
||||
out_features=out_features, out_indices=out_indices, stage_names=stage_names
|
||||
)
|
||||
# Verify that the aligned out_features and out_indices are valid
|
||||
verify_out_features_out_indices(
|
||||
out_features=output_features,
|
||||
out_indices=output_indices,
|
||||
stage_names=stage_names,
|
||||
)
|
||||
return output_features, output_indices
|
||||
|
||||
|
||||
def find_pruneable_heads_and_indices(
|
||||
heads: list[int],
|
||||
n_heads: int,
|
||||
head_size: int,
|
||||
already_pruned_heads: set[int],
|
||||
) -> tuple[set[int], torch.LongTensor]:
|
||||
mask = torch.ones(n_heads, head_size, dtype=torch.bool)
|
||||
|
||||
heads = set(heads) - already_pruned_heads
|
||||
|
||||
for head in heads:
|
||||
# Shift the head index left by however many smaller heads
|
||||
# were already removed earlier.
|
||||
shifted_head = head - sum(1 for h in already_pruned_heads if h < head)
|
||||
mask[shifted_head] = False
|
||||
|
||||
index = torch.arange(n_heads * head_size)[mask.view(-1)].long()
|
||||
return heads, index
|
||||
@@ -0,0 +1,25 @@
|
||||
from typing import Optional, Any
|
||||
|
||||
import torch
|
||||
|
||||
from surya.settings import settings
|
||||
|
||||
|
||||
class ModelLoader:
|
||||
def __init__(self, checkpoint: Optional[str] = None):
|
||||
self.checkpoint = checkpoint
|
||||
|
||||
def model(
|
||||
self,
|
||||
device: torch.device | str | None = settings.TORCH_DEVICE_MODEL,
|
||||
dtype: Optional[torch.dtype | str] = settings.MODEL_DTYPE,
|
||||
attention_implementation: Optional[str] = None,
|
||||
) -> Any:
|
||||
raise NotImplementedError()
|
||||
|
||||
def processor(
|
||||
self,
|
||||
device: torch.device | str | None = settings.TORCH_DEVICE_MODEL,
|
||||
dtype: Optional[torch.dtype | str] = settings.MODEL_DTYPE,
|
||||
) -> Any:
|
||||
raise NotImplementedError()
|
||||
@@ -0,0 +1,3 @@
|
||||
from surya.common.order.predictor import OrderPredictor, load_order_predictor
|
||||
|
||||
__all__ = ["OrderPredictor", "load_order_predictor"]
|
||||
@@ -0,0 +1,137 @@
|
||||
"""Autoregressive reading-order head (inference only).
|
||||
|
||||
Box tokens (geometry + label) cross-attend to the FULL rf-detr encoder feature map, then an AR
|
||||
decoder emits the reading-order permutation as indices into the canonically (raster) sorted box
|
||||
sequence. Constrained greedy decode -> a valid permutation (every box once, none invented).
|
||||
|
||||
Vendored from training/models/rtdetr/order_ar.py (training code stripped). The 19-class layout
|
||||
taxonomy must match the layout detector the features come from.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
# 19-class layout taxonomy (sorted), must match the fast_layout detector's classes.
|
||||
LAYOUT_CLASSES = sorted(
|
||||
[
|
||||
"Caption",
|
||||
"Footnote",
|
||||
"Equation-Block",
|
||||
"List-Group",
|
||||
"Page-Header",
|
||||
"Page-Footer",
|
||||
"Image",
|
||||
"Section-Header",
|
||||
"Table",
|
||||
"Text",
|
||||
"Complex-Block",
|
||||
"Code-Block",
|
||||
"Form",
|
||||
"Table-Of-Contents",
|
||||
"Figure",
|
||||
"Chemical-Block",
|
||||
"Diagram",
|
||||
"Bibliography",
|
||||
"Blank-Page",
|
||||
]
|
||||
)
|
||||
N_LABELS = len(LAYOUT_CLASSES)
|
||||
MAX_BOXES = 128
|
||||
|
||||
|
||||
def box_features(boxes_1000):
|
||||
"""boxes 0-1000 [N,4] (x0,y0,x1,y1) -> [N,8] normalized (x0,y0,x1,y1,cx,cy,w,h)."""
|
||||
b = np.asarray(boxes_1000, dtype=np.float32) / 1000.0
|
||||
x0, y0, x1, y1 = b[:, 0], b[:, 1], b[:, 2], b[:, 3]
|
||||
cx, cy, w, h = (x0 + x1) / 2, (y0 + y1) / 2, (x1 - x0), (y1 - y0)
|
||||
return np.stack([x0, y0, x1, y1, cx, cy, w, h], axis=1)
|
||||
|
||||
|
||||
def canonical_order(boxes_1000, n_bands=24):
|
||||
"""Deterministic y-banded raster order: order[p] = original index of the box at raster pos p."""
|
||||
b = np.asarray(boxes_1000, dtype=np.float32)
|
||||
band_h = 1000.0 / n_bands
|
||||
keys = [
|
||||
(int(b[i, 1] // band_h), float(b[i, 0]), float(b[i, 1]), float(b[i, 2]))
|
||||
for i in range(len(b))
|
||||
]
|
||||
return sorted(range(len(b)), key=lambda i: keys[i])
|
||||
|
||||
|
||||
class ReadingOrderAR(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
d=128,
|
||||
layers=3,
|
||||
heads=4,
|
||||
feat_dim=256,
|
||||
feat_hw=28,
|
||||
max_boxes=MAX_BOXES,
|
||||
dropout=0.0,
|
||||
):
|
||||
super().__init__()
|
||||
self.d = d
|
||||
self.max_boxes = max_boxes
|
||||
self.use_feat = bool(feat_dim)
|
||||
self.geom = nn.Linear(8, d)
|
||||
self.lab = nn.Embedding(N_LABELS, d)
|
||||
if self.use_feat:
|
||||
self.feat_proj = nn.Linear(feat_dim, d)
|
||||
self.feat_pos = nn.Parameter(torch.zeros(1, feat_hw * feat_hw, d))
|
||||
ctx_layer = nn.TransformerDecoderLayer(
|
||||
d, heads, d * 4, batch_first=True, dropout=dropout
|
||||
)
|
||||
self.ctx = nn.TransformerDecoder(ctx_layer, layers)
|
||||
else:
|
||||
enc_layer = nn.TransformerEncoderLayer(
|
||||
d, heads, d * 4, batch_first=True, dropout=dropout
|
||||
)
|
||||
self.enc = nn.TransformerEncoder(enc_layer, layers)
|
||||
dec_layer = nn.TransformerDecoderLayer(
|
||||
d, heads, d * 4, batch_first=True, dropout=dropout
|
||||
)
|
||||
self.dec = nn.TransformerDecoder(dec_layer, layers)
|
||||
self.bos = nn.Parameter(torch.zeros(d))
|
||||
self.step = nn.Embedding(max_boxes + 1, d)
|
||||
self.out = nn.Linear(d, max_boxes)
|
||||
|
||||
def encode(self, feats, labels, mask, fmap=None):
|
||||
x = self.geom(feats) + self.lab(labels)
|
||||
if self.use_feat and fmap is not None:
|
||||
f = self.feat_proj(fmap) + self.feat_pos[:, : fmap.shape[1]]
|
||||
x = self.ctx(x, f, tgt_key_padding_mask=~mask)
|
||||
else:
|
||||
x = self.enc(x, src_key_padding_mask=~mask)
|
||||
return x
|
||||
|
||||
@torch.no_grad()
|
||||
def decode(self, feats, labels, mask, fmap=None):
|
||||
"""Greedy constrained decode -> list (per batch item) of raster-position permutations."""
|
||||
memory = self.encode(feats, labels, mask, fmap)
|
||||
B = memory.shape[0]
|
||||
K = mask.sum(1)
|
||||
out = []
|
||||
for b in range(B):
|
||||
k = int(K[b].item())
|
||||
mem = memory[b : b + 1, :k]
|
||||
used = torch.zeros(k, dtype=torch.bool, device=memory.device)
|
||||
din = [self.bos + self.step.weight[0]]
|
||||
seq = []
|
||||
for t in range(k):
|
||||
x = torch.stack(din, 0).unsqueeze(0)
|
||||
causal = torch.triu(
|
||||
torch.ones(t + 1, t + 1, device=memory.device, dtype=torch.bool), 1
|
||||
)
|
||||
h = self.dec(x, mem, tgt_mask=causal)
|
||||
logits = self.out(h[0, -1]).clone()
|
||||
logits[k:] = float("-inf")
|
||||
logits[:k] = logits[:k].masked_fill(used, float("-inf"))
|
||||
nxt = int(logits.argmax().item())
|
||||
seq.append(nxt)
|
||||
used[nxt] = True
|
||||
din.append(mem[0, nxt] + self.step.weight[min(t + 1, self.max_boxes)])
|
||||
out.append(seq)
|
||||
return out
|
||||
@@ -0,0 +1,117 @@
|
||||
"""OrderPredictor — runs the AR reading-order head on rf-detr detections + the encoder feature map.
|
||||
|
||||
Given, per page, the rf-detr projector feature map and the detected boxes (pixel xyxy) + labels,
|
||||
returns a reading-order position for each detection (0 = read first). Used by FastLayoutPredictor
|
||||
so layout always returns order.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from typing import List, Optional
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from surya.common.order.order_ar import (
|
||||
ReadingOrderAR,
|
||||
canonical_order,
|
||||
box_features,
|
||||
LAYOUT_CLASSES,
|
||||
MAX_BOXES,
|
||||
)
|
||||
|
||||
|
||||
class OrderPredictor:
|
||||
def __init__(self, model_dir: str, device: str = "cpu"):
|
||||
self.device = torch.device(device)
|
||||
ckpt_path = os.path.join(model_dir, "order_ar.pt")
|
||||
ck = torch.load(ckpt_path, map_location="cpu", weights_only=False)
|
||||
self.feat_dim = int(ck.get("feat_dim", 256))
|
||||
self.feat_hw = int(ck.get("feat_hw", 28))
|
||||
self.res = int(ck.get("res", 448))
|
||||
self.model = ReadingOrderAR(
|
||||
d=int(ck.get("d", 128)),
|
||||
layers=int(ck.get("layers", 3)),
|
||||
feat_dim=self.feat_dim,
|
||||
feat_hw=self.feat_hw,
|
||||
dropout=0.0,
|
||||
)
|
||||
self.model.load_state_dict(ck["model"])
|
||||
self.model.eval().to(self.device)
|
||||
|
||||
@torch.inference_mode()
|
||||
def order_page(self, feature_map, boxes_xyxy, labels, width, height) -> List[int]:
|
||||
"""feature_map: [C,F,F] tensor (rf-detr projector output for this page).
|
||||
boxes_xyxy: [N,4] pixel coords. labels: list of label strings (canonical class names).
|
||||
Returns position[i] for each detection i (0 = read first)."""
|
||||
n = len(boxes_xyxy)
|
||||
if n == 0:
|
||||
return []
|
||||
if n == 1:
|
||||
return [0]
|
||||
if n > MAX_BOXES: # fall back to raster order beyond the trained vocab width
|
||||
return _raster_positions(boxes_xyxy)
|
||||
|
||||
boxes = np.asarray(boxes_xyxy, dtype=np.float32)
|
||||
b1000 = np.empty_like(boxes)
|
||||
b1000[:, [0, 2]] = boxes[:, [0, 2]] / max(1.0, width) * 1000.0
|
||||
b1000[:, [1, 3]] = boxes[:, [1, 3]] / max(1.0, height) * 1000.0
|
||||
|
||||
order = canonical_order(b1000) # raster pos -> original idx
|
||||
b_raster = b1000[order]
|
||||
lab_raster = [
|
||||
LAYOUT_CLASSES.index(labels[p])
|
||||
if labels[p] in LAYOUT_CLASSES
|
||||
else LAYOUT_CLASSES.index("Text")
|
||||
for p in order
|
||||
]
|
||||
|
||||
feats = torch.from_numpy(box_features(b_raster)).unsqueeze(0).to(self.device)
|
||||
labs = torch.tensor(lab_raster, dtype=torch.long, device=self.device).unsqueeze(
|
||||
0
|
||||
)
|
||||
mask = torch.ones(1, n, dtype=torch.bool, device=self.device)
|
||||
# [C,F,F] -> [1, HW, C]
|
||||
fmap = (
|
||||
feature_map.reshape(self.feat_dim, -1)
|
||||
.transpose(0, 1)
|
||||
.unsqueeze(0)
|
||||
.to(self.device, dtype=torch.float32)
|
||||
)
|
||||
pred = self.model.decode(feats, labs, mask, fmap)[
|
||||
0
|
||||
] # raster positions, in reading order
|
||||
|
||||
# raster pos p -> original idx order[p]; reading sequence of original indices:
|
||||
reading = [order[p] for p in pred]
|
||||
position = [0] * n
|
||||
for rank, orig_idx in enumerate(reading):
|
||||
position[orig_idx] = rank
|
||||
return position
|
||||
|
||||
|
||||
def _raster_positions(boxes_xyxy) -> List[int]:
|
||||
"""Plain top-to-bottom, left-to-right fallback."""
|
||||
idx = sorted(
|
||||
range(len(boxes_xyxy)), key=lambda i: (boxes_xyxy[i][1], boxes_xyxy[i][0])
|
||||
)
|
||||
position = [0] * len(boxes_xyxy)
|
||||
for rank, i in enumerate(idx):
|
||||
position[i] = rank
|
||||
return position
|
||||
|
||||
|
||||
def load_order_predictor(checkpoint: Optional[str] = None, device: str = "cpu"):
|
||||
"""Resolve + load the order predictor, or return None if no checkpoint is configured/available."""
|
||||
from surya.common.rfdetr_torch import resolve_model_dir
|
||||
from surya.settings import settings
|
||||
|
||||
ckpt = checkpoint or getattr(settings, "FAST_ORDER_MODEL_CHECKPOINT", None)
|
||||
if not ckpt:
|
||||
return None
|
||||
try:
|
||||
model_dir = resolve_model_dir(ckpt)
|
||||
return OrderPredictor(model_dir, device=device)
|
||||
except Exception:
|
||||
return None
|
||||
@@ -0,0 +1,122 @@
|
||||
import copy
|
||||
from typing import List, Optional
|
||||
|
||||
import numpy as np
|
||||
from pydantic import BaseModel, field_validator, computed_field
|
||||
import numbers
|
||||
|
||||
|
||||
class PolygonBox(BaseModel):
|
||||
polygon: List[List[float]]
|
||||
confidence: Optional[float] = None
|
||||
|
||||
@field_validator("polygon", mode="before")
|
||||
@classmethod
|
||||
def convert_bbox_to_polygon(cls, value):
|
||||
if isinstance(value, (list, tuple)) and len(value) == 4:
|
||||
if all(isinstance(x, numbers.Number) for x in value):
|
||||
value = [float(v) for v in value]
|
||||
x_min, y_min, x_max, y_max = value
|
||||
polygon = [
|
||||
[x_min, y_min],
|
||||
[x_max, y_min],
|
||||
[x_max, y_max],
|
||||
[x_min, y_max],
|
||||
]
|
||||
return polygon
|
||||
elif all(
|
||||
isinstance(point, (list, tuple)) and len(point) == 2 for point in value
|
||||
):
|
||||
value = [[float(v) for v in point] for point in value]
|
||||
return value
|
||||
elif isinstance(value, np.ndarray):
|
||||
if value.shape == (4, 2):
|
||||
return value.tolist()
|
||||
|
||||
raise ValueError(
|
||||
f"Input must be either a bbox [x_min, y_min, x_max, y_max] or a polygon with 4 corners [(x,y), (x,y), (x,y), (x,y)]. All values must be numeric. You passed {value} of type {type(value)}. The first value is of type {type(value[0])}."
|
||||
)
|
||||
|
||||
@property
|
||||
def height(self):
|
||||
return self.bbox[3] - self.bbox[1]
|
||||
|
||||
@property
|
||||
def width(self):
|
||||
return self.bbox[2] - self.bbox[0]
|
||||
|
||||
@property
|
||||
def area(self):
|
||||
return self.width * self.height
|
||||
|
||||
@computed_field
|
||||
@property
|
||||
def bbox(self) -> List[float]:
|
||||
x_coords = [point[0] for point in self.polygon]
|
||||
y_coords = [point[1] for point in self.polygon]
|
||||
return [min(x_coords), min(y_coords), max(x_coords), max(y_coords)]
|
||||
|
||||
def rescale(self, processor_size, image_size):
|
||||
# Point is in x, y format
|
||||
page_width, page_height = processor_size
|
||||
|
||||
img_width, img_height = image_size
|
||||
width_scaler = img_width / page_width
|
||||
height_scaler = img_height / page_height
|
||||
|
||||
for corner in self.polygon:
|
||||
corner[0] = int(corner[0] * width_scaler)
|
||||
corner[1] = int(corner[1] * height_scaler)
|
||||
|
||||
def round(self, divisor):
|
||||
for corner in self.polygon:
|
||||
corner[0] = int(corner[0] / divisor) * divisor
|
||||
corner[1] = int(corner[1] / divisor) * divisor
|
||||
|
||||
def fit_to_bounds(self, bounds):
|
||||
new_corners = copy.deepcopy(self.polygon)
|
||||
for corner in new_corners:
|
||||
corner[0] = max(min(corner[0], bounds[2]), bounds[0])
|
||||
corner[1] = max(min(corner[1], bounds[3]), bounds[1])
|
||||
self.polygon = new_corners
|
||||
|
||||
def expand(self, x_margin: float, y_margin: float):
|
||||
new_polygon = []
|
||||
x_margin = x_margin * self.width
|
||||
y_margin = y_margin * self.height
|
||||
for idx, poly in enumerate(self.polygon):
|
||||
if idx == 0:
|
||||
new_polygon.append([int(poly[0] - x_margin), int(poly[1] - y_margin)])
|
||||
elif idx == 1:
|
||||
new_polygon.append([int(poly[0] + x_margin), int(poly[1] - y_margin)])
|
||||
elif idx == 2:
|
||||
new_polygon.append([int(poly[0] + x_margin), int(poly[1] + y_margin)])
|
||||
elif idx == 3:
|
||||
new_polygon.append([int(poly[0] - x_margin), int(poly[1] + y_margin)])
|
||||
self.polygon = new_polygon
|
||||
|
||||
def intersection_area(self, other, x_margin=0, y_margin=0):
|
||||
x_overlap = self.x_overlap(other, x_margin)
|
||||
y_overlap = self.y_overlap(other, y_margin)
|
||||
return x_overlap * y_overlap
|
||||
|
||||
def x_overlap(self, other, x_margin=0):
|
||||
return max(
|
||||
0,
|
||||
min(self.bbox[2] + x_margin, other.bbox[2] + x_margin)
|
||||
- max(self.bbox[0] - x_margin, other.bbox[0] - x_margin),
|
||||
)
|
||||
|
||||
def y_overlap(self, other, y_margin=0):
|
||||
return max(
|
||||
0,
|
||||
min(self.bbox[3] + y_margin, other.bbox[3] + y_margin)
|
||||
- max(self.bbox[1] - y_margin, other.bbox[1] - y_margin),
|
||||
)
|
||||
|
||||
@property
|
||||
def center(self):
|
||||
return [(self.bbox[0] + self.bbox[2]) / 2, (self.bbox[1] + self.bbox[3]) / 2]
|
||||
|
||||
def __hash__(self):
|
||||
return hash(tuple(self.bbox))
|
||||
@@ -0,0 +1,57 @@
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
|
||||
from surya.common.load import ModelLoader
|
||||
from surya.settings import settings
|
||||
|
||||
|
||||
class BasePredictor:
|
||||
model_loader_cls = ModelLoader
|
||||
batch_size: Optional[int] = None
|
||||
default_batch_sizes = {"cpu": 1, "mps": 1, "cuda": 1}
|
||||
torch_dtype = settings.MODEL_DTYPE
|
||||
|
||||
@property
|
||||
def disable_tqdm(self) -> bool:
|
||||
return self._disable_tqdm
|
||||
|
||||
@disable_tqdm.setter
|
||||
def disable_tqdm(self, value: bool) -> None:
|
||||
self._disable_tqdm = bool(value)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
checkpoint: Optional[str] = None,
|
||||
device: torch.device | str | None = settings.TORCH_DEVICE_MODEL,
|
||||
dtype: Optional[torch.dtype | str] = None,
|
||||
attention_implementation: Optional[str] = None,
|
||||
):
|
||||
if dtype is None:
|
||||
dtype = self.torch_dtype
|
||||
|
||||
loader = self.model_loader_cls(checkpoint)
|
||||
self.model = loader.model(device, dtype, attention_implementation)
|
||||
self.processor = loader.processor()
|
||||
self._disable_tqdm = settings.DISABLE_TQDM
|
||||
|
||||
def to(self, device_dtype: torch.device | str | None = None):
|
||||
if hasattr(self, "model") and self.model:
|
||||
self.model.to(device_dtype)
|
||||
return
|
||||
# Predictors that don't own a torch model (e.g. VLM-backed predictors that
|
||||
# rely on an external server) treat .to() as a no-op.
|
||||
if hasattr(self, "manager") and self.manager is not None:
|
||||
return
|
||||
raise ValueError("Model not loaded")
|
||||
|
||||
def get_batch_size(self):
|
||||
batch_size = self.batch_size
|
||||
if batch_size is None:
|
||||
batch_size = self.default_batch_sizes["cpu"]
|
||||
if settings.TORCH_DEVICE_MODEL in self.default_batch_sizes:
|
||||
batch_size = self.default_batch_sizes[settings.TORCH_DEVICE_MODEL]
|
||||
return batch_size
|
||||
|
||||
def __call__(self, *args, **kwargs):
|
||||
raise NotImplementedError()
|
||||
@@ -0,0 +1,22 @@
|
||||
from typing import Optional
|
||||
|
||||
from transformers import PreTrainedModel
|
||||
from transformers.utils import is_flash_attn_2_available
|
||||
|
||||
|
||||
class SuryaPreTrainedModel(PreTrainedModel):
|
||||
# No-op if we pass attention, so we can set attention however we want in the config
|
||||
def _check_and_adjust_attn_implementation(
|
||||
self, attn_implementation: Optional[str], **kwargs
|
||||
):
|
||||
if attn_implementation is None:
|
||||
try:
|
||||
self._sdpa_can_dispatch(True)
|
||||
attn_implementation = "sdpa"
|
||||
except (ValueError, ImportError):
|
||||
attn_implementation = "eager"
|
||||
|
||||
if self._supports_flash_attn and is_flash_attn_2_available():
|
||||
attn_implementation = "flash_attention_2"
|
||||
|
||||
return attn_implementation
|
||||
@@ -0,0 +1,10 @@
|
||||
"""Vendored, slimmed, detection-only copy of Roboflow's RF-DETR for surya's fast detectors.
|
||||
|
||||
Avoids a runtime dependency on the `rfdetr` package (and its heavy transitive deps:
|
||||
roboflow, rf100vl, albumentations, supervision, peft). Pure PyTorch; runs on cpu/mps/cuda.
|
||||
See `predictor.RFDetrDetector` for the inference entry point.
|
||||
"""
|
||||
|
||||
from surya.common.rfdetr.predictor import RFDetrDetector
|
||||
|
||||
__all__ = ["RFDetrDetector"]
|
||||
@@ -0,0 +1,16 @@
|
||||
# ------------------------------------------------------------------------
|
||||
# RF-DETR
|
||||
# Copyright (c) 2025 Roboflow. All Rights Reserved.
|
||||
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
|
||||
# ------------------------------------------------------------------------
|
||||
# Copied from LW-DETR (https://github.com/Atten4Vis/LW-DETR)
|
||||
# Copyright (c) 2024 Baidu. All Rights Reserved.
|
||||
# ------------------------------------------------------------------------
|
||||
# Copied from Conditional DETR (https://github.com/Atten4Vis/ConditionalDETR)
|
||||
# Copyright (c) 2021 Microsoft. All Rights Reserved.
|
||||
# ------------------------------------------------------------------------
|
||||
# Copied from DETR (https://github.com/facebookresearch/detr)
|
||||
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
|
||||
# ------------------------------------------------------------------------
|
||||
|
||||
from surya.common.rfdetr.models.lwdetr import PostProcess, build_model
|
||||
@@ -0,0 +1,105 @@
|
||||
# ------------------------------------------------------------------------
|
||||
# RF-DETR
|
||||
# Copyright (c) 2025 Roboflow. All Rights Reserved.
|
||||
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
|
||||
# ------------------------------------------------------------------------
|
||||
# Copied and modified from LW-DETR (https://github.com/Atten4Vis/LW-DETR)
|
||||
# Copyright (c) 2024 Baidu. All Rights Reserved.
|
||||
# ------------------------------------------------------------------------
|
||||
|
||||
from typing import Callable, Dict, List
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from surya.common.rfdetr.models.backbone.backbone import *
|
||||
from surya.common.rfdetr.models.position_encoding import build_position_encoding
|
||||
from surya.common.rfdetr.util.misc import NestedTensor
|
||||
|
||||
|
||||
class Joiner(nn.Sequential):
|
||||
def __init__(self, backbone, position_embedding):
|
||||
super().__init__(backbone, position_embedding)
|
||||
self._export = False
|
||||
|
||||
def forward(self, tensor_list: NestedTensor):
|
||||
""" """
|
||||
x = self[0](tensor_list)
|
||||
pos = []
|
||||
for x_ in x:
|
||||
pos.append(self[1](x_, align_dim_orders=False).to(x_.tensors.dtype))
|
||||
return x, pos
|
||||
|
||||
def export(self):
|
||||
self._export = True
|
||||
self._forward_origin = self.forward
|
||||
self.forward = self.forward_export
|
||||
for name, m in self.named_modules():
|
||||
if hasattr(m, "export") and isinstance(m.export, Callable) and hasattr(m, "_export") and not m._export:
|
||||
m.export()
|
||||
|
||||
def forward_export(self, inputs: torch.Tensor):
|
||||
feats, masks = self[0](inputs)
|
||||
poss = []
|
||||
for feat, mask in zip(feats, masks):
|
||||
poss.append(self[1](mask, align_dim_orders=False).to(feat.dtype))
|
||||
return feats, None, poss
|
||||
|
||||
|
||||
def build_backbone(
|
||||
encoder,
|
||||
vit_encoder_num_layers,
|
||||
pretrained_encoder,
|
||||
window_block_indexes,
|
||||
drop_path,
|
||||
out_channels,
|
||||
out_feature_indexes,
|
||||
projector_scale,
|
||||
use_cls_token,
|
||||
hidden_dim,
|
||||
position_embedding,
|
||||
freeze_encoder,
|
||||
layer_norm,
|
||||
target_shape,
|
||||
rms_norm,
|
||||
backbone_lora,
|
||||
force_no_pretrain,
|
||||
gradient_checkpointing,
|
||||
load_dinov2_weights,
|
||||
patch_size,
|
||||
num_windows,
|
||||
positional_encoding_size,
|
||||
):
|
||||
"""
|
||||
Useful args:
|
||||
- encoder: encoder name
|
||||
- lr_encoder:
|
||||
- dilation
|
||||
- use_checkpoint: for swin only for now
|
||||
|
||||
"""
|
||||
position_embedding = build_position_encoding(hidden_dim, position_embedding)
|
||||
|
||||
backbone = Backbone(
|
||||
encoder,
|
||||
pretrained_encoder,
|
||||
window_block_indexes=window_block_indexes,
|
||||
drop_path=drop_path,
|
||||
out_channels=out_channels,
|
||||
out_feature_indexes=out_feature_indexes,
|
||||
projector_scale=projector_scale,
|
||||
use_cls_token=use_cls_token,
|
||||
layer_norm=layer_norm,
|
||||
freeze_encoder=freeze_encoder,
|
||||
target_shape=target_shape,
|
||||
rms_norm=rms_norm,
|
||||
backbone_lora=backbone_lora,
|
||||
gradient_checkpointing=gradient_checkpointing,
|
||||
load_dinov2_weights=load_dinov2_weights,
|
||||
patch_size=patch_size,
|
||||
num_windows=num_windows,
|
||||
positional_encoding_size=positional_encoding_size,
|
||||
)
|
||||
|
||||
model = Joiner(backbone, position_embedding)
|
||||
return model
|
||||
@@ -0,0 +1,201 @@
|
||||
# ------------------------------------------------------------------------
|
||||
# RF-DETR
|
||||
# Copyright (c) 2025 Roboflow. All Rights Reserved.
|
||||
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
|
||||
# ------------------------------------------------------------------------
|
||||
# Copied and modified from LW-DETR (https://github.com/Atten4Vis/LW-DETR)
|
||||
# Copyright (c) 2024 Baidu. All Rights Reserved.
|
||||
# ------------------------------------------------------------------------
|
||||
# Modified from Conditional DETR (https://github.com/Atten4Vis/ConditionalDETR)
|
||||
# Copyright (c) 2021 Microsoft. All Rights Reserved.
|
||||
# ------------------------------------------------------------------------
|
||||
# Copied from DETR (https://github.com/facebookresearch/detr)
|
||||
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
|
||||
# ------------------------------------------------------------------------
|
||||
|
||||
"""
|
||||
Backbone modules.
|
||||
"""
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from surya.common.rfdetr.models.backbone.base import BackboneBase
|
||||
from surya.common.rfdetr.models.backbone.dinov2 import DinoV2
|
||||
from surya.common.rfdetr.models.backbone.projector import MultiScaleProjector
|
||||
from surya.common.rfdetr.util.logger import get_logger
|
||||
from surya.common.rfdetr.util.misc import NestedTensor
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
__all__ = ["Backbone"]
|
||||
|
||||
|
||||
class Backbone(BackboneBase):
|
||||
"""backbone."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: str,
|
||||
pretrained_encoder: str = None,
|
||||
window_block_indexes: list = None,
|
||||
drop_path=0.0,
|
||||
out_channels=256,
|
||||
out_feature_indexes: list = None,
|
||||
projector_scale: list = None,
|
||||
use_cls_token: bool = False,
|
||||
freeze_encoder: bool = False,
|
||||
layer_norm: bool = False,
|
||||
target_shape: tuple[int, int] = (640, 640),
|
||||
rms_norm: bool = False,
|
||||
backbone_lora: bool = False,
|
||||
gradient_checkpointing: bool = False,
|
||||
load_dinov2_weights: bool = True,
|
||||
patch_size: int = 14,
|
||||
num_windows: int = 4,
|
||||
positional_encoding_size: int = 0,
|
||||
):
|
||||
super().__init__()
|
||||
# an example name here would be "dinov2_base" or "dinov2_registers_windowed_base"
|
||||
# if "registers" is in the name, then use_registers is set to True, otherwise it is set to False
|
||||
# similarly, if "windowed" is in the name, then use_windowed_attn is set to True, otherwise it is set to False
|
||||
# the last part of the name should be the size
|
||||
# and the start should be dinov2
|
||||
name_parts = name.split("_")
|
||||
assert name_parts[0] == "dinov2"
|
||||
# name_parts[-1]
|
||||
use_registers = False
|
||||
if "registers" in name_parts:
|
||||
use_registers = True
|
||||
name_parts.remove("registers")
|
||||
use_windowed_attn = False
|
||||
if "windowed" in name_parts:
|
||||
use_windowed_attn = True
|
||||
name_parts.remove("windowed")
|
||||
assert len(name_parts) == 2, (
|
||||
"name should be dinov2, then either registers, windowed, both, or none, then the size"
|
||||
)
|
||||
self.encoder = DinoV2(
|
||||
size=name_parts[-1],
|
||||
out_feature_indexes=out_feature_indexes,
|
||||
shape=target_shape,
|
||||
use_registers=use_registers,
|
||||
use_windowed_attn=use_windowed_attn,
|
||||
gradient_checkpointing=gradient_checkpointing,
|
||||
load_dinov2_weights=load_dinov2_weights,
|
||||
patch_size=patch_size,
|
||||
num_windows=num_windows,
|
||||
positional_encoding_size=positional_encoding_size,
|
||||
drop_path_rate=drop_path,
|
||||
)
|
||||
# build encoder + projector as backbone module
|
||||
if freeze_encoder:
|
||||
for param in self.encoder.parameters():
|
||||
param.requires_grad = False
|
||||
|
||||
self.projector_scale = projector_scale
|
||||
assert len(self.projector_scale) > 0
|
||||
# x[0]
|
||||
assert sorted(self.projector_scale) == self.projector_scale, (
|
||||
"only support projector scale P3/P4/P5/P6 in ascending order."
|
||||
)
|
||||
level2scalefactor = dict(P3=2.0, P4=1.0, P5=0.5, P6=0.25)
|
||||
scale_factors = [level2scalefactor[lvl] for lvl in self.projector_scale]
|
||||
|
||||
self.projector = MultiScaleProjector(
|
||||
in_channels=self.encoder._out_feature_channels,
|
||||
out_channels=out_channels,
|
||||
scale_factors=scale_factors,
|
||||
layer_norm=layer_norm,
|
||||
rms_norm=rms_norm,
|
||||
)
|
||||
|
||||
self._export = False
|
||||
|
||||
def export(self):
|
||||
self._export = True
|
||||
self._forward_origin = self.forward
|
||||
self.forward = self.forward_export
|
||||
|
||||
def forward(self, tensor_list: NestedTensor):
|
||||
""" """
|
||||
# (H, W, B, C)
|
||||
feats = self.encoder(tensor_list.tensors)
|
||||
feats = self.projector(feats)
|
||||
# x: [(B, C, H, W)]
|
||||
out = []
|
||||
for feat in feats:
|
||||
m = tensor_list.mask
|
||||
assert m is not None
|
||||
mask = F.interpolate(m[None].float(), size=feat.shape[-2:]).to(torch.bool)[0]
|
||||
out.append(NestedTensor(feat, mask))
|
||||
return out
|
||||
|
||||
def forward_export(self, tensors: torch.Tensor):
|
||||
feats = self.encoder(tensors)
|
||||
feats = self.projector(feats)
|
||||
out_feats = []
|
||||
out_masks = []
|
||||
for feat in feats:
|
||||
# x: [(B, C, H, W)]
|
||||
b, _, h, w = feat.shape
|
||||
out_masks.append(torch.zeros((b, h, w), dtype=torch.bool, device=feat.device))
|
||||
out_feats.append(feat)
|
||||
return out_feats, out_masks
|
||||
|
||||
def get_named_param_lr_pairs(self, args, prefix: str = "backbone.0"):
|
||||
num_layers = args.out_feature_indexes[-1] + 1
|
||||
backbone_key = "backbone.0.encoder"
|
||||
named_param_lr_pairs = {}
|
||||
for n, p in self.named_parameters():
|
||||
n = prefix + "." + n
|
||||
if backbone_key in n and p.requires_grad:
|
||||
lr = (
|
||||
args.lr_encoder
|
||||
* get_dinov2_lr_decay_rate(
|
||||
n,
|
||||
lr_decay_rate=args.lr_vit_layer_decay,
|
||||
num_layers=num_layers,
|
||||
)
|
||||
* args.lr_component_decay**2
|
||||
)
|
||||
wd = args.weight_decay * get_dinov2_weight_decay_rate(n)
|
||||
named_param_lr_pairs[n] = {
|
||||
"params": p,
|
||||
"lr": lr,
|
||||
"weight_decay": wd,
|
||||
}
|
||||
return named_param_lr_pairs
|
||||
|
||||
|
||||
def get_dinov2_lr_decay_rate(name, lr_decay_rate=1.0, num_layers=12):
|
||||
"""
|
||||
Calculate lr decay rate for different ViT blocks.
|
||||
|
||||
Args:
|
||||
name (string): parameter name.
|
||||
lr_decay_rate (float): base lr decay rate.
|
||||
num_layers (int): number of ViT blocks.
|
||||
Returns:
|
||||
lr decay rate for the given parameter.
|
||||
"""
|
||||
layer_id = num_layers + 1
|
||||
if name.startswith("backbone"):
|
||||
if "embeddings" in name:
|
||||
layer_id = 0
|
||||
elif ".layer." in name and ".residual." not in name:
|
||||
layer_id = int(name[name.find(".layer.") :].split(".")[2]) + 1
|
||||
return lr_decay_rate ** (num_layers + 1 - layer_id)
|
||||
|
||||
|
||||
def get_dinov2_weight_decay_rate(name, weight_decay_rate=1.0):
|
||||
if (
|
||||
("gamma" in name)
|
||||
or ("pos_embed" in name)
|
||||
or ("rel_pos" in name)
|
||||
or ("bias" in name)
|
||||
or ("norm" in name)
|
||||
or ("embeddings" in name)
|
||||
):
|
||||
weight_decay_rate = 0.0
|
||||
return weight_decay_rate
|
||||
@@ -0,0 +1,18 @@
|
||||
# ------------------------------------------------------------------------
|
||||
# RF-DETR
|
||||
# Copyright (c) 2025 Roboflow. All Rights Reserved.
|
||||
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
|
||||
# ------------------------------------------------------------------------
|
||||
# Copied and modified from LW-DETR (https://github.com/Atten4Vis/LW-DETR)
|
||||
# Copyright (c) 2024 Baidu. All Rights Reserved.
|
||||
# ------------------------------------------------------------------------
|
||||
|
||||
from torch import nn
|
||||
|
||||
|
||||
class BackboneBase(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
def get_named_param_lr_pairs(self, args, prefix: str):
|
||||
raise NotImplementedError
|
||||
@@ -0,0 +1,220 @@
|
||||
# ------------------------------------------------------------------------
|
||||
# RF-DETR
|
||||
# Copyright (c) 2025 Roboflow. All Rights Reserved.
|
||||
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
|
||||
# ------------------------------------------------------------------------
|
||||
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import types
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from transformers import AutoBackbone
|
||||
|
||||
from surya.common.rfdetr.models.backbone.dinov2_with_windowed_attn import WindowedDinov2WithRegistersBackbone, \
|
||||
WindowedDinov2WithRegistersConfig
|
||||
|
||||
from surya.common.rfdetr.util.logger import get_logger
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
size_to_width = {
|
||||
"tiny": 192,
|
||||
"small": 384,
|
||||
"base": 768,
|
||||
"large": 1024,
|
||||
}
|
||||
|
||||
size_to_config = {
|
||||
"small": "dinov2_small.json",
|
||||
"base": "dinov2_base.json",
|
||||
"large": "dinov2_large.json",
|
||||
}
|
||||
|
||||
size_to_config_with_registers = {
|
||||
"small": "dinov2_with_registers_small.json",
|
||||
"base": "dinov2_with_registers_base.json",
|
||||
"large": "dinov2_with_registers_large.json",
|
||||
}
|
||||
|
||||
|
||||
def get_config(size, use_registers):
|
||||
config_dict = size_to_config_with_registers if use_registers else size_to_config
|
||||
current_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
configs_dir = os.path.join(current_dir, "dinov2_configs")
|
||||
config_path = os.path.join(configs_dir, config_dict[size])
|
||||
with open(config_path, "r") as f:
|
||||
dino_config = json.load(f)
|
||||
return dino_config
|
||||
|
||||
|
||||
class DinoV2(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
shape=(640, 640),
|
||||
out_feature_indexes=[2, 4, 5, 9],
|
||||
size="base",
|
||||
use_registers=True,
|
||||
use_windowed_attn=True,
|
||||
gradient_checkpointing=False,
|
||||
load_dinov2_weights=True,
|
||||
patch_size=14,
|
||||
num_windows=4,
|
||||
positional_encoding_size=37,
|
||||
drop_path_rate=0.0,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
name = f"facebook/dinov2-with-registers-{size}" if use_registers else f"facebook/dinov2-{size}"
|
||||
|
||||
self.shape = shape
|
||||
self.patch_size = patch_size
|
||||
self.num_windows = num_windows
|
||||
|
||||
# Create the encoder
|
||||
|
||||
if not use_windowed_attn:
|
||||
assert not gradient_checkpointing, "Gradient checkpointing is not supported for non-windowed attention"
|
||||
assert load_dinov2_weights, "Using non-windowed attention requires loading dinov2 weights from hub"
|
||||
if drop_path_rate > 0.0:
|
||||
logger.warning(
|
||||
"drop_path_rate > 0.0 is not supported for non-windowed DinoV2 backbones. "
|
||||
"drop_path will be ignored."
|
||||
)
|
||||
self.encoder = AutoBackbone.from_pretrained(
|
||||
name,
|
||||
out_features=[f"stage{i}" for i in out_feature_indexes],
|
||||
return_dict=False,
|
||||
)
|
||||
else:
|
||||
window_block_indexes = set(range(out_feature_indexes[-1] + 1))
|
||||
window_block_indexes.difference_update(out_feature_indexes)
|
||||
window_block_indexes = list(window_block_indexes)
|
||||
|
||||
dino_config = get_config(size, use_registers)
|
||||
|
||||
dino_config["return_dict"] = False
|
||||
dino_config["out_features"] = [f"stage{i}" for i in out_feature_indexes]
|
||||
dino_config["drop_path_rate"] = drop_path_rate
|
||||
|
||||
implied_resolution = positional_encoding_size * patch_size
|
||||
|
||||
if implied_resolution != dino_config["image_size"]:
|
||||
logger.warning(
|
||||
"Using a different number of positional encodings than DINOv2, which means we're not loading DINOv2 backbone weights. This is not a problem if finetuning a pretrained RF-DETR model."
|
||||
)
|
||||
dino_config["image_size"] = implied_resolution
|
||||
load_dinov2_weights = False
|
||||
|
||||
if patch_size != 14:
|
||||
logger.warning(
|
||||
f"Using patch size {patch_size} instead of 14, which means we're not loading DINOv2 backbone weights. This is not a problem if finetuning a pretrained RF-DETR model."
|
||||
)
|
||||
dino_config["patch_size"] = patch_size
|
||||
load_dinov2_weights = False
|
||||
|
||||
if use_registers:
|
||||
windowed_dino_config = WindowedDinov2WithRegistersConfig(
|
||||
**dino_config,
|
||||
num_windows=num_windows,
|
||||
window_block_indexes=window_block_indexes,
|
||||
gradient_checkpointing=gradient_checkpointing,
|
||||
)
|
||||
else:
|
||||
windowed_dino_config = WindowedDinov2WithRegistersConfig(
|
||||
**dino_config,
|
||||
num_windows=num_windows,
|
||||
window_block_indexes=window_block_indexes,
|
||||
num_register_tokens=0,
|
||||
gradient_checkpointing=gradient_checkpointing,
|
||||
)
|
||||
self.encoder = (
|
||||
WindowedDinov2WithRegistersBackbone.from_pretrained(
|
||||
name,
|
||||
config=windowed_dino_config,
|
||||
)
|
||||
if load_dinov2_weights
|
||||
else WindowedDinov2WithRegistersBackbone(windowed_dino_config)
|
||||
)
|
||||
|
||||
self._out_feature_channels = [size_to_width[size]] * len(out_feature_indexes)
|
||||
self._export = False
|
||||
|
||||
def export(self):
|
||||
if self._export:
|
||||
return
|
||||
self._export = True
|
||||
shape = self.shape
|
||||
|
||||
def make_new_interpolated_pos_encoding(position_embeddings, patch_size, height, width):
|
||||
|
||||
num_positions = position_embeddings.shape[1] - 1
|
||||
dim = position_embeddings.shape[-1]
|
||||
height = height // patch_size
|
||||
width = width // patch_size
|
||||
|
||||
class_pos_embed = position_embeddings[:, 0]
|
||||
patch_pos_embed = position_embeddings[:, 1:]
|
||||
|
||||
# Reshape and permute
|
||||
patch_pos_embed = patch_pos_embed.reshape(
|
||||
1, int(math.sqrt(num_positions)), int(math.sqrt(num_positions)), dim
|
||||
)
|
||||
patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2)
|
||||
|
||||
# Use bilinear interpolation without antialias
|
||||
patch_pos_embed = F.interpolate(
|
||||
patch_pos_embed,
|
||||
size=(height, width),
|
||||
mode="bicubic",
|
||||
align_corners=False,
|
||||
antialias=True,
|
||||
)
|
||||
|
||||
# Reshape back
|
||||
patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).reshape(1, -1, dim)
|
||||
return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1)
|
||||
|
||||
# If the shape of self.encoder.embeddings.position_embeddings
|
||||
# matches the shape of your new tensor, use copy_:
|
||||
with torch.no_grad():
|
||||
new_positions = make_new_interpolated_pos_encoding(
|
||||
self.encoder.embeddings.position_embeddings,
|
||||
self.encoder.config.patch_size,
|
||||
shape[0],
|
||||
shape[1],
|
||||
)
|
||||
# Create a new Parameter with the new size
|
||||
old_interpolate_pos_encoding = self.encoder.embeddings.interpolate_pos_encoding
|
||||
|
||||
def new_interpolate_pos_encoding(self_mod, embeddings, height, width):
|
||||
num_patches = embeddings.shape[1] - 1
|
||||
num_positions = self_mod.position_embeddings.shape[1] - 1
|
||||
if num_patches == num_positions and height == width:
|
||||
return self_mod.position_embeddings
|
||||
return old_interpolate_pos_encoding(embeddings, height, width)
|
||||
|
||||
self.encoder.embeddings.position_embeddings = nn.Parameter(new_positions)
|
||||
self.encoder.embeddings.interpolate_pos_encoding = types.MethodType(
|
||||
new_interpolate_pos_encoding, self.encoder.embeddings
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
block_size = self.patch_size * self.num_windows
|
||||
assert x.shape[2] % block_size == 0 and x.shape[3] % block_size == 0, (
|
||||
f"Backbone requires input shape to be divisible by {block_size}, but got {x.shape}"
|
||||
)
|
||||
x = self.encoder(x)
|
||||
return list(x[0])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
model = DinoV2()
|
||||
model.export()
|
||||
x = torch.randn(1, 3, 640, 640)
|
||||
logger.info(model(x))
|
||||
for j in model(x):
|
||||
logger.info(j.shape)
|
||||
@@ -0,0 +1,24 @@
|
||||
{
|
||||
"architectures": [
|
||||
"Dinov2Model"
|
||||
],
|
||||
"attention_probs_dropout_prob": 0.0,
|
||||
"drop_path_rate": 0.0,
|
||||
"hidden_act": "gelu",
|
||||
"hidden_dropout_prob": 0.0,
|
||||
"hidden_size": 768,
|
||||
"image_size": 518,
|
||||
"initializer_range": 0.02,
|
||||
"layer_norm_eps": 1e-06,
|
||||
"layerscale_value": 1.0,
|
||||
"mlp_ratio": 4,
|
||||
"model_type": "dinov2",
|
||||
"num_attention_heads": 12,
|
||||
"num_channels": 3,
|
||||
"num_hidden_layers": 12,
|
||||
"patch_size": 14,
|
||||
"qkv_bias": true,
|
||||
"torch_dtype": "float32",
|
||||
"transformers_version": "4.31.0.dev0",
|
||||
"use_swiglu_ffn": false
|
||||
}
|
||||
@@ -0,0 +1,24 @@
|
||||
{
|
||||
"architectures": [
|
||||
"Dinov2Model"
|
||||
],
|
||||
"attention_probs_dropout_prob": 0.0,
|
||||
"drop_path_rate": 0.0,
|
||||
"hidden_act": "gelu",
|
||||
"hidden_dropout_prob": 0.0,
|
||||
"hidden_size": 1024,
|
||||
"image_size": 518,
|
||||
"initializer_range": 0.02,
|
||||
"layer_norm_eps": 1e-06,
|
||||
"layerscale_value": 1.0,
|
||||
"mlp_ratio": 4,
|
||||
"model_type": "dinov2",
|
||||
"num_attention_heads": 16,
|
||||
"num_channels": 3,
|
||||
"num_hidden_layers": 24,
|
||||
"patch_size": 14,
|
||||
"qkv_bias": true,
|
||||
"torch_dtype": "float32",
|
||||
"transformers_version": "4.31.0.dev0",
|
||||
"use_swiglu_ffn": false
|
||||
}
|
||||
@@ -0,0 +1,24 @@
|
||||
{
|
||||
"architectures": [
|
||||
"Dinov2Model"
|
||||
],
|
||||
"attention_probs_dropout_prob": 0.0,
|
||||
"drop_path_rate": 0.0,
|
||||
"hidden_act": "gelu",
|
||||
"hidden_dropout_prob": 0.0,
|
||||
"hidden_size": 384,
|
||||
"image_size": 518,
|
||||
"initializer_range": 0.02,
|
||||
"layer_norm_eps": 1e-06,
|
||||
"layerscale_value": 1.0,
|
||||
"mlp_ratio": 4,
|
||||
"model_type": "dinov2",
|
||||
"num_attention_heads": 6,
|
||||
"num_channels": 3,
|
||||
"num_hidden_layers": 12,
|
||||
"patch_size": 14,
|
||||
"qkv_bias": true,
|
||||
"torch_dtype": "float32",
|
||||
"transformers_version": "4.32.0.dev0",
|
||||
"use_swiglu_ffn": false
|
||||
}
|
||||
@@ -0,0 +1,50 @@
|
||||
{
|
||||
"apply_layernorm": true,
|
||||
"architectures": [
|
||||
"Dinov2WithRegistersModel"
|
||||
],
|
||||
"attention_probs_dropout_prob": 0.0,
|
||||
"drop_path_rate": 0.0,
|
||||
"hidden_act": "gelu",
|
||||
"hidden_dropout_prob": 0.0,
|
||||
"hidden_size": 768,
|
||||
"image_size": 518,
|
||||
"initializer_range": 0.02,
|
||||
"interpolate_antialias": true,
|
||||
"interpolate_offset": 0.0,
|
||||
"layer_norm_eps": 1e-06,
|
||||
"layerscale_value": 1.0,
|
||||
"mlp_ratio": 4,
|
||||
"model_type": "dinov2_with_registers",
|
||||
"num_attention_heads": 12,
|
||||
"num_channels": 3,
|
||||
"num_hidden_layers": 12,
|
||||
"num_register_tokens": 4,
|
||||
"out_features": [
|
||||
"stage12"
|
||||
],
|
||||
"out_indices": [
|
||||
12
|
||||
],
|
||||
"patch_size": 14,
|
||||
"qkv_bias": true,
|
||||
"reshape_hidden_states": true,
|
||||
"stage_names": [
|
||||
"stem",
|
||||
"stage1",
|
||||
"stage2",
|
||||
"stage3",
|
||||
"stage4",
|
||||
"stage5",
|
||||
"stage6",
|
||||
"stage7",
|
||||
"stage8",
|
||||
"stage9",
|
||||
"stage10",
|
||||
"stage11",
|
||||
"stage12"
|
||||
],
|
||||
"torch_dtype": "float32",
|
||||
"transformers_version": "4.48.0.dev0",
|
||||
"use_swiglu_ffn": false
|
||||
}
|
||||
@@ -0,0 +1,50 @@
|
||||
{
|
||||
"apply_layernorm": true,
|
||||
"architectures": [
|
||||
"Dinov2WithRegistersModel"
|
||||
],
|
||||
"attention_probs_dropout_prob": 0.0,
|
||||
"drop_path_rate": 0.0,
|
||||
"hidden_act": "gelu",
|
||||
"hidden_dropout_prob": 0.0,
|
||||
"hidden_size": 1024,
|
||||
"image_size": 518,
|
||||
"initializer_range": 0.02,
|
||||
"interpolate_antialias": true,
|
||||
"interpolate_offset": 0.0,
|
||||
"layer_norm_eps": 1e-06,
|
||||
"layerscale_value": 1.0,
|
||||
"mlp_ratio": 4,
|
||||
"model_type": "dinov2_with_registers",
|
||||
"num_attention_heads": 16,
|
||||
"num_channels": 3,
|
||||
"num_hidden_layers": 24,
|
||||
"num_register_tokens": 4,
|
||||
"out_features": [
|
||||
"stage12"
|
||||
],
|
||||
"out_indices": [
|
||||
12
|
||||
],
|
||||
"patch_size": 14,
|
||||
"qkv_bias": true,
|
||||
"reshape_hidden_states": true,
|
||||
"stage_names": [
|
||||
"stem",
|
||||
"stage1",
|
||||
"stage2",
|
||||
"stage3",
|
||||
"stage4",
|
||||
"stage5",
|
||||
"stage6",
|
||||
"stage7",
|
||||
"stage8",
|
||||
"stage9",
|
||||
"stage10",
|
||||
"stage11",
|
||||
"stage12"
|
||||
],
|
||||
"torch_dtype": "float32",
|
||||
"transformers_version": "4.48.0.dev0",
|
||||
"use_swiglu_ffn": false
|
||||
}
|
||||
@@ -0,0 +1,50 @@
|
||||
{
|
||||
"apply_layernorm": true,
|
||||
"architectures": [
|
||||
"Dinov2WithRegistersModel"
|
||||
],
|
||||
"attention_probs_dropout_prob": 0.0,
|
||||
"drop_path_rate": 0.0,
|
||||
"hidden_act": "gelu",
|
||||
"hidden_dropout_prob": 0.0,
|
||||
"hidden_size": 384,
|
||||
"image_size": 518,
|
||||
"initializer_range": 0.02,
|
||||
"interpolate_antialias": true,
|
||||
"interpolate_offset": 0.0,
|
||||
"layer_norm_eps": 1e-06,
|
||||
"layerscale_value": 1.0,
|
||||
"mlp_ratio": 4,
|
||||
"model_type": "dinov2_with_registers",
|
||||
"num_attention_heads": 6,
|
||||
"num_channels": 3,
|
||||
"num_hidden_layers": 12,
|
||||
"num_register_tokens": 4,
|
||||
"out_features": [
|
||||
"stage12"
|
||||
],
|
||||
"out_indices": [
|
||||
12
|
||||
],
|
||||
"patch_size": 14,
|
||||
"qkv_bias": true,
|
||||
"reshape_hidden_states": true,
|
||||
"stage_names": [
|
||||
"stem",
|
||||
"stage1",
|
||||
"stage2",
|
||||
"stage3",
|
||||
"stage4",
|
||||
"stage5",
|
||||
"stage6",
|
||||
"stage7",
|
||||
"stage8",
|
||||
"stage9",
|
||||
"stage10",
|
||||
"stage11",
|
||||
"stage12"
|
||||
],
|
||||
"torch_dtype": "float32",
|
||||
"transformers_version": "4.48.0.dev0",
|
||||
"use_swiglu_ffn": false
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,315 @@
|
||||
# ------------------------------------------------------------------------
|
||||
# RF-DETR
|
||||
# Copyright (c) 2025 Roboflow. All Rights Reserved.
|
||||
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
|
||||
# ------------------------------------------------------------------------
|
||||
# Copied and modified from LW-DETR (https://github.com/Atten4Vis/LW-DETR)
|
||||
# Copyright (c) 2024 Baidu. All Rights Reserved.
|
||||
# ------------------------------------------------------------------------
|
||||
# Modified from ViTDet (https://github.com/facebookresearch/detectron2/tree/main/projects/ViTDet)
|
||||
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
|
||||
# ------------------------------------------------------------------------
|
||||
|
||||
"""
|
||||
Projector
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class LayerNorm(nn.Module):
|
||||
"""
|
||||
A LayerNorm variant, popularized by Transformers, that performs point-wise mean and
|
||||
variance normalization over the channel dimension for inputs that have shape
|
||||
(batch_size, channels, height, width).
|
||||
https://github.com/facebookresearch/ConvNeXt/blob/d1fa8f6fef0a165b27399986cc2bdacc92777e40/models/convnext.py#L119
|
||||
"""
|
||||
|
||||
def __init__(self, normalized_shape, eps=1e-6):
|
||||
super().__init__()
|
||||
self.weight = nn.Parameter(torch.ones(normalized_shape))
|
||||
self.bias = nn.Parameter(torch.zeros(normalized_shape))
|
||||
self.eps = eps
|
||||
self.normalized_shape = (normalized_shape,)
|
||||
|
||||
def forward(self, x):
|
||||
"""
|
||||
LayerNorm forward
|
||||
TODO: this is a hack to avoid overflow when using fp16
|
||||
"""
|
||||
x = x.permute(0, 2, 3, 1)
|
||||
x = F.layer_norm(x, (x.size(3),), self.weight, self.bias, self.eps)
|
||||
x = x.permute(0, 3, 1, 2)
|
||||
return x
|
||||
|
||||
|
||||
def get_norm(norm, out_channels):
|
||||
"""
|
||||
Args:
|
||||
norm (str or callable): either one of BN, SyncBN, FrozenBN, GN;
|
||||
or a callable that takes a channel number and returns
|
||||
the normalization layer as a nn.Module.
|
||||
Returns:
|
||||
nn.Module or None: the normalization layer
|
||||
"""
|
||||
if norm is None:
|
||||
return None
|
||||
if isinstance(norm, str):
|
||||
if len(norm) == 0:
|
||||
return None
|
||||
norm = {
|
||||
"LN": lambda channels: LayerNorm(channels),
|
||||
}[norm]
|
||||
return norm(out_channels)
|
||||
|
||||
|
||||
def get_activation(name, inplace=False):
|
||||
"""get activation"""
|
||||
if name == "silu":
|
||||
module = nn.SiLU(inplace=inplace)
|
||||
elif name == "relu":
|
||||
module = nn.ReLU(inplace=inplace)
|
||||
elif name in ["LeakyReLU", "leakyrelu", "lrelu"]:
|
||||
module = nn.LeakyReLU(0.1, inplace=inplace)
|
||||
elif name is None:
|
||||
module = nn.Identity()
|
||||
else:
|
||||
raise AttributeError("Unsupported act type: {}".format(name))
|
||||
return module
|
||||
|
||||
|
||||
class ConvX(nn.Module):
|
||||
"""Conv-bn module"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_planes,
|
||||
out_planes,
|
||||
kernel=3,
|
||||
stride=1,
|
||||
groups=1,
|
||||
dilation=1,
|
||||
act="relu",
|
||||
layer_norm=False,
|
||||
rms_norm=False,
|
||||
):
|
||||
super(ConvX, self).__init__()
|
||||
if not isinstance(kernel, tuple):
|
||||
kernel = (kernel, kernel)
|
||||
padding = (kernel[0] // 2, kernel[1] // 2)
|
||||
self.conv = nn.Conv2d(
|
||||
in_planes,
|
||||
out_planes,
|
||||
kernel_size=kernel,
|
||||
stride=stride,
|
||||
padding=padding,
|
||||
groups=groups,
|
||||
dilation=dilation,
|
||||
bias=False,
|
||||
)
|
||||
if rms_norm:
|
||||
self.bn = nn.RMSNorm(out_planes)
|
||||
else:
|
||||
self.bn = get_norm("LN", out_planes) if layer_norm else nn.BatchNorm2d(out_planes)
|
||||
self.act = get_activation(act, inplace=True)
|
||||
|
||||
def forward(self, x):
|
||||
"""forward"""
|
||||
out = self.act(self.bn(self.conv(x.contiguous())))
|
||||
return out
|
||||
|
||||
|
||||
class Bottleneck(nn.Module):
|
||||
"""Standard bottleneck."""
|
||||
|
||||
def __init__(self, c1, c2, shortcut=True, g=1, k=(3, 3), e=0.5, act="silu", layer_norm=False, rms_norm=False):
|
||||
"""ch_in, ch_out, shortcut, groups, kernels, expand"""
|
||||
super().__init__()
|
||||
c_ = int(c2 * e) # hidden channels
|
||||
self.cv1 = ConvX(c1, c_, k[0], 1, act=act, layer_norm=layer_norm, rms_norm=rms_norm)
|
||||
self.cv2 = ConvX(c_, c2, k[1], 1, groups=g, act=act, layer_norm=layer_norm, rms_norm=rms_norm)
|
||||
self.add = shortcut and c1 == c2
|
||||
|
||||
def forward(self, x):
|
||||
"""'forward()' applies the YOLOv5 FPN to input data."""
|
||||
return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))
|
||||
|
||||
|
||||
class C2f(nn.Module):
|
||||
"""Faster Implementation of CSP Bottleneck with 2 convolutions."""
|
||||
|
||||
def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5, act="silu", layer_norm=False, rms_norm=False):
|
||||
"""ch_in, ch_out, number, shortcut, groups, expansion"""
|
||||
super().__init__()
|
||||
self.c = int(c2 * e) # hidden channels
|
||||
self.cv1 = ConvX(c1, 2 * self.c, 1, 1, act=act, layer_norm=layer_norm, rms_norm=rms_norm)
|
||||
self.cv2 = ConvX(
|
||||
(2 + n) * self.c, c2, 1, act=act, layer_norm=layer_norm, rms_norm=rms_norm
|
||||
) # optional act=FReLU(c2)
|
||||
self.m = nn.ModuleList(
|
||||
Bottleneck(self.c, self.c, shortcut, g, k=(3, 3), e=1.0, act=act, layer_norm=layer_norm, rms_norm=rms_norm)
|
||||
for _ in range(n)
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
"""Forward pass using split() instead of chunk()."""
|
||||
y = list(self.cv1(x).split((self.c, self.c), 1))
|
||||
y.extend(m(y[-1]) for m in self.m)
|
||||
return self.cv2(torch.cat(y, 1))
|
||||
|
||||
|
||||
class MultiScaleProjector(nn.Module):
|
||||
"""
|
||||
This module implements MultiScaleProjector in :paper:`lwdetr`.
|
||||
It creates pyramid features built on top of the input feature map.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
out_channels,
|
||||
scale_factors,
|
||||
num_blocks=3,
|
||||
layer_norm=False,
|
||||
rms_norm=False,
|
||||
survival_prob=1.0,
|
||||
force_drop_last_n_features=0,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
net (Backbone): module representing the subnetwork backbone.
|
||||
Must be a subclass of :class:`Backbone`.
|
||||
out_channels (int): number of channels in the output feature maps.
|
||||
scale_factors (list[float]): list of scaling factors to upsample or downsample
|
||||
the input features for creating pyramid features.
|
||||
"""
|
||||
super(MultiScaleProjector, self).__init__()
|
||||
|
||||
self.scale_factors = scale_factors
|
||||
self.survival_prob = survival_prob
|
||||
self.force_drop_last_n_features = force_drop_last_n_features
|
||||
|
||||
stages_sampling = []
|
||||
stages = []
|
||||
# use_bias = norm == ""
|
||||
self.use_extra_pool = False
|
||||
for scale in scale_factors:
|
||||
stages_sampling.append([])
|
||||
for in_dim in in_channels:
|
||||
layers = []
|
||||
|
||||
# if in_dim > 512:
|
||||
# layers.append(ConvX(in_dim, in_dim // 2, kernel=1))
|
||||
# in_dim = in_dim // 2
|
||||
|
||||
if scale == 4.0:
|
||||
layers.extend(
|
||||
[
|
||||
nn.ConvTranspose2d(in_dim, in_dim // 2, kernel_size=2, stride=2),
|
||||
get_norm("LN", in_dim // 2),
|
||||
nn.GELU(),
|
||||
nn.ConvTranspose2d(in_dim // 2, in_dim // 4, kernel_size=2, stride=2),
|
||||
]
|
||||
)
|
||||
# in_dim // 4
|
||||
elif scale == 2.0:
|
||||
# a hack to reduce the FLOPs and Params when the dimension of output feature is too large
|
||||
# if in_dim > 512:
|
||||
# layers = [
|
||||
# ConvX(in_dim, in_dim // 2, kernel=1),
|
||||
# nn.ConvTranspose2d(in_dim // 2, in_dim // 4, kernel_size=2, stride=2),
|
||||
# ]
|
||||
# out_dim = in_dim // 4
|
||||
# else:
|
||||
layers.extend(
|
||||
[
|
||||
nn.ConvTranspose2d(in_dim, in_dim // 2, kernel_size=2, stride=2),
|
||||
]
|
||||
)
|
||||
# in_dim // 2
|
||||
elif scale == 1.0:
|
||||
pass
|
||||
elif scale == 0.5:
|
||||
layers.extend(
|
||||
[
|
||||
ConvX(in_dim, in_dim, 3, 2, layer_norm=layer_norm),
|
||||
]
|
||||
)
|
||||
elif scale == 0.25:
|
||||
self.use_extra_pool = True
|
||||
continue
|
||||
else:
|
||||
raise NotImplementedError("Unsupported scale_factor:{}".format(scale))
|
||||
layers = nn.Sequential(*layers)
|
||||
stages_sampling[-1].append(layers)
|
||||
stages_sampling[-1] = nn.ModuleList(stages_sampling[-1])
|
||||
|
||||
in_dim = int(sum(in_channel // max(1, scale) for in_channel in in_channels))
|
||||
layers = [
|
||||
C2f(in_dim, out_channels, num_blocks, layer_norm=layer_norm),
|
||||
get_norm("LN", out_channels),
|
||||
]
|
||||
layers = nn.Sequential(*layers)
|
||||
stages.append(layers)
|
||||
|
||||
self.stages_sampling = nn.ModuleList(stages_sampling)
|
||||
self.stages = nn.ModuleList(stages)
|
||||
|
||||
def forward(self, x):
|
||||
"""
|
||||
Args:
|
||||
x: Tensor of shape (N,C,H,W). H, W must be a multiple of ``self.size_divisibility``.
|
||||
Returns:
|
||||
dict[str->Tensor]:
|
||||
mapping from feature map name to pyramid feature map tensor
|
||||
in high to low resolution order. Returned feature names follow the FPN
|
||||
convention: "p<stage>", where stage has stride = 2 ** stage e.g.,
|
||||
["p2", "p3", ..., "p6"].
|
||||
"""
|
||||
num_features = len(x)
|
||||
if self.survival_prob < 1.0 and self.training:
|
||||
final_drop_prob = 1 - self.survival_prob
|
||||
drop_p = np.random.uniform()
|
||||
for i in range(1, num_features):
|
||||
critical_drop_prob = i * (final_drop_prob / (num_features - 1))
|
||||
if drop_p < critical_drop_prob:
|
||||
x[i][:] = 0
|
||||
elif self.force_drop_last_n_features > 0:
|
||||
for i in range(self.force_drop_last_n_features):
|
||||
# don't do it inplace to ensure the compiler can optimize out the backbone layers
|
||||
x[-(i + 1)] = torch.zeros_like(x[-(i + 1)])
|
||||
|
||||
results = []
|
||||
# x list of len(out_features_indexes)
|
||||
for i, stage in enumerate(self.stages):
|
||||
feat_fuse = []
|
||||
for j, stage_sampling in enumerate(self.stages_sampling[i]):
|
||||
feat_fuse.append(stage_sampling(x[j]))
|
||||
if len(feat_fuse) > 1:
|
||||
feat_fuse = torch.cat(feat_fuse, dim=1)
|
||||
else:
|
||||
feat_fuse = feat_fuse[0]
|
||||
results.append(stage(feat_fuse))
|
||||
if self.use_extra_pool:
|
||||
results.append(F.max_pool2d(results[-1], kernel_size=1, stride=2, padding=0))
|
||||
return results
|
||||
|
||||
|
||||
class SimpleProjector(nn.Module):
|
||||
def __init__(self, in_dim, out_dim, factor_kernel=False):
|
||||
super(SimpleProjector, self).__init__()
|
||||
if not factor_kernel:
|
||||
self.convx1 = ConvX(in_dim, in_dim * 2, layer_norm=True, act="silu")
|
||||
self.convx2 = ConvX(in_dim * 2, out_dim, layer_norm=True, act="silu")
|
||||
else:
|
||||
self.convx1 = ConvX(in_dim, out_dim, kernel=(3, 1), layer_norm=True, act="silu")
|
||||
self.convx2 = ConvX(out_dim, out_dim, kernel=(1, 3), layer_norm=True, act="silu")
|
||||
self.ln = get_norm("LN", out_dim)
|
||||
|
||||
def forward(self, x):
|
||||
"""forward"""
|
||||
out = self.ln(self.convx2(self.convx1(x[0])))
|
||||
return [out]
|
||||
@@ -0,0 +1,481 @@
|
||||
# ------------------------------------------------------------------------
|
||||
# RF-DETR
|
||||
# Copyright (c) 2025 Roboflow. All Rights Reserved.
|
||||
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
|
||||
# ------------------------------------------------------------------------
|
||||
# Copied and modified from LW-DETR (https://github.com/Atten4Vis/LW-DETR)
|
||||
# Copyright (c) 2024 Baidu. All Rights Reserved.
|
||||
# ------------------------------------------------------------------------
|
||||
# Modified from Conditional DETR (https://github.com/Atten4Vis/ConditionalDETR)
|
||||
# Copyright (c) 2021 Microsoft. All Rights Reserved.
|
||||
# ------------------------------------------------------------------------
|
||||
# Modified from DETR (https://github.com/facebookresearch/detr)
|
||||
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
|
||||
# ------------------------------------------------------------------------
|
||||
# Modified from Deformable DETR (https://github.com/fundamentalvision/Deformable-DETR)
|
||||
# Copyright (c) 2020 SenseTime. All Rights Reserved.
|
||||
# ------------------------------------------------------------------------
|
||||
|
||||
"""
|
||||
LW-DETR model and criterion classes
|
||||
"""
|
||||
|
||||
import copy
|
||||
import math
|
||||
from typing import Callable, Optional
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
|
||||
from surya.common.rfdetr.models.backbone import build_backbone
|
||||
from surya.common.rfdetr.models.transformer import build_transformer
|
||||
from surya.common.rfdetr.util import box_ops
|
||||
from surya.common.rfdetr.util.misc import (
|
||||
NestedTensor,
|
||||
nested_tensor_from_tensor_list,
|
||||
)
|
||||
|
||||
|
||||
class LWDETR(nn.Module):
|
||||
"""This is the Group DETR v3 module that performs object detection"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
backbone,
|
||||
transformer,
|
||||
segmentation_head,
|
||||
num_classes,
|
||||
num_queries,
|
||||
aux_loss=False,
|
||||
group_detr=1,
|
||||
two_stage=False,
|
||||
lite_refpoint_refine=False,
|
||||
bbox_reparam=False,
|
||||
):
|
||||
"""Initializes the model.
|
||||
Parameters:
|
||||
backbone: torch module of the backbone to be used. See backbone.py
|
||||
transformer: torch module of the transformer architecture. See transformer.py
|
||||
num_classes: number of object classes
|
||||
num_queries: number of object queries, ie detection slot. This is the maximal number of objects
|
||||
Conditional DETR can detect in a single image. For COCO, we recommend 100 queries.
|
||||
aux_loss: True if auxiliary decoding losses (loss at each decoder layer) are to be used.
|
||||
group_detr: Number of groups to speed detr training. Default is 1.
|
||||
lite_refpoint_refine: TODO
|
||||
"""
|
||||
super().__init__()
|
||||
self.num_queries = num_queries
|
||||
self.transformer = transformer
|
||||
hidden_dim = transformer.d_model
|
||||
self.class_embed = nn.Linear(hidden_dim, num_classes)
|
||||
self.bbox_embed = MLP(hidden_dim, hidden_dim, 4, 3)
|
||||
self.segmentation_head = segmentation_head
|
||||
|
||||
query_dim = 4
|
||||
self.refpoint_embed = nn.Embedding(num_queries * group_detr, query_dim)
|
||||
self.query_feat = nn.Embedding(num_queries * group_detr, hidden_dim)
|
||||
nn.init.constant_(self.refpoint_embed.weight.data, 0)
|
||||
|
||||
self.backbone = backbone
|
||||
self.aux_loss = aux_loss
|
||||
self.group_detr = group_detr
|
||||
|
||||
# iter update
|
||||
self.lite_refpoint_refine = lite_refpoint_refine
|
||||
if not self.lite_refpoint_refine:
|
||||
self.transformer.decoder.bbox_embed = self.bbox_embed
|
||||
else:
|
||||
self.transformer.decoder.bbox_embed = None
|
||||
|
||||
self.bbox_reparam = bbox_reparam
|
||||
|
||||
# init prior_prob setting for focal loss
|
||||
prior_prob = 0.01
|
||||
bias_value = -math.log((1 - prior_prob) / prior_prob)
|
||||
self.class_embed.bias.data = torch.ones(num_classes) * bias_value
|
||||
|
||||
# init bbox_mebed
|
||||
nn.init.constant_(self.bbox_embed.layers[-1].weight.data, 0)
|
||||
nn.init.constant_(self.bbox_embed.layers[-1].bias.data, 0)
|
||||
|
||||
# two_stage
|
||||
self.two_stage = two_stage
|
||||
if self.two_stage:
|
||||
self.transformer.enc_out_bbox_embed = nn.ModuleList(
|
||||
[copy.deepcopy(self.bbox_embed) for _ in range(group_detr)]
|
||||
)
|
||||
self.transformer.enc_out_class_embed = nn.ModuleList(
|
||||
[copy.deepcopy(self.class_embed) for _ in range(group_detr)]
|
||||
)
|
||||
|
||||
self._export = False
|
||||
|
||||
def reinitialize_detection_head(self, num_classes):
|
||||
base = self.class_embed.weight.shape[0]
|
||||
num_repeats = int(math.ceil(num_classes / base))
|
||||
self.class_embed.weight.data = self.class_embed.weight.data.repeat(num_repeats, 1)
|
||||
self.class_embed.weight.data = self.class_embed.weight.data[:num_classes]
|
||||
self.class_embed.bias.data = self.class_embed.bias.data.repeat(num_repeats)
|
||||
self.class_embed.bias.data = self.class_embed.bias.data[:num_classes]
|
||||
|
||||
if self.two_stage:
|
||||
for enc_out_class_embed in self.transformer.enc_out_class_embed:
|
||||
enc_out_class_embed.weight.data = enc_out_class_embed.weight.data.repeat(num_repeats, 1)
|
||||
enc_out_class_embed.weight.data = enc_out_class_embed.weight.data[:num_classes]
|
||||
enc_out_class_embed.bias.data = enc_out_class_embed.bias.data.repeat(num_repeats)
|
||||
enc_out_class_embed.bias.data = enc_out_class_embed.bias.data[:num_classes]
|
||||
|
||||
def export(self):
|
||||
self._export = True
|
||||
self._forward_origin = self.forward
|
||||
self.forward = self.forward_export
|
||||
for name, m in self.named_modules():
|
||||
if hasattr(m, "export") and isinstance(m.export, Callable) and hasattr(m, "_export") and not m._export:
|
||||
m.export()
|
||||
|
||||
def forward(self, samples: NestedTensor, targets=None):
|
||||
"""The forward expects a NestedTensor, which consists of:
|
||||
- samples.tensor: batched images, of shape [batch_size x 3 x H x W]
|
||||
- samples.mask: a binary mask of shape [batch_size x H x W], containing 1 on padded pixels
|
||||
|
||||
It returns a dict with the following elements:
|
||||
- "pred_logits": the classification logits (including no-object) for all queries.
|
||||
Shape= [batch_size x num_queries x num_classes]
|
||||
- "pred_boxes": The normalized boxes coordinates for all queries, represented as
|
||||
(center_x, center_y, width, height). These values are normalized in [0, 1],
|
||||
relative to the size of each individual image (disregarding possible padding).
|
||||
See PostProcess for information on how to retrieve the unnormalized bounding box.
|
||||
- "aux_outputs": Optional, only returned when auxiliary losses are activated. It is a list of
|
||||
dictionaries containing the two above keys for each decoder layer.
|
||||
"""
|
||||
if isinstance(samples, (list, torch.Tensor)):
|
||||
samples = nested_tensor_from_tensor_list(samples)
|
||||
features, poss = self.backbone(samples)
|
||||
|
||||
srcs = []
|
||||
masks = []
|
||||
for l, feat in enumerate(features):
|
||||
src, mask = feat.decompose()
|
||||
srcs.append(src)
|
||||
masks.append(mask)
|
||||
assert mask is not None
|
||||
|
||||
if self.training:
|
||||
refpoint_embed_weight = self.refpoint_embed.weight
|
||||
query_feat_weight = self.query_feat.weight
|
||||
else:
|
||||
# only use one group in inference
|
||||
refpoint_embed_weight = self.refpoint_embed.weight[: self.num_queries]
|
||||
query_feat_weight = self.query_feat.weight[: self.num_queries]
|
||||
|
||||
if self.segmentation_head is not None:
|
||||
seg_head_fwd = self.segmentation_head.sparse_forward if self.training else self.segmentation_head.forward
|
||||
|
||||
hs, ref_unsigmoid, hs_enc, ref_enc = self.transformer(
|
||||
srcs, masks, poss, refpoint_embed_weight, query_feat_weight
|
||||
)
|
||||
|
||||
if hs is not None:
|
||||
if self.bbox_reparam:
|
||||
outputs_coord_delta = self.bbox_embed(hs)
|
||||
outputs_coord_cxcy = outputs_coord_delta[..., :2] * ref_unsigmoid[..., 2:] + ref_unsigmoid[..., :2]
|
||||
outputs_coord_wh = outputs_coord_delta[..., 2:].exp() * ref_unsigmoid[..., 2:]
|
||||
outputs_coord = torch.concat([outputs_coord_cxcy, outputs_coord_wh], dim=-1)
|
||||
else:
|
||||
outputs_coord = (self.bbox_embed(hs) + ref_unsigmoid).sigmoid()
|
||||
|
||||
outputs_class = self.class_embed(hs)
|
||||
|
||||
if self.segmentation_head is not None:
|
||||
outputs_masks = seg_head_fwd(features[0].tensors, hs, samples.tensors.shape[-2:])
|
||||
|
||||
out = {"pred_logits": outputs_class[-1], "pred_boxes": outputs_coord[-1]}
|
||||
if self.segmentation_head is not None:
|
||||
out["pred_masks"] = outputs_masks[-1]
|
||||
if self.aux_loss:
|
||||
out["aux_outputs"] = self._set_aux_loss(
|
||||
outputs_class,
|
||||
outputs_coord,
|
||||
outputs_masks if self.segmentation_head is not None else None,
|
||||
)
|
||||
|
||||
if self.two_stage:
|
||||
group_detr = self.group_detr if self.training else 1
|
||||
hs_enc_list = hs_enc.chunk(group_detr, dim=1)
|
||||
cls_enc = []
|
||||
for g_idx in range(group_detr):
|
||||
cls_enc_gidx = self.transformer.enc_out_class_embed[g_idx](hs_enc_list[g_idx])
|
||||
cls_enc.append(cls_enc_gidx)
|
||||
|
||||
cls_enc = torch.cat(cls_enc, dim=1)
|
||||
|
||||
if self.segmentation_head is not None:
|
||||
masks_enc = seg_head_fwd(
|
||||
features[0].tensors,
|
||||
[
|
||||
hs_enc,
|
||||
],
|
||||
samples.tensors.shape[-2:],
|
||||
skip_blocks=True,
|
||||
)[0]
|
||||
|
||||
if hs is not None:
|
||||
out["enc_outputs"] = {"pred_logits": cls_enc, "pred_boxes": ref_enc}
|
||||
if self.segmentation_head is not None:
|
||||
out["enc_outputs"]["pred_masks"] = masks_enc
|
||||
else:
|
||||
out = {"pred_logits": cls_enc, "pred_boxes": ref_enc}
|
||||
if self.segmentation_head is not None:
|
||||
out["pred_masks"] = masks_enc
|
||||
|
||||
return out
|
||||
|
||||
def forward_export(self, tensors):
|
||||
srcs, _, poss = self.backbone(tensors)
|
||||
# only use one group in inference
|
||||
refpoint_embed_weight = self.refpoint_embed.weight[: self.num_queries]
|
||||
query_feat_weight = self.query_feat.weight[: self.num_queries]
|
||||
|
||||
hs, ref_unsigmoid, hs_enc, ref_enc = self.transformer(
|
||||
srcs, None, poss, refpoint_embed_weight, query_feat_weight
|
||||
)
|
||||
|
||||
outputs_masks = None
|
||||
|
||||
if hs is not None:
|
||||
if self.bbox_reparam:
|
||||
outputs_coord_delta = self.bbox_embed(hs)
|
||||
outputs_coord_cxcy = outputs_coord_delta[..., :2] * ref_unsigmoid[..., 2:] + ref_unsigmoid[..., :2]
|
||||
outputs_coord_wh = outputs_coord_delta[..., 2:].exp() * ref_unsigmoid[..., 2:]
|
||||
outputs_coord = torch.concat([outputs_coord_cxcy, outputs_coord_wh], dim=-1)
|
||||
else:
|
||||
outputs_coord = (self.bbox_embed(hs) + ref_unsigmoid).sigmoid()
|
||||
outputs_class = self.class_embed(hs)
|
||||
if self.segmentation_head is not None:
|
||||
outputs_masks = self.segmentation_head(
|
||||
srcs[0],
|
||||
[
|
||||
hs,
|
||||
],
|
||||
tensors.shape[-2:],
|
||||
)[0]
|
||||
else:
|
||||
assert self.two_stage, "if not using decoder, two_stage must be True"
|
||||
outputs_class = self.transformer.enc_out_class_embed[0](hs_enc)
|
||||
outputs_coord = ref_enc
|
||||
if self.segmentation_head is not None:
|
||||
outputs_masks = self.segmentation_head(
|
||||
srcs[0],
|
||||
[
|
||||
hs_enc,
|
||||
],
|
||||
tensors.shape[-2:],
|
||||
skip_blocks=True,
|
||||
)[0]
|
||||
|
||||
if outputs_masks is not None:
|
||||
return outputs_coord, outputs_class, outputs_masks
|
||||
else:
|
||||
return outputs_coord, outputs_class
|
||||
|
||||
@torch.jit.unused
|
||||
def _set_aux_loss(self, outputs_class, outputs_coord, outputs_masks):
|
||||
# this is a workaround to make torchscript happy, as torchscript
|
||||
# doesn't support dictionary with non-homogeneous values, such
|
||||
# as a dict having both a Tensor and a list.
|
||||
if outputs_masks is not None:
|
||||
return [
|
||||
{"pred_logits": a, "pred_boxes": b, "pred_masks": c}
|
||||
for a, b, c in zip(outputs_class[:-1], outputs_coord[:-1], outputs_masks[:-1])
|
||||
]
|
||||
else:
|
||||
return [{"pred_logits": a, "pred_boxes": b} for a, b in zip(outputs_class[:-1], outputs_coord[:-1])]
|
||||
|
||||
def _get_backbone_encoder_layers(self) -> Optional[nn.ModuleList]:
|
||||
"""Resolve the list of transformer blocks/layers from backbone[0].encoder.
|
||||
|
||||
Supports multiple backbone architectures:
|
||||
- encoder.blocks (standard ViT)
|
||||
- encoder.trunk.blocks (aimv2)
|
||||
- encoder.encoder.encoder.layer (HuggingFace DinoV2)
|
||||
|
||||
Returns:
|
||||
List of transformer layers, or None if not found.
|
||||
"""
|
||||
enc = self.backbone[0].encoder
|
||||
if hasattr(enc, "blocks"):
|
||||
return enc.blocks
|
||||
if hasattr(enc, "trunk") and hasattr(enc.trunk, "blocks"):
|
||||
return enc.trunk.blocks
|
||||
if hasattr(enc, "encoder") and hasattr(enc.encoder, "encoder") and hasattr(enc.encoder.encoder, "layer"):
|
||||
return enc.encoder.encoder.layer
|
||||
return None
|
||||
|
||||
def update_drop_path(self, drop_path_rate: float, vit_encoder_num_layers: int) -> None:
|
||||
"""Update drop_path rates for backbone encoder layers with linear schedule.
|
||||
|
||||
Applies a linear schedule where the first layer has drop_path_rate=0 and the last
|
||||
layer has drop_path_rate=drop_path_rate. Intermediate layers are interpolated linearly.
|
||||
|
||||
Args:
|
||||
drop_path_rate: Maximum drop path rate (applied to last layer).
|
||||
vit_encoder_num_layers: Number of encoder layers to update.
|
||||
"""
|
||||
layers = self._get_backbone_encoder_layers()
|
||||
if layers is None:
|
||||
return
|
||||
n = min(vit_encoder_num_layers, len(layers))
|
||||
dp_rates = [x.item() for x in torch.linspace(0, drop_path_rate, n)]
|
||||
for i in range(n):
|
||||
if hasattr(layers[i], "drop_path") and hasattr(layers[i].drop_path, "drop_prob"):
|
||||
layers[i].drop_path.drop_prob = dp_rates[i]
|
||||
|
||||
def update_dropout(self, drop_rate):
|
||||
for module in self.transformer.modules():
|
||||
if isinstance(module, nn.Dropout):
|
||||
module.p = drop_rate
|
||||
|
||||
|
||||
class PostProcess(nn.Module):
|
||||
"""This module converts the model's output into the format expected by the coco api"""
|
||||
|
||||
def __init__(self, num_select=300) -> None:
|
||||
super().__init__()
|
||||
self.num_select = num_select
|
||||
|
||||
@torch.no_grad()
|
||||
def forward(self, outputs, target_sizes):
|
||||
"""Perform the computation
|
||||
Parameters:
|
||||
outputs: raw outputs of the model
|
||||
target_sizes: tensor of dimension [batch_size x 2] containing the size of each images of the batch
|
||||
For evaluation, this must be the original image size (before any data augmentation)
|
||||
For visualization, this should be the image size after data augment, but before padding
|
||||
"""
|
||||
out_logits, out_bbox = outputs["pred_logits"], outputs["pred_boxes"]
|
||||
out_masks = outputs.get("pred_masks", None)
|
||||
|
||||
assert len(out_logits) == len(target_sizes)
|
||||
assert target_sizes.shape[1] == 2
|
||||
|
||||
prob = out_logits.sigmoid()
|
||||
topk_values, topk_indexes = torch.topk(prob.view(out_logits.shape[0], -1), self.num_select, dim=1)
|
||||
scores = topk_values
|
||||
topk_boxes = topk_indexes // out_logits.shape[2]
|
||||
labels = topk_indexes % out_logits.shape[2]
|
||||
boxes = box_ops.box_cxcywh_to_xyxy(out_bbox)
|
||||
boxes = torch.gather(boxes, 1, topk_boxes.unsqueeze(-1).repeat(1, 1, 4))
|
||||
|
||||
# and from relative [0, 1] to absolute [0, height] coordinates
|
||||
img_h, img_w = target_sizes.unbind(1)
|
||||
scale_fct = torch.stack([img_w, img_h, img_w, img_h], dim=1)
|
||||
boxes = boxes * scale_fct[:, None, :]
|
||||
|
||||
# Optionally gather masks corresponding to the same top-K queries and resize to original size
|
||||
results = []
|
||||
if out_masks is not None:
|
||||
for i in range(out_masks.shape[0]):
|
||||
res_i = {"scores": scores[i], "labels": labels[i], "boxes": boxes[i]}
|
||||
k_idx = topk_boxes[i]
|
||||
masks_i = torch.gather(
|
||||
out_masks[i],
|
||||
0,
|
||||
k_idx.unsqueeze(-1).unsqueeze(-1).repeat(1, out_masks.shape[-2], out_masks.shape[-1]),
|
||||
) # [K, Hm, Wm]
|
||||
h, w = target_sizes[i].tolist()
|
||||
masks_i = F.interpolate(
|
||||
masks_i.unsqueeze(1),
|
||||
size=(int(h), int(w)),
|
||||
mode="bilinear",
|
||||
align_corners=False,
|
||||
) # [K,1,H,W]
|
||||
res_i["masks"] = masks_i > 0.0
|
||||
results.append(res_i)
|
||||
else:
|
||||
results = [{"scores": s, "labels": l, "boxes": b} for s, l, b in zip(scores, labels, boxes)]
|
||||
|
||||
return results
|
||||
|
||||
|
||||
class MLP(nn.Module):
|
||||
"""Very simple multi-layer perceptron (also called FFN)"""
|
||||
|
||||
def __init__(self, input_dim, hidden_dim, output_dim, num_layers):
|
||||
super().__init__()
|
||||
self.num_layers = num_layers
|
||||
h = [hidden_dim] * (num_layers - 1)
|
||||
self.layers = nn.ModuleList(nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim]))
|
||||
|
||||
def forward(self, x):
|
||||
for i, layer in enumerate(self.layers):
|
||||
x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
|
||||
return x
|
||||
|
||||
|
||||
def build_model(args):
|
||||
# the `num_classes` naming here is somewhat misleading.
|
||||
# it indeed corresponds to `max_obj_id + 1`, where max_obj_id
|
||||
# is the maximum id for a class in your dataset. For example,
|
||||
# COCO has a max_obj_id of 90, so we pass `num_classes` to be 91.
|
||||
# As another example, for a dataset that has a single class with id 1,
|
||||
# you should pass `num_classes` to be 2 (max_obj_id + 1).
|
||||
# For more details on this, check the following discussion
|
||||
# https://github.com/facebookresearch/detr/issues/108#issuecomment-650269223
|
||||
num_classes = args.num_classes + 1
|
||||
torch.device(args.device)
|
||||
|
||||
backbone = build_backbone(
|
||||
encoder=args.encoder,
|
||||
vit_encoder_num_layers=args.vit_encoder_num_layers,
|
||||
pretrained_encoder=args.pretrained_encoder,
|
||||
window_block_indexes=args.window_block_indexes,
|
||||
drop_path=args.drop_path,
|
||||
out_channels=args.hidden_dim,
|
||||
out_feature_indexes=args.out_feature_indexes,
|
||||
projector_scale=args.projector_scale,
|
||||
use_cls_token=args.use_cls_token,
|
||||
hidden_dim=args.hidden_dim,
|
||||
position_embedding=args.position_embedding,
|
||||
freeze_encoder=args.freeze_encoder,
|
||||
layer_norm=args.layer_norm,
|
||||
target_shape=(
|
||||
args.shape
|
||||
if hasattr(args, "shape")
|
||||
else ((args.resolution, args.resolution) if hasattr(args, "resolution") else (640, 640))
|
||||
),
|
||||
rms_norm=args.rms_norm,
|
||||
backbone_lora=args.backbone_lora,
|
||||
force_no_pretrain=args.force_no_pretrain,
|
||||
gradient_checkpointing=args.gradient_checkpointing,
|
||||
load_dinov2_weights=args.pretrain_weights is None,
|
||||
patch_size=args.patch_size,
|
||||
num_windows=args.num_windows,
|
||||
positional_encoding_size=args.positional_encoding_size,
|
||||
)
|
||||
if args.encoder_only:
|
||||
return backbone[0].encoder, None, None
|
||||
if args.backbone_only:
|
||||
return backbone, None, None
|
||||
|
||||
args.num_feature_levels = len(args.projector_scale)
|
||||
transformer = build_transformer(args)
|
||||
|
||||
# detection-only vendored copy: segmentation head removed
|
||||
segmentation_head = None
|
||||
|
||||
model = LWDETR(
|
||||
backbone,
|
||||
transformer,
|
||||
segmentation_head,
|
||||
num_classes=num_classes,
|
||||
num_queries=args.num_queries,
|
||||
aux_loss=args.aux_loss,
|
||||
group_detr=args.group_detr,
|
||||
two_stage=args.two_stage,
|
||||
lite_refpoint_refine=args.lite_refpoint_refine,
|
||||
bbox_reparam=args.bbox_reparam,
|
||||
)
|
||||
return model
|
||||
|
||||
|
||||
@@ -0,0 +1,5 @@
|
||||
# ------------------------------------------------------------------------
|
||||
# RF-DETR
|
||||
# Copyright (c) 2025 Roboflow. All Rights Reserved.
|
||||
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
|
||||
# ------------------------------------------------------------------------
|
||||
@@ -0,0 +1,18 @@
|
||||
# ------------------------------------------------------------------------
|
||||
# RF-DETR
|
||||
# Copyright (c) 2025 Roboflow. All Rights Reserved.
|
||||
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
|
||||
# ------------------------------------------------------------------------
|
||||
# Copied and modified from LW-DETR (https://github.com/Atten4Vis/LW-DETR)
|
||||
# Copyright (c) 2024 Baidu. All Rights Reserved.
|
||||
# ------------------------------------------------------------------------------------------------
|
||||
# Modified from Deformable DETR
|
||||
# Copyright (c) 2020 SenseTime. All Rights Reserved.
|
||||
# ------------------------------------------------------------------------------------------------
|
||||
# Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
|
||||
# ------------------------------------------------------------------------------------------------
|
||||
"""
|
||||
ms_deform_attn_func
|
||||
"""
|
||||
|
||||
from surya.common.rfdetr.models.ops.functions.ms_deform_attn_func import ms_deform_attn_core_pytorch
|
||||
@@ -0,0 +1,47 @@
|
||||
# ------------------------------------------------------------------------
|
||||
# RF-DETR
|
||||
# Copyright (c) 2025 Roboflow. All Rights Reserved.
|
||||
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
|
||||
# ------------------------------------------------------------------------
|
||||
# Copied and modified from LW-DETR (https://github.com/Atten4Vis/LW-DETR)
|
||||
# Copyright (c) 2024 Baidu. All Rights Reserved.
|
||||
# ------------------------------------------------------------------------------------------------
|
||||
# Modified from Deformable DETR
|
||||
# Copyright (c) 2020 SenseTime. All Rights Reserved.
|
||||
# ------------------------------------------------------------------------------------------------
|
||||
# Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
|
||||
# ------------------------------------------------------------------------------------------------
|
||||
"""
|
||||
ms_deform_attn_func
|
||||
"""
|
||||
|
||||
from __future__ import absolute_import, division, print_function
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
def ms_deform_attn_core_pytorch(value, value_spatial_shapes, sampling_locations, attention_weights):
|
||||
""" "for debug and test only, need to use cuda version instead"""
|
||||
# B, n_heads, head_dim, N
|
||||
B, n_heads, head_dim, _ = value.shape
|
||||
_, Len_q, n_heads, L, P, _ = sampling_locations.shape
|
||||
value_list = value.split([H * W for H, W in value_spatial_shapes], dim=3)
|
||||
sampling_grids = 2 * sampling_locations - 1
|
||||
sampling_value_list = []
|
||||
for lid_, (H, W) in enumerate(value_spatial_shapes):
|
||||
# B, n_heads, head_dim, H, W
|
||||
value_l_ = value_list[lid_].view(B * n_heads, head_dim, H, W)
|
||||
# B, Len_q, n_heads, P, 2 -> B, n_heads, Len_q, P, 2 -> B*n_heads, Len_q, P, 2
|
||||
sampling_grid_l_ = sampling_grids[:, :, :, lid_].transpose(1, 2).flatten(0, 1)
|
||||
# B*n_heads, head_dim, Len_q, P
|
||||
sampling_value_l_ = F.grid_sample(
|
||||
value_l_, sampling_grid_l_, mode="bilinear", padding_mode="zeros", align_corners=False
|
||||
)
|
||||
sampling_value_list.append(sampling_value_l_)
|
||||
# (B, Len_q, n_heads, L * P) -> (B, n_heads, Len_q, L, P) -> (B*n_heads, 1, Len_q, L*P)
|
||||
attention_weights = attention_weights.transpose(1, 2).reshape(B * n_heads, 1, Len_q, L * P)
|
||||
# B*n_heads, head_dim, Len_q, L*P
|
||||
sampling_value_list = torch.stack(sampling_value_list, dim=-2).flatten(-2)
|
||||
output = (sampling_value_list * attention_weights).sum(-1).view(B, n_heads * head_dim, Len_q)
|
||||
return output.transpose(1, 2).contiguous()
|
||||
@@ -0,0 +1,13 @@
|
||||
# ------------------------------------------------------------------------
|
||||
# RF-DETR
|
||||
# Copyright (c) 2025 Roboflow. All Rights Reserved.
|
||||
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
|
||||
# ------------------------------------------------------------------------
|
||||
# Copied and modified from Deformable DETR (https://github.com/fundamentalvision/Deformable-DETR)
|
||||
# Copyright (c) 2020 SenseTime. All Rights Reserved.
|
||||
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
|
||||
# ------------------------------------------------------------------------
|
||||
# Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
|
||||
# ------------------------------------------------------------------------
|
||||
|
||||
from surya.common.rfdetr.models.ops.modules.ms_deform_attn import MSDeformAttn
|
||||
@@ -0,0 +1,152 @@
|
||||
# ------------------------------------------------------------------------
|
||||
# RF-DETR
|
||||
# Copyright (c) 2025 Roboflow. All Rights Reserved.
|
||||
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
|
||||
# ------------------------------------------------------------------------
|
||||
# Copied and modified from LW-DETR (https://github.com/Atten4Vis/LW-DETR)
|
||||
# Copyright (c) 2024 Baidu. All Rights Reserved.
|
||||
# ------------------------------------------------------------------------------------------------
|
||||
# Modified from Deformable DETR
|
||||
# Copyright (c) 2020 SenseTime. All Rights Reserved.
|
||||
# ------------------------------------------------------------------------------------------------
|
||||
# Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
|
||||
# ------------------------------------------------------------------------------------------------
|
||||
"""
|
||||
Multi-Scale Deformable Attention Module
|
||||
"""
|
||||
|
||||
from __future__ import absolute_import, division, print_function
|
||||
|
||||
import math
|
||||
import warnings
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
from torch.nn.init import constant_, xavier_uniform_
|
||||
|
||||
from surya.common.rfdetr.models.ops.functions import ms_deform_attn_core_pytorch
|
||||
|
||||
|
||||
def _is_power_of_2(n):
|
||||
if (not isinstance(n, int)) or (n < 0):
|
||||
raise ValueError("invalid input for _is_power_of_2: {} (type: {})".format(n, type(n)))
|
||||
return (n & (n - 1) == 0) and n != 0
|
||||
|
||||
|
||||
class MSDeformAttn(nn.Module):
|
||||
"""Multi-Scale Deformable Attention Module"""
|
||||
|
||||
def __init__(self, d_model=256, n_levels=4, n_heads=8, n_points=4):
|
||||
"""
|
||||
Multi-Scale Deformable Attention Module
|
||||
:param d_model hidden dimension
|
||||
:param n_levels number of feature levels
|
||||
:param n_heads number of attention heads
|
||||
:param n_points number of sampling points per attention head per feature level
|
||||
"""
|
||||
super().__init__()
|
||||
if d_model % n_heads != 0:
|
||||
raise ValueError("d_model must be divisible by n_heads, but got {} and {}".format(d_model, n_heads))
|
||||
_d_per_head = d_model // n_heads
|
||||
# you'd better set _d_per_head to a power of 2 which is more efficient in our CUDA implementation
|
||||
if not _is_power_of_2(_d_per_head):
|
||||
warnings.warn(
|
||||
"You'd better set d_model in MSDeformAttn to make the "
|
||||
"dimension of each attention head a power of 2 "
|
||||
"which is more efficient in our CUDA implementation."
|
||||
)
|
||||
|
||||
self.im2col_step = 64
|
||||
|
||||
self.d_model = d_model
|
||||
self.n_levels = n_levels
|
||||
self.n_heads = n_heads
|
||||
self.n_points = n_points
|
||||
|
||||
self.sampling_offsets = nn.Linear(d_model, n_heads * n_levels * n_points * 2)
|
||||
self.attention_weights = nn.Linear(d_model, n_heads * n_levels * n_points)
|
||||
self.value_proj = nn.Linear(d_model, d_model)
|
||||
self.output_proj = nn.Linear(d_model, d_model)
|
||||
|
||||
self._reset_parameters()
|
||||
|
||||
self._export = False
|
||||
|
||||
def export(self):
|
||||
"""export mode"""
|
||||
self._export = True
|
||||
|
||||
def _reset_parameters(self):
|
||||
constant_(self.sampling_offsets.weight.data, 0.0)
|
||||
thetas = torch.arange(self.n_heads, dtype=torch.float32) * (2.0 * math.pi / self.n_heads)
|
||||
grid_init = torch.stack([thetas.cos(), thetas.sin()], -1)
|
||||
grid_init = (
|
||||
(grid_init / grid_init.abs().max(-1, keepdim=True)[0])
|
||||
.view(self.n_heads, 1, 1, 2)
|
||||
.repeat(1, self.n_levels, self.n_points, 1)
|
||||
)
|
||||
for i in range(self.n_points):
|
||||
grid_init[:, :, i, :] *= i + 1
|
||||
with torch.no_grad():
|
||||
self.sampling_offsets.bias = nn.Parameter(grid_init.view(-1))
|
||||
constant_(self.attention_weights.weight.data, 0.0)
|
||||
constant_(self.attention_weights.bias.data, 0.0)
|
||||
xavier_uniform_(self.value_proj.weight.data)
|
||||
constant_(self.value_proj.bias.data, 0.0)
|
||||
xavier_uniform_(self.output_proj.weight.data)
|
||||
constant_(self.output_proj.bias.data, 0.0)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
query,
|
||||
reference_points,
|
||||
input_flatten,
|
||||
input_spatial_shapes,
|
||||
input_level_start_index,
|
||||
input_padding_mask=None,
|
||||
):
|
||||
r"""
|
||||
:param query (N, Length_{query}, C)
|
||||
:param reference_points (N, Length_{query}, n_levels, 2), range in [0, 1], top-left (0,0), bottom-right (1, 1), including padding area
|
||||
or (N, Length_{query}, n_levels, 4), add additional (w, h) to form reference boxes
|
||||
:param input_flatten (N, \sum_{l=0}^{L-1} H_l \cdot W_l, C)
|
||||
:param input_spatial_shapes (n_levels, 2), [(H_0, W_0), (H_1, W_1), ..., (H_{L-1}, W_{L-1})]
|
||||
:param input_level_start_index (n_levels, ), [0, H_0*W_0, H_0*W_0+H_1*W_1, H_0*W_0+H_1*W_1+H_2*W_2, ..., H_0*W_0+H_1*W_1+...+H_{L-1}*W_{L-1}]
|
||||
:param input_padding_mask (N, \sum_{l=0}^{L-1} H_l \cdot W_l), True for padding elements, False for non-padding elements
|
||||
|
||||
:return output (N, Length_{query}, C)
|
||||
"""
|
||||
N, Len_q, _ = query.shape
|
||||
N, Len_in, _ = input_flatten.shape
|
||||
assert (input_spatial_shapes[:, 0] * input_spatial_shapes[:, 1]).sum() == Len_in
|
||||
|
||||
value = self.value_proj(input_flatten)
|
||||
if input_padding_mask is not None:
|
||||
value = value.masked_fill(input_padding_mask[..., None], float(0))
|
||||
|
||||
sampling_offsets = self.sampling_offsets(query).view(N, Len_q, self.n_heads, self.n_levels, self.n_points, 2)
|
||||
attention_weights = self.attention_weights(query).view(N, Len_q, self.n_heads, self.n_levels * self.n_points)
|
||||
|
||||
# N, Len_q, n_heads, n_levels, n_points, 2
|
||||
if reference_points.shape[-1] == 2:
|
||||
offset_normalizer = torch.stack([input_spatial_shapes[..., 1], input_spatial_shapes[..., 0]], -1)
|
||||
sampling_locations = (
|
||||
reference_points[:, :, None, :, None, :]
|
||||
+ sampling_offsets / offset_normalizer[None, None, None, :, None, :]
|
||||
)
|
||||
elif reference_points.shape[-1] == 4:
|
||||
sampling_locations = (
|
||||
reference_points[:, :, None, :, None, :2]
|
||||
+ sampling_offsets / self.n_points * reference_points[:, :, None, :, None, 2:] * 0.5
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
"Last dim of reference_points must be 2 or 4, but get {} instead.".format(reference_points.shape[-1])
|
||||
)
|
||||
attention_weights = F.softmax(attention_weights, -1)
|
||||
|
||||
value = value.transpose(1, 2).contiguous().view(N, self.n_heads, self.d_model // self.n_heads, Len_in)
|
||||
output = ms_deform_attn_core_pytorch(value, input_spatial_shapes, sampling_locations, attention_weights)
|
||||
output = self.output_proj(output)
|
||||
return output
|
||||
@@ -0,0 +1,155 @@
|
||||
# ------------------------------------------------------------------------
|
||||
# RF-DETR
|
||||
# Copyright (c) 2025 Roboflow. All Rights Reserved.
|
||||
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
|
||||
# ------------------------------------------------------------------------
|
||||
# Copied and modified from LW-DETR (https://github.com/Atten4Vis/LW-DETR)
|
||||
# Copyright (c) 2024 Baidu. All Rights Reserved.
|
||||
# ------------------------------------------------------------------------
|
||||
# Modified from Conditional DETR (https://github.com/Atten4Vis/ConditionalDETR)
|
||||
# Copyright (c) 2021 Microsoft. All Rights Reserved.
|
||||
# ------------------------------------------------------------------------
|
||||
# Copied from DETR (https://github.com/facebookresearch/detr)
|
||||
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
|
||||
# ------------------------------------------------------------------------
|
||||
|
||||
"""
|
||||
Various positional encodings for the transformer.
|
||||
"""
|
||||
|
||||
import math
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from surya.common.rfdetr.util.misc import NestedTensor
|
||||
|
||||
|
||||
class PositionEmbeddingSine(nn.Module):
|
||||
"""
|
||||
This is a more standard version of the position embedding, very similar to the one
|
||||
used by the Attention is all you need paper, generalized to work on images.
|
||||
"""
|
||||
|
||||
def __init__(self, num_pos_feats=64, temperature=10000, normalize=False, scale=None):
|
||||
super().__init__()
|
||||
self.num_pos_feats = num_pos_feats
|
||||
self.temperature = temperature
|
||||
self.normalize = normalize
|
||||
if scale is not None and normalize is False:
|
||||
raise ValueError("normalize should be True if scale is passed")
|
||||
if scale is None:
|
||||
scale = 2 * math.pi
|
||||
self.scale = scale
|
||||
self._export = False
|
||||
|
||||
def export(self):
|
||||
self._export = True
|
||||
self._forward_origin = self.forward
|
||||
self.forward = self.forward_export
|
||||
|
||||
def forward(self, tensor_list: NestedTensor, align_dim_orders=True):
|
||||
x = tensor_list.tensors
|
||||
mask = tensor_list.mask
|
||||
assert mask is not None
|
||||
not_mask = ~mask
|
||||
y_embed = not_mask.cumsum(1, dtype=torch.float32)
|
||||
x_embed = not_mask.cumsum(2, dtype=torch.float32)
|
||||
if self.normalize:
|
||||
eps = 1e-6
|
||||
y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale
|
||||
x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale
|
||||
|
||||
dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device)
|
||||
dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats)
|
||||
|
||||
pos_x = x_embed[:, :, :, None] / dim_t
|
||||
pos_y = y_embed[:, :, :, None] / dim_t
|
||||
pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3)
|
||||
pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3)
|
||||
if align_dim_orders:
|
||||
pos = torch.cat((pos_y, pos_x), dim=3).permute(1, 2, 0, 3)
|
||||
# return: (H, W, bs, C)
|
||||
else:
|
||||
pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2)
|
||||
# return: (bs, C, H, W)
|
||||
return pos
|
||||
|
||||
def forward_export(self, mask: torch.Tensor, align_dim_orders=True):
|
||||
assert mask is not None
|
||||
not_mask = ~mask
|
||||
y_embed = not_mask.cumsum(1, dtype=torch.float32)
|
||||
x_embed = not_mask.cumsum(2, dtype=torch.float32)
|
||||
if self.normalize:
|
||||
eps = 1e-6
|
||||
y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale
|
||||
x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale
|
||||
|
||||
dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=mask.device)
|
||||
dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats)
|
||||
|
||||
pos_x = x_embed[:, :, :, None] / dim_t
|
||||
pos_y = y_embed[:, :, :, None] / dim_t
|
||||
pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3)
|
||||
pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3)
|
||||
if align_dim_orders:
|
||||
pos = torch.cat((pos_y, pos_x), dim=3).permute(1, 2, 0, 3)
|
||||
# return: (H, W, bs, C)
|
||||
else:
|
||||
pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2)
|
||||
# return: (bs, C, H, W)
|
||||
return pos
|
||||
|
||||
|
||||
class PositionEmbeddingLearned(nn.Module):
|
||||
"""
|
||||
Absolute pos embedding, learned.
|
||||
"""
|
||||
|
||||
def __init__(self, num_pos_feats=256):
|
||||
super().__init__()
|
||||
self.row_embed = nn.Embedding(50, num_pos_feats)
|
||||
self.col_embed = nn.Embedding(50, num_pos_feats)
|
||||
self.reset_parameters()
|
||||
self._export = False
|
||||
|
||||
def export(self):
|
||||
raise NotImplementedError
|
||||
|
||||
def reset_parameters(self):
|
||||
nn.init.uniform_(self.row_embed.weight)
|
||||
nn.init.uniform_(self.col_embed.weight)
|
||||
|
||||
def forward(self, tensor_list: NestedTensor):
|
||||
x = tensor_list.tensors
|
||||
h, w = x.shape[:2]
|
||||
i = torch.arange(w, device=x.device)
|
||||
j = torch.arange(h, device=x.device)
|
||||
x_emb = self.col_embed(i)
|
||||
y_emb = self.row_embed(j)
|
||||
pos = (
|
||||
torch.cat(
|
||||
[
|
||||
x_emb.unsqueeze(0).repeat(h, 1, 1),
|
||||
y_emb.unsqueeze(1).repeat(1, w, 1),
|
||||
],
|
||||
dim=-1,
|
||||
)
|
||||
.unsqueeze(2)
|
||||
.repeat(1, 1, x.shape[2], 1)
|
||||
)
|
||||
# return: (H, W, bs, C)
|
||||
return pos
|
||||
|
||||
|
||||
def build_position_encoding(hidden_dim, position_embedding):
|
||||
N_steps = hidden_dim // 2
|
||||
if position_embedding in ("v2", "sine"):
|
||||
# TODO find a better way of exposing other arguments
|
||||
position_embedding = PositionEmbeddingSine(N_steps, normalize=True)
|
||||
elif position_embedding in ("v3", "learned"):
|
||||
position_embedding = PositionEmbeddingLearned(N_steps)
|
||||
else:
|
||||
raise ValueError(f"not supported {position_embedding}")
|
||||
|
||||
return position_embedding
|
||||
@@ -0,0 +1,659 @@
|
||||
# ------------------------------------------------------------------------
|
||||
# RF-DETR
|
||||
# Copyright (c) 2025 Roboflow. All Rights Reserved.
|
||||
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
|
||||
# ------------------------------------------------------------------------
|
||||
# Copied and modified from LW-DETR (https://github.com/Atten4Vis/LW-DETR)
|
||||
# Copyright (c) 2024 Baidu. All Rights Reserved.
|
||||
# ------------------------------------------------------------------------
|
||||
# Modified from Conditional DETR (https://github.com/Atten4Vis/ConditionalDETR)
|
||||
# Copyright (c) 2021 Microsoft. All Rights Reserved.
|
||||
# ------------------------------------------------------------------------
|
||||
# Modified from DETR (https://github.com/facebookresearch/detr)
|
||||
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
|
||||
# ------------------------------------------------------------------------
|
||||
"""
|
||||
Transformer class
|
||||
"""
|
||||
|
||||
import copy
|
||||
import math
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import Tensor, nn
|
||||
|
||||
from surya.common.rfdetr.models.ops.modules import MSDeformAttn
|
||||
|
||||
|
||||
class MLP(nn.Module):
|
||||
"""Very simple multi-layer perceptron (also called FFN)"""
|
||||
|
||||
def __init__(self, input_dim, hidden_dim, output_dim, num_layers):
|
||||
super().__init__()
|
||||
self.num_layers = num_layers
|
||||
h = [hidden_dim] * (num_layers - 1)
|
||||
self.layers = nn.ModuleList(nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim]))
|
||||
|
||||
def forward(self, x):
|
||||
for i, layer in enumerate(self.layers):
|
||||
x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
|
||||
return x
|
||||
|
||||
|
||||
def gen_sineembed_for_position(pos_tensor, dim=128):
|
||||
# n_query, bs, _ = pos_tensor.size()
|
||||
# sineembed_tensor = torch.zeros(n_query, bs, 256)
|
||||
scale = 2 * math.pi
|
||||
dim_t = torch.arange(dim, dtype=pos_tensor.dtype, device=pos_tensor.device)
|
||||
dim_t = 10000 ** (2 * (dim_t // 2) / dim)
|
||||
x_embed = pos_tensor[:, :, 0] * scale
|
||||
y_embed = pos_tensor[:, :, 1] * scale
|
||||
pos_x = x_embed[:, :, None] / dim_t
|
||||
pos_y = y_embed[:, :, None] / dim_t
|
||||
pos_x = torch.stack((pos_x[:, :, 0::2].sin(), pos_x[:, :, 1::2].cos()), dim=3).flatten(2)
|
||||
pos_y = torch.stack((pos_y[:, :, 0::2].sin(), pos_y[:, :, 1::2].cos()), dim=3).flatten(2)
|
||||
if pos_tensor.size(-1) == 2:
|
||||
pos = torch.cat((pos_y, pos_x), dim=2)
|
||||
elif pos_tensor.size(-1) == 4:
|
||||
w_embed = pos_tensor[:, :, 2] * scale
|
||||
pos_w = w_embed[:, :, None] / dim_t
|
||||
pos_w = torch.stack((pos_w[:, :, 0::2].sin(), pos_w[:, :, 1::2].cos()), dim=3).flatten(2)
|
||||
|
||||
h_embed = pos_tensor[:, :, 3] * scale
|
||||
pos_h = h_embed[:, :, None] / dim_t
|
||||
pos_h = torch.stack((pos_h[:, :, 0::2].sin(), pos_h[:, :, 1::2].cos()), dim=3).flatten(2)
|
||||
|
||||
pos = torch.cat((pos_y, pos_x, pos_w, pos_h), dim=2)
|
||||
else:
|
||||
raise ValueError("Unknown pos_tensor shape(-1):{}".format(pos_tensor.size(-1)))
|
||||
return pos
|
||||
|
||||
|
||||
def gen_encoder_output_proposals(memory, memory_padding_mask, spatial_shapes, unsigmoid=True):
|
||||
r"""
|
||||
Input:
|
||||
- memory: bs, \sum{hw}, d_model
|
||||
- memory_padding_mask: bs, \sum{hw}
|
||||
- spatial_shapes: nlevel, 2
|
||||
Output:
|
||||
- output_memory: bs, \sum{hw}, d_model
|
||||
- output_proposals: bs, \sum{hw}, 4
|
||||
"""
|
||||
N_, S_, C_ = memory.shape
|
||||
proposals = []
|
||||
_cur = 0
|
||||
for lvl, (H_, W_) in enumerate(spatial_shapes):
|
||||
if memory_padding_mask is not None:
|
||||
mask_flatten_ = memory_padding_mask[:, _cur : (_cur + H_ * W_)].view(N_, H_, W_, 1)
|
||||
valid_H = torch.sum(~mask_flatten_[:, :, 0, 0], 1)
|
||||
valid_W = torch.sum(~mask_flatten_[:, 0, :, 0], 1)
|
||||
else:
|
||||
valid_H = torch.tensor([H_ for _ in range(N_)], device=memory.device)
|
||||
valid_W = torch.tensor([W_ for _ in range(N_)], device=memory.device)
|
||||
|
||||
grid_y, grid_x = torch.meshgrid(
|
||||
torch.linspace(0, H_ - 1, H_, dtype=torch.float32, device=memory.device),
|
||||
torch.linspace(0, W_ - 1, W_, dtype=torch.float32, device=memory.device),
|
||||
indexing="ij",
|
||||
)
|
||||
grid = torch.cat([grid_x.unsqueeze(-1), grid_y.unsqueeze(-1)], -1) # H_, W_, 2
|
||||
|
||||
scale = torch.cat([valid_W.unsqueeze(-1), valid_H.unsqueeze(-1)], 1).view(N_, 1, 1, 2)
|
||||
grid = (grid.unsqueeze(0).expand(N_, -1, -1, -1) + 0.5) / scale
|
||||
|
||||
wh = torch.ones_like(grid) * 0.05 * (2.0**lvl)
|
||||
|
||||
proposal = torch.cat((grid, wh), -1).view(N_, -1, 4)
|
||||
proposals.append(proposal)
|
||||
_cur += H_ * W_
|
||||
|
||||
output_proposals = torch.cat(proposals, 1)
|
||||
output_proposals_valid = ((output_proposals > 0.01) & (output_proposals < 0.99)).all(-1, keepdim=True)
|
||||
|
||||
if unsigmoid:
|
||||
output_proposals = torch.log(output_proposals / (1 - output_proposals)) # unsigmoid
|
||||
if memory_padding_mask is not None:
|
||||
output_proposals = output_proposals.masked_fill(memory_padding_mask.unsqueeze(-1), float("inf"))
|
||||
output_proposals = output_proposals.masked_fill(~output_proposals_valid, float("inf"))
|
||||
else:
|
||||
if memory_padding_mask is not None:
|
||||
output_proposals = output_proposals.masked_fill(memory_padding_mask.unsqueeze(-1), float(0))
|
||||
output_proposals = output_proposals.masked_fill(~output_proposals_valid, float(0))
|
||||
|
||||
output_memory = memory
|
||||
if memory_padding_mask is not None:
|
||||
output_memory = output_memory.masked_fill(memory_padding_mask.unsqueeze(-1), float(0))
|
||||
output_memory = output_memory.masked_fill(~output_proposals_valid, float(0))
|
||||
|
||||
return output_memory.to(memory.dtype), output_proposals.to(memory.dtype)
|
||||
|
||||
|
||||
class Transformer(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
d_model=512,
|
||||
sa_nhead=8,
|
||||
ca_nhead=8,
|
||||
num_queries=300,
|
||||
num_decoder_layers=6,
|
||||
dim_feedforward=2048,
|
||||
dropout=0.0,
|
||||
activation="relu",
|
||||
normalize_before=False,
|
||||
return_intermediate_dec=False,
|
||||
group_detr=1,
|
||||
two_stage=False,
|
||||
num_feature_levels=4,
|
||||
dec_n_points=4,
|
||||
lite_refpoint_refine=False,
|
||||
decoder_norm_type="LN",
|
||||
bbox_reparam=False,
|
||||
):
|
||||
super().__init__()
|
||||
self.encoder = None
|
||||
|
||||
decoder_layer = TransformerDecoderLayer(
|
||||
d_model,
|
||||
sa_nhead,
|
||||
ca_nhead,
|
||||
dim_feedforward,
|
||||
dropout,
|
||||
activation,
|
||||
normalize_before,
|
||||
group_detr=group_detr,
|
||||
num_feature_levels=num_feature_levels,
|
||||
dec_n_points=dec_n_points,
|
||||
skip_self_attn=False,
|
||||
)
|
||||
assert decoder_norm_type in ["LN", "Identity"]
|
||||
norm = {
|
||||
"LN": lambda channels: nn.LayerNorm(channels),
|
||||
"Identity": lambda channels: nn.Identity(),
|
||||
}
|
||||
decoder_norm = norm[decoder_norm_type](d_model)
|
||||
|
||||
self.decoder = TransformerDecoder(
|
||||
decoder_layer,
|
||||
num_decoder_layers,
|
||||
decoder_norm,
|
||||
return_intermediate=return_intermediate_dec,
|
||||
d_model=d_model,
|
||||
lite_refpoint_refine=lite_refpoint_refine,
|
||||
bbox_reparam=bbox_reparam,
|
||||
)
|
||||
|
||||
self.two_stage = two_stage
|
||||
if two_stage:
|
||||
self.enc_output = nn.ModuleList([nn.Linear(d_model, d_model) for _ in range(group_detr)])
|
||||
self.enc_output_norm = nn.ModuleList([nn.LayerNorm(d_model) for _ in range(group_detr)])
|
||||
|
||||
self._reset_parameters()
|
||||
|
||||
self.num_queries = num_queries
|
||||
self.d_model = d_model
|
||||
self.dec_layers = num_decoder_layers
|
||||
self.group_detr = group_detr
|
||||
self.num_feature_levels = num_feature_levels
|
||||
self.bbox_reparam = bbox_reparam
|
||||
|
||||
self._export = False
|
||||
|
||||
def export(self):
|
||||
self._export = True
|
||||
|
||||
def _reset_parameters(self):
|
||||
for p in self.parameters():
|
||||
if p.dim() > 1:
|
||||
nn.init.xavier_uniform_(p)
|
||||
for m in self.modules():
|
||||
if isinstance(m, MSDeformAttn):
|
||||
m._reset_parameters()
|
||||
|
||||
def get_valid_ratio(self, mask):
|
||||
_, H, W = mask.shape
|
||||
valid_H = torch.sum(~mask[:, :, 0], 1)
|
||||
valid_W = torch.sum(~mask[:, 0, :], 1)
|
||||
valid_ratio_h = valid_H.float() / H
|
||||
valid_ratio_w = valid_W.float() / W
|
||||
valid_ratio = torch.stack([valid_ratio_w, valid_ratio_h], -1)
|
||||
return valid_ratio
|
||||
|
||||
def forward(self, srcs, masks, pos_embeds, refpoint_embed, query_feat):
|
||||
src_flatten = []
|
||||
mask_flatten = [] if masks is not None else None
|
||||
lvl_pos_embed_flatten = []
|
||||
spatial_shapes = []
|
||||
valid_ratios = [] if masks is not None else None
|
||||
for lvl, (src, pos_embed) in enumerate(zip(srcs, pos_embeds)):
|
||||
bs, c, h, w = src.shape
|
||||
spatial_shape = (h, w)
|
||||
spatial_shapes.append(spatial_shape)
|
||||
|
||||
src = src.flatten(2).transpose(1, 2) # bs, hw, c
|
||||
pos_embed = pos_embed.flatten(2).transpose(1, 2) # bs, hw, c
|
||||
lvl_pos_embed_flatten.append(pos_embed)
|
||||
src_flatten.append(src)
|
||||
if masks is not None:
|
||||
mask = masks[lvl].flatten(1) # bs, hw
|
||||
mask_flatten.append(mask)
|
||||
memory = torch.cat(src_flatten, 1) # bs, \sum{hxw}, c
|
||||
if masks is not None:
|
||||
mask_flatten = torch.cat(mask_flatten, 1) # bs, \sum{hxw}
|
||||
valid_ratios = torch.stack([self.get_valid_ratio(m) for m in masks], 1)
|
||||
lvl_pos_embed_flatten = torch.cat(lvl_pos_embed_flatten, 1) # bs, \sum{hxw}, c
|
||||
spatial_shapes = torch.as_tensor(spatial_shapes, dtype=torch.long, device=memory.device)
|
||||
level_start_index = torch.cat((spatial_shapes.new_zeros((1,)), spatial_shapes.prod(1).cumsum(0)[:-1]))
|
||||
|
||||
if self.two_stage:
|
||||
output_memory, output_proposals = gen_encoder_output_proposals(
|
||||
memory, mask_flatten, spatial_shapes, unsigmoid=not self.bbox_reparam
|
||||
)
|
||||
# group detr for first stage
|
||||
refpoint_embed_ts, memory_ts, boxes_ts = [], [], []
|
||||
group_detr = self.group_detr if self.training else 1
|
||||
for g_idx in range(group_detr):
|
||||
output_memory_gidx = self.enc_output_norm[g_idx](self.enc_output[g_idx](output_memory))
|
||||
|
||||
enc_outputs_class_unselected_gidx = self.enc_out_class_embed[g_idx](output_memory_gidx)
|
||||
if self.bbox_reparam:
|
||||
enc_outputs_coord_delta_gidx = self.enc_out_bbox_embed[g_idx](output_memory_gidx)
|
||||
enc_outputs_coord_cxcy_gidx = (
|
||||
enc_outputs_coord_delta_gidx[..., :2] * output_proposals[..., 2:] + output_proposals[..., :2]
|
||||
)
|
||||
enc_outputs_coord_wh_gidx = enc_outputs_coord_delta_gidx[..., 2:].exp() * output_proposals[..., 2:]
|
||||
enc_outputs_coord_unselected_gidx = torch.concat(
|
||||
[enc_outputs_coord_cxcy_gidx, enc_outputs_coord_wh_gidx], dim=-1
|
||||
)
|
||||
else:
|
||||
enc_outputs_coord_unselected_gidx = (
|
||||
self.enc_out_bbox_embed[g_idx](output_memory_gidx) + output_proposals
|
||||
) # (bs, \sum{hw}, 4) unsigmoid
|
||||
|
||||
topk = min(self.num_queries, enc_outputs_class_unselected_gidx.shape[-2])
|
||||
topk_proposals_gidx = torch.topk(enc_outputs_class_unselected_gidx.max(-1)[0], topk, dim=1)[1] # bs, nq
|
||||
|
||||
refpoint_embed_gidx_undetach = torch.gather(
|
||||
enc_outputs_coord_unselected_gidx, 1, topk_proposals_gidx.unsqueeze(-1).repeat(1, 1, 4)
|
||||
) # unsigmoid
|
||||
# for decoder layer, detached as initial ones, (bs, nq, 4)
|
||||
refpoint_embed_gidx = refpoint_embed_gidx_undetach.detach()
|
||||
|
||||
# get memory tgt
|
||||
tgt_undetach_gidx = torch.gather(
|
||||
output_memory_gidx, 1, topk_proposals_gidx.unsqueeze(-1).repeat(1, 1, self.d_model)
|
||||
)
|
||||
|
||||
refpoint_embed_ts.append(refpoint_embed_gidx)
|
||||
memory_ts.append(tgt_undetach_gidx)
|
||||
boxes_ts.append(refpoint_embed_gidx_undetach)
|
||||
# concat on dim=1, the nq dimension, (bs, nq, d) --> (bs, nq, d)
|
||||
refpoint_embed_ts = torch.cat(refpoint_embed_ts, dim=1)
|
||||
# (bs, nq, d)
|
||||
memory_ts = torch.cat(memory_ts, dim=1) # .transpose(0, 1)
|
||||
boxes_ts = torch.cat(boxes_ts, dim=1) # .transpose(0, 1)
|
||||
|
||||
if self.dec_layers > 0:
|
||||
tgt = query_feat.unsqueeze(0).repeat(bs, 1, 1)
|
||||
refpoint_embed = refpoint_embed.unsqueeze(0).repeat(bs, 1, 1)
|
||||
if self.two_stage:
|
||||
ts_len = refpoint_embed_ts.shape[-2]
|
||||
refpoint_embed_ts_subset = refpoint_embed[..., :ts_len, :]
|
||||
refpoint_embed_subset = refpoint_embed[..., ts_len:, :]
|
||||
|
||||
if self.bbox_reparam:
|
||||
refpoint_embed_cxcy = refpoint_embed_ts_subset[..., :2] * refpoint_embed_ts[..., 2:]
|
||||
refpoint_embed_cxcy = refpoint_embed_cxcy + refpoint_embed_ts[..., :2]
|
||||
refpoint_embed_wh = refpoint_embed_ts_subset[..., 2:].exp() * refpoint_embed_ts[..., 2:]
|
||||
refpoint_embed_ts_subset = torch.concat([refpoint_embed_cxcy, refpoint_embed_wh], dim=-1)
|
||||
else:
|
||||
refpoint_embed_ts_subset = refpoint_embed_ts_subset + refpoint_embed_ts
|
||||
|
||||
refpoint_embed = torch.concat([refpoint_embed_ts_subset, refpoint_embed_subset], dim=-2)
|
||||
|
||||
hs, references = self.decoder(
|
||||
tgt,
|
||||
memory,
|
||||
memory_key_padding_mask=mask_flatten,
|
||||
pos=lvl_pos_embed_flatten,
|
||||
refpoints_unsigmoid=refpoint_embed,
|
||||
level_start_index=level_start_index,
|
||||
spatial_shapes=spatial_shapes,
|
||||
valid_ratios=valid_ratios.to(memory.dtype) if valid_ratios is not None else valid_ratios,
|
||||
)
|
||||
else:
|
||||
assert self.two_stage, "if not using decoder, two_stage must be True"
|
||||
hs = None
|
||||
references = None
|
||||
|
||||
if self.two_stage:
|
||||
if self.bbox_reparam:
|
||||
return hs, references, memory_ts, boxes_ts
|
||||
else:
|
||||
return hs, references, memory_ts, boxes_ts.sigmoid()
|
||||
return hs, references, None, None
|
||||
|
||||
|
||||
class TransformerDecoder(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
decoder_layer,
|
||||
num_layers,
|
||||
norm=None,
|
||||
return_intermediate=False,
|
||||
d_model=256,
|
||||
lite_refpoint_refine=False,
|
||||
bbox_reparam=False,
|
||||
):
|
||||
super().__init__()
|
||||
self.layers = _get_clones(decoder_layer, num_layers)
|
||||
self.num_layers = num_layers
|
||||
self.d_model = d_model
|
||||
self.norm = norm
|
||||
self.return_intermediate = return_intermediate
|
||||
self.lite_refpoint_refine = lite_refpoint_refine
|
||||
self.bbox_reparam = bbox_reparam
|
||||
|
||||
self.ref_point_head = MLP(2 * d_model, d_model, d_model, 2)
|
||||
|
||||
self._export = False
|
||||
|
||||
def export(self):
|
||||
self._export = True
|
||||
|
||||
def refpoints_refine(self, refpoints_unsigmoid, new_refpoints_delta):
|
||||
if self.bbox_reparam:
|
||||
new_refpoints_cxcy = (
|
||||
new_refpoints_delta[..., :2] * refpoints_unsigmoid[..., 2:] + refpoints_unsigmoid[..., :2]
|
||||
)
|
||||
new_refpoints_wh = new_refpoints_delta[..., 2:].exp() * refpoints_unsigmoid[..., 2:]
|
||||
new_refpoints_unsigmoid = torch.concat([new_refpoints_cxcy, new_refpoints_wh], dim=-1)
|
||||
else:
|
||||
new_refpoints_unsigmoid = refpoints_unsigmoid + new_refpoints_delta
|
||||
return new_refpoints_unsigmoid
|
||||
|
||||
def forward(
|
||||
self,
|
||||
tgt,
|
||||
memory,
|
||||
tgt_mask: Optional[Tensor] = None,
|
||||
memory_mask: Optional[Tensor] = None,
|
||||
tgt_key_padding_mask: Optional[Tensor] = None,
|
||||
memory_key_padding_mask: Optional[Tensor] = None,
|
||||
pos: Optional[Tensor] = None,
|
||||
refpoints_unsigmoid: Optional[Tensor] = None,
|
||||
# for memory
|
||||
level_start_index: Optional[Tensor] = None, # num_levels
|
||||
spatial_shapes: Optional[Tensor] = None, # bs, num_levels, 2
|
||||
valid_ratios: Optional[Tensor] = None,
|
||||
):
|
||||
output = tgt
|
||||
|
||||
intermediate = []
|
||||
hs_refpoints_unsigmoid = [refpoints_unsigmoid]
|
||||
|
||||
def get_reference(refpoints):
|
||||
# [num_queries, batch_size, 4]
|
||||
obj_center = refpoints[..., :4]
|
||||
|
||||
if self._export:
|
||||
query_sine_embed = gen_sineembed_for_position(obj_center, self.d_model / 2) # bs, nq, 256*2
|
||||
refpoints_input = obj_center[:, :, None] # bs, nq, 1, 4
|
||||
else:
|
||||
refpoints_input = (
|
||||
obj_center[:, :, None] * torch.cat([valid_ratios, valid_ratios], -1)[:, None]
|
||||
) # bs, nq, nlevel, 4
|
||||
query_sine_embed = gen_sineembed_for_position(
|
||||
refpoints_input[:, :, 0, :], self.d_model / 2
|
||||
) # bs, nq, 256*2
|
||||
query_pos = self.ref_point_head(query_sine_embed)
|
||||
return obj_center, refpoints_input, query_pos, query_sine_embed
|
||||
|
||||
# always use init refpoints
|
||||
if self.lite_refpoint_refine:
|
||||
if self.bbox_reparam:
|
||||
obj_center, refpoints_input, query_pos, query_sine_embed = get_reference(refpoints_unsigmoid)
|
||||
else:
|
||||
obj_center, refpoints_input, query_pos, query_sine_embed = get_reference(refpoints_unsigmoid.sigmoid())
|
||||
|
||||
for layer_id, layer in enumerate(self.layers):
|
||||
# iter refine each layer
|
||||
if not self.lite_refpoint_refine:
|
||||
if self.bbox_reparam:
|
||||
obj_center, refpoints_input, query_pos, query_sine_embed = get_reference(refpoints_unsigmoid)
|
||||
else:
|
||||
obj_center, refpoints_input, query_pos, query_sine_embed = get_reference(
|
||||
refpoints_unsigmoid.sigmoid()
|
||||
)
|
||||
|
||||
# For the first decoder layer, we do not apply transformation over p_s
|
||||
pos_transformation = 1
|
||||
|
||||
query_pos = query_pos * pos_transformation
|
||||
|
||||
output = layer(
|
||||
output,
|
||||
memory,
|
||||
tgt_mask=tgt_mask,
|
||||
memory_mask=memory_mask,
|
||||
tgt_key_padding_mask=tgt_key_padding_mask,
|
||||
memory_key_padding_mask=memory_key_padding_mask,
|
||||
pos=pos,
|
||||
query_pos=query_pos,
|
||||
query_sine_embed=query_sine_embed,
|
||||
is_first=(layer_id == 0),
|
||||
reference_points=refpoints_input,
|
||||
spatial_shapes=spatial_shapes,
|
||||
level_start_index=level_start_index,
|
||||
)
|
||||
|
||||
if not self.lite_refpoint_refine:
|
||||
# box iterative update
|
||||
new_refpoints_delta = self.bbox_embed(output)
|
||||
new_refpoints_unsigmoid = self.refpoints_refine(refpoints_unsigmoid, new_refpoints_delta)
|
||||
if layer_id != self.num_layers - 1:
|
||||
hs_refpoints_unsigmoid.append(new_refpoints_unsigmoid)
|
||||
refpoints_unsigmoid = new_refpoints_unsigmoid.detach()
|
||||
|
||||
if self.return_intermediate:
|
||||
intermediate.append(self.norm(output))
|
||||
|
||||
if self.norm is not None:
|
||||
output = self.norm(output)
|
||||
if self.return_intermediate:
|
||||
intermediate.pop()
|
||||
intermediate.append(output)
|
||||
|
||||
if self.return_intermediate:
|
||||
if self._export:
|
||||
# to shape: B, N, C
|
||||
hs = intermediate[-1]
|
||||
if self.bbox_embed is not None:
|
||||
ref = hs_refpoints_unsigmoid[-1]
|
||||
else:
|
||||
ref = refpoints_unsigmoid
|
||||
return hs, ref
|
||||
# box iterative update
|
||||
if self.bbox_embed is not None:
|
||||
return [
|
||||
torch.stack(intermediate),
|
||||
torch.stack(hs_refpoints_unsigmoid),
|
||||
]
|
||||
else:
|
||||
return [torch.stack(intermediate), refpoints_unsigmoid.unsqueeze(0)]
|
||||
|
||||
return output.unsqueeze(0)
|
||||
|
||||
|
||||
class TransformerDecoderLayer(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
d_model,
|
||||
sa_nhead,
|
||||
ca_nhead,
|
||||
dim_feedforward=2048,
|
||||
dropout=0.1,
|
||||
activation="relu",
|
||||
normalize_before=False,
|
||||
group_detr=1,
|
||||
num_feature_levels=4,
|
||||
dec_n_points=4,
|
||||
skip_self_attn=False,
|
||||
):
|
||||
super().__init__()
|
||||
# Decoder Self-Attention
|
||||
self.self_attn = nn.MultiheadAttention(embed_dim=d_model, num_heads=sa_nhead, dropout=dropout, batch_first=True)
|
||||
self.dropout1 = nn.Dropout(dropout)
|
||||
self.norm1 = nn.LayerNorm(d_model)
|
||||
|
||||
# Decoder Cross-Attention
|
||||
self.cross_attn = MSDeformAttn(d_model, n_levels=num_feature_levels, n_heads=ca_nhead, n_points=dec_n_points)
|
||||
|
||||
self.nhead = ca_nhead
|
||||
|
||||
# Implementation of Feedforward model
|
||||
self.linear1 = nn.Linear(d_model, dim_feedforward)
|
||||
self.dropout = nn.Dropout(dropout)
|
||||
self.linear2 = nn.Linear(dim_feedforward, d_model)
|
||||
|
||||
self.norm2 = nn.LayerNorm(d_model)
|
||||
self.norm3 = nn.LayerNorm(d_model)
|
||||
|
||||
self.dropout2 = nn.Dropout(dropout)
|
||||
self.dropout3 = nn.Dropout(dropout)
|
||||
|
||||
self.activation = _get_activation_fn(activation)
|
||||
self.normalize_before = normalize_before
|
||||
self.group_detr = group_detr
|
||||
|
||||
def with_pos_embed(self, tensor, pos: Optional[Tensor]):
|
||||
return tensor if pos is None else tensor + pos
|
||||
|
||||
def forward_post(
|
||||
self,
|
||||
tgt,
|
||||
memory,
|
||||
tgt_mask: Optional[Tensor] = None,
|
||||
memory_mask: Optional[Tensor] = None,
|
||||
tgt_key_padding_mask: Optional[Tensor] = None,
|
||||
memory_key_padding_mask: Optional[Tensor] = None,
|
||||
pos: Optional[Tensor] = None,
|
||||
query_pos: Optional[Tensor] = None,
|
||||
query_sine_embed=None,
|
||||
is_first=False,
|
||||
reference_points=None,
|
||||
spatial_shapes=None,
|
||||
level_start_index=None,
|
||||
):
|
||||
bs, num_queries, _ = tgt.shape
|
||||
|
||||
# ========== Begin of Self-Attention =============
|
||||
# Apply projections here
|
||||
# shape: batch_size x num_queries x 256
|
||||
q = k = tgt + query_pos
|
||||
v = tgt
|
||||
if self.training:
|
||||
q = torch.cat(q.split(num_queries // self.group_detr, dim=1), dim=0)
|
||||
k = torch.cat(k.split(num_queries // self.group_detr, dim=1), dim=0)
|
||||
v = torch.cat(v.split(num_queries // self.group_detr, dim=1), dim=0)
|
||||
|
||||
tgt2 = self.self_attn(q, k, v, attn_mask=tgt_mask, key_padding_mask=tgt_key_padding_mask, need_weights=False)[0]
|
||||
|
||||
if self.training:
|
||||
tgt2 = torch.cat(tgt2.split(bs, dim=0), dim=1)
|
||||
# ========== End of Self-Attention =============
|
||||
|
||||
tgt = tgt + self.dropout1(tgt2)
|
||||
tgt = self.norm1(tgt)
|
||||
|
||||
# ========== Begin of Cross-Attention =============
|
||||
tgt2 = self.cross_attn(
|
||||
self.with_pos_embed(tgt, query_pos),
|
||||
reference_points,
|
||||
memory,
|
||||
spatial_shapes,
|
||||
level_start_index,
|
||||
memory_key_padding_mask,
|
||||
)
|
||||
# ========== End of Cross-Attention =============
|
||||
|
||||
tgt = tgt + self.dropout2(tgt2)
|
||||
tgt = self.norm2(tgt)
|
||||
tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
|
||||
tgt = tgt + self.dropout3(tgt2)
|
||||
tgt = self.norm3(tgt)
|
||||
return tgt
|
||||
|
||||
def forward(
|
||||
self,
|
||||
tgt,
|
||||
memory,
|
||||
tgt_mask: Optional[Tensor] = None,
|
||||
memory_mask: Optional[Tensor] = None,
|
||||
tgt_key_padding_mask: Optional[Tensor] = None,
|
||||
memory_key_padding_mask: Optional[Tensor] = None,
|
||||
pos: Optional[Tensor] = None,
|
||||
query_pos: Optional[Tensor] = None,
|
||||
query_sine_embed=None,
|
||||
is_first=False,
|
||||
reference_points=None,
|
||||
spatial_shapes=None,
|
||||
level_start_index=None,
|
||||
):
|
||||
return self.forward_post(
|
||||
tgt,
|
||||
memory,
|
||||
tgt_mask,
|
||||
memory_mask,
|
||||
tgt_key_padding_mask,
|
||||
memory_key_padding_mask,
|
||||
pos,
|
||||
query_pos,
|
||||
query_sine_embed,
|
||||
is_first,
|
||||
reference_points,
|
||||
spatial_shapes,
|
||||
level_start_index,
|
||||
)
|
||||
|
||||
|
||||
def _get_clones(module, N):
|
||||
return nn.ModuleList([copy.deepcopy(module) for i in range(N)])
|
||||
|
||||
|
||||
def build_transformer(args):
|
||||
|
||||
try:
|
||||
two_stage = args.two_stage
|
||||
except:
|
||||
two_stage = False
|
||||
|
||||
return Transformer(
|
||||
d_model=args.hidden_dim,
|
||||
sa_nhead=args.sa_nheads,
|
||||
ca_nhead=args.ca_nheads,
|
||||
num_queries=args.num_queries,
|
||||
dropout=args.dropout,
|
||||
dim_feedforward=args.dim_feedforward,
|
||||
num_decoder_layers=args.dec_layers,
|
||||
return_intermediate_dec=True,
|
||||
group_detr=args.group_detr,
|
||||
two_stage=two_stage,
|
||||
num_feature_levels=args.num_feature_levels,
|
||||
dec_n_points=args.dec_n_points,
|
||||
lite_refpoint_refine=args.lite_refpoint_refine,
|
||||
decoder_norm_type=args.decoder_norm,
|
||||
bbox_reparam=args.bbox_reparam,
|
||||
)
|
||||
|
||||
|
||||
def _get_activation_fn(activation):
|
||||
"""Return an activation function given a string"""
|
||||
if activation == "relu":
|
||||
return F.relu
|
||||
if activation == "gelu":
|
||||
return F.gelu
|
||||
if activation == "glu":
|
||||
return F.glu
|
||||
raise RuntimeError(f"activation should be relu/gelu, not {activation}.")
|
||||
@@ -0,0 +1,200 @@
|
||||
"""Vendored rf-detr (Roboflow RF-DETR) detection inference — no rfdetr package dependency.
|
||||
|
||||
Model definition copied (slimmed, detection-only) from the rfdetr package under
|
||||
``surya/common/rfdetr/models`` + ``util``. This module is the thin runtime wrapper:
|
||||
build the LWDETR architecture, load a fine-tuned checkpoint, and run predict() with the
|
||||
same preprocessing/post-processing the rfdetr package uses (ImageNet-normalize, resize to
|
||||
the model resolution, sigmoid top-k decode). Pure PyTorch — runs on cpu/mps/cuda.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
from typing import List, Optional
|
||||
|
||||
import torch
|
||||
import torchvision.transforms.functional as TF
|
||||
from PIL import Image
|
||||
|
||||
from surya.common.rfdetr.models import PostProcess, build_model
|
||||
|
||||
IMAGENET_MEAN = [0.485, 0.456, 0.406]
|
||||
IMAGENET_STD = [0.229, 0.224, 0.225]
|
||||
|
||||
LARGE_ARGS = { 'amp': True,
|
||||
'aug_config': None,
|
||||
'aux_loss': True,
|
||||
'backbone_lora': False,
|
||||
'backbone_only': False,
|
||||
'batch_size': 2,
|
||||
'bbox_loss_coef': 5,
|
||||
'bbox_reparam': True,
|
||||
'ca_nheads': 16,
|
||||
'checkpoint_interval': 10,
|
||||
'clip_max_norm': 0.1,
|
||||
'cls_loss_coef': 1.0,
|
||||
'coco_path': None,
|
||||
'cutoff_epoch': 0,
|
||||
'dataset_dir': None,
|
||||
'dataset_file': 'coco',
|
||||
'dec_layers': 4,
|
||||
'dec_n_points': 2,
|
||||
'decoder_norm': 'LN',
|
||||
'dim_feedforward': 2048,
|
||||
'dist_url': 'env://',
|
||||
'do_benchmark': False,
|
||||
'do_random_resize_via_padding': False,
|
||||
'dont_save_weights': False,
|
||||
'drop_mode': 'standard',
|
||||
'drop_path': 0,
|
||||
'drop_schedule': 'constant',
|
||||
'dropout': 0,
|
||||
'early_stopping': True,
|
||||
'early_stopping_min_delta': 0.001,
|
||||
'early_stopping_patience': 10,
|
||||
'early_stopping_use_ema': False,
|
||||
'ema_decay': 0.9997,
|
||||
'ema_tau': 0,
|
||||
'encoder': 'dinov2_windowed_small',
|
||||
'encoder_only': False,
|
||||
'epochs': 12,
|
||||
'eval': False,
|
||||
'expanded_scales': False,
|
||||
'focal_alpha': 0.25,
|
||||
'force_no_pretrain': False,
|
||||
'fp16_eval': False,
|
||||
'freeze_batch_norm': False,
|
||||
'freeze_encoder': False,
|
||||
'giou_loss_coef': 2,
|
||||
'grad_accum_steps': 1,
|
||||
'gradient_checkpointing': False,
|
||||
'group_detr': 13,
|
||||
'hidden_dim': 256,
|
||||
'ia_bce_loss': True,
|
||||
'layer_norm': True,
|
||||
'license': 'Apache-2.0',
|
||||
'lite_refpoint_refine': True,
|
||||
'lr': 0.0001,
|
||||
'lr_component_decay': 1.0,
|
||||
'lr_drop': 11,
|
||||
'lr_encoder': 0.00015,
|
||||
'lr_min_factor': 0.0,
|
||||
'lr_scheduler': 'step',
|
||||
'lr_vit_layer_decay': 0.8,
|
||||
'mask_downsample_ratio': 4,
|
||||
'multi_scale': False,
|
||||
'num_classes': 90,
|
||||
'num_feature_levels': 1,
|
||||
'num_queries': 300,
|
||||
'num_select': 100,
|
||||
'num_windows': 2,
|
||||
'num_workers': 2,
|
||||
'out_feature_indexes': [3, 6, 9, 12],
|
||||
'output_dir': 'output',
|
||||
'patch_size': 16,
|
||||
'position_embedding': 'sine',
|
||||
'positional_encoding_size': 44,
|
||||
'pretrain_exclude_keys': None,
|
||||
'pretrain_keys_modify_to_load': None,
|
||||
'pretrained_distiller': None,
|
||||
'pretrained_encoder': None,
|
||||
'print_freq': 10,
|
||||
'projector_scale': ['P4'],
|
||||
'resolution': 704,
|
||||
'resume': '',
|
||||
'rms_norm': False,
|
||||
'sa_nheads': 8,
|
||||
'seed': 42,
|
||||
'segmentation_head': False,
|
||||
'set_cost_bbox': 5,
|
||||
'set_cost_class': 2,
|
||||
'set_cost_giou': 2,
|
||||
'square_resize_div_64': False,
|
||||
'start_epoch': 0,
|
||||
'sum_group_losses': False,
|
||||
'sync_bn': True,
|
||||
'two_stage': True,
|
||||
'use_cls_token': False,
|
||||
'use_ema': False,
|
||||
'use_position_supervised_loss': False,
|
||||
'use_varifocal_loss': False,
|
||||
'vit_encoder_num_layers': 12,
|
||||
'warmup_epochs': 1,
|
||||
'weight_decay': 0.0001,
|
||||
'window_block_indexes': None,
|
||||
'world_size': 1}
|
||||
|
||||
|
||||
class RFDetrDetector:
|
||||
"""Builds RF-DETR-Large and runs detection. ``predict`` returns, per image, a dict with
|
||||
``boxes`` (xyxy pixels), ``scores``, ``labels`` (0-indexed class ids) as CPU tensors."""
|
||||
|
||||
def __init__(self, weights_path: str, device: str = "cpu",
|
||||
arch_args: Optional[dict] = None):
|
||||
args = SimpleNamespace(**{**LARGE_ARGS, **(arch_args or {})})
|
||||
args.device = device
|
||||
# Truthy so build_backbone sets load_dinov2_weights=False (no DINOv2 hub download);
|
||||
# the fine-tuned checkpoint is loaded manually below, not by build_model.
|
||||
args.pretrain_weights = weights_path
|
||||
self.device = torch.device(device)
|
||||
self.resolution = int(args.resolution)
|
||||
self.num_select = int(args.num_select)
|
||||
|
||||
model = build_model(args)
|
||||
|
||||
ckpt = torch.load(weights_path, map_location="cpu", weights_only=False)
|
||||
state = ckpt["model"]
|
||||
# Match the checkpoint's class count (head was configured for the default 90).
|
||||
ckpt_num_classes = state["class_embed.bias"].shape[0]
|
||||
if ckpt_num_classes != args.num_classes + 1:
|
||||
model.reinitialize_detection_head(ckpt_num_classes)
|
||||
# Trim group-detr query params to the desired query count (no-op if already matching).
|
||||
num_desired = args.num_queries * args.group_detr
|
||||
for name in list(state.keys()):
|
||||
if name.endswith("refpoint_embed.weight") or name.endswith("query_feat.weight"):
|
||||
state[name] = state[name][:num_desired]
|
||||
model.load_state_dict(state, strict=False)
|
||||
|
||||
self.model = model.eval().to(self.device)
|
||||
self.postprocess = PostProcess(num_select=self.num_select)
|
||||
|
||||
# Optional capture of the encoder feature map (projector output, [B,C,F,F]) so the
|
||||
# reading-order head can cross-attend to it. Hook the LWDETR backbone projector.
|
||||
self._feat = {}
|
||||
for name, mod in self.model.named_modules():
|
||||
if name.endswith("projector"):
|
||||
mod.register_forward_hook(
|
||||
lambda m, i, o: self._feat.__setitem__(
|
||||
"f", (o[0] if isinstance(o, (tuple, list)) else o)))
|
||||
break
|
||||
|
||||
@torch.inference_mode()
|
||||
def predict(self, images: List[Image.Image], threshold: float = 0.4,
|
||||
return_features: bool = False) -> List[dict]:
|
||||
if not images:
|
||||
return []
|
||||
tensors, sizes = [], []
|
||||
for img in images:
|
||||
img = img.convert("RGB")
|
||||
sizes.append((img.height, img.width)) # (h, w) for PostProcess scaling
|
||||
t = TF.to_tensor(img).to(self.device)
|
||||
t = TF.normalize(t, IMAGENET_MEAN, IMAGENET_STD)
|
||||
t = TF.resize(t, (self.resolution, self.resolution))
|
||||
tensors.append(t)
|
||||
outputs = self.model(torch.stack(tensors, 0))
|
||||
feats = self._feat.get("f") if return_features else None # [B,C,F,F]
|
||||
target_sizes = torch.tensor(sizes, device=self.device)
|
||||
results = self.postprocess(outputs, target_sizes=target_sizes)
|
||||
|
||||
out = []
|
||||
for i, res in enumerate(results):
|
||||
keep = res["scores"] > threshold
|
||||
d = {
|
||||
"boxes": res["boxes"][keep].cpu(),
|
||||
"scores": res["scores"][keep].cpu(),
|
||||
"labels": res["labels"][keep].cpu(),
|
||||
}
|
||||
if return_features and feats is not None:
|
||||
d["features"] = feats[i].cpu() # [C,F,F] for this page
|
||||
out.append(d)
|
||||
return out
|
||||
@@ -0,0 +1,16 @@
|
||||
# ------------------------------------------------------------------------
|
||||
# RF-DETR
|
||||
# Copyright (c) 2025 Roboflow. All Rights Reserved.
|
||||
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
|
||||
# ------------------------------------------------------------------------
|
||||
# Copied and modified from LW-DETR (https://github.com/Atten4Vis/LW-DETR)
|
||||
# Copyright (c) 2024 Baidu. All Rights Reserved.
|
||||
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
|
||||
# ------------------------------------------------------------------------
|
||||
# Conditional DETR
|
||||
# Copyright (c) 2021 Microsoft. All Rights Reserved.
|
||||
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
|
||||
# ------------------------------------------------------------------------
|
||||
# Copied from DETR (https://github.com/facebookresearch/detr)
|
||||
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
|
||||
# ------------------------------------------------------------------------
|
||||
@@ -0,0 +1,161 @@
|
||||
# ------------------------------------------------------------------------
|
||||
# RF-DETR
|
||||
# Copyright (c) 2025 Roboflow. All Rights Reserved.
|
||||
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
|
||||
# ------------------------------------------------------------------------
|
||||
# Copied and modified from LW-DETR (https://github.com/Atten4Vis/LW-DETR)
|
||||
# Copyright (c) 2024 Baidu. All Rights Reserved.
|
||||
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
|
||||
# ------------------------------------------------------------------------
|
||||
# Conditional DETR
|
||||
# Copyright (c) 2021 Microsoft. All Rights Reserved.
|
||||
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
|
||||
# ------------------------------------------------------------------------
|
||||
# Copied from DETR (https://github.com/facebookresearch/detr)
|
||||
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
|
||||
# ------------------------------------------------------------------------
|
||||
|
||||
"""
|
||||
Utilities for bounding box manipulation and GIoU.
|
||||
"""
|
||||
|
||||
from typing import Tuple
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torchvision.ops.boxes import box_area
|
||||
|
||||
|
||||
def box_cxcywh_to_xyxy(x: torch.Tensor) -> torch.Tensor:
|
||||
x_c, y_c, w, h = x.unbind(-1)
|
||||
b = [
|
||||
(x_c - 0.5 * w.clamp(min=0.0)),
|
||||
(y_c - 0.5 * h.clamp(min=0.0)),
|
||||
(x_c + 0.5 * w.clamp(min=0.0)),
|
||||
(y_c + 0.5 * h.clamp(min=0.0)),
|
||||
]
|
||||
return torch.stack(b, dim=-1)
|
||||
|
||||
|
||||
def box_xyxy_to_cxcywh(x: torch.Tensor) -> torch.Tensor:
|
||||
x0, y0, x1, y1 = x.unbind(-1)
|
||||
b = [(x0 + x1) / 2, (y0 + y1) / 2, (x1 - x0), (y1 - y0)]
|
||||
return torch.stack(b, dim=-1)
|
||||
|
||||
|
||||
# modified from torchvision to also return the union
|
||||
def box_iou(boxes1: torch.Tensor, boxes2: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Returns:
|
||||
iou: the NxM matrix containing the pairwise
|
||||
IoU values for every element in boxes1 and boxes2
|
||||
union: the NxM matrix containing the pairwise
|
||||
union values for every element in boxes1 and boxes2
|
||||
"""
|
||||
area1 = box_area(boxes1)
|
||||
area2 = box_area(boxes2)
|
||||
|
||||
lt = torch.max(boxes1[:, None, :2], boxes2[:, :2]) # [N,M,2]
|
||||
rb = torch.min(boxes1[:, None, 2:], boxes2[:, 2:]) # [N,M,2]
|
||||
|
||||
wh = (rb - lt).clamp(min=0) # [N,M,2]
|
||||
inter = wh[:, :, 0] * wh[:, :, 1] # [N,M]
|
||||
|
||||
union = area1[:, None] + area2 - inter
|
||||
|
||||
iou = inter / union
|
||||
return iou, union
|
||||
|
||||
|
||||
def generalized_box_iou(boxes1: torch.Tensor, boxes2: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Generalized IoU from https://giou.stanford.edu/
|
||||
|
||||
The boxes should be in [x0, y0, x1, y1] format
|
||||
|
||||
Returns a [N, M] pairwise matrix, where N = len(boxes1)
|
||||
and M = len(boxes2)
|
||||
"""
|
||||
# degenerate boxes gives inf / nan results
|
||||
# so do an early check
|
||||
iou, union = box_iou(boxes1, boxes2)
|
||||
|
||||
lt = torch.min(boxes1[:, None, :2], boxes2[:, :2])
|
||||
rb = torch.max(boxes1[:, None, 2:], boxes2[:, 2:])
|
||||
|
||||
wh = (rb - lt).clamp(min=0) # [N,M,2]
|
||||
area = wh[:, :, 0] * wh[:, :, 1]
|
||||
|
||||
return iou - (area - union) / area
|
||||
|
||||
|
||||
def masks_to_boxes(masks: torch.Tensor) -> torch.Tensor:
|
||||
"""Compute the bounding boxes around the provided masks
|
||||
|
||||
The masks should be in format [N, H, W] where N is the number of masks, (H, W) are the spatial dimensions.
|
||||
|
||||
Returns a [N, 4] tensors, with the boxes in xyxy format
|
||||
"""
|
||||
if masks.numel() == 0:
|
||||
return torch.zeros((0, 4), device=masks.device)
|
||||
|
||||
h, w = masks.shape[-2:]
|
||||
|
||||
y = torch.arange(0, h, dtype=torch.float)
|
||||
x = torch.arange(0, w, dtype=torch.float)
|
||||
y, x = torch.meshgrid(y, x, indexing="ij")
|
||||
|
||||
x_mask = masks * x.unsqueeze(0)
|
||||
x_max = x_mask.flatten(1).max(-1)[0]
|
||||
x_min = x_mask.masked_fill(~(masks.bool()), 1e8).flatten(1).min(-1)[0]
|
||||
|
||||
y_mask = masks * y.unsqueeze(0)
|
||||
y_max = y_mask.flatten(1).max(-1)[0]
|
||||
y_min = y_mask.masked_fill(~(masks.bool()), 1e8).flatten(1).min(-1)[0]
|
||||
|
||||
return torch.stack([x_min, y_min, x_max, y_max], 1)
|
||||
|
||||
|
||||
def batch_dice_loss(inputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Compute the DICE loss, similar to generalized IOU for masks
|
||||
Args:
|
||||
inputs: A float tensor of arbitrary shape.
|
||||
The predictions for each example.
|
||||
targets: A float tensor with the same shape as inputs. Stores the binary
|
||||
classification label for each element in inputs
|
||||
(0 for the negative class and 1 for the positive class).
|
||||
"""
|
||||
inputs = inputs.sigmoid()
|
||||
inputs = inputs.flatten(1)
|
||||
numerator = 2 * torch.einsum("nc,mc->nm", inputs, targets)
|
||||
denominator = inputs.sum(-1)[:, None] + targets.sum(-1)[None, :]
|
||||
loss = 1 - (numerator + 1) / (denominator + 1)
|
||||
return loss
|
||||
|
||||
|
||||
batch_dice_loss_jit = torch.jit.script(batch_dice_loss) # type: torch.jit.ScriptModule
|
||||
|
||||
|
||||
def batch_sigmoid_ce_loss(inputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Args:
|
||||
inputs: A float tensor of arbitrary shape.
|
||||
The predictions for each example.
|
||||
targets: A float tensor with the same shape as inputs. Stores the binary
|
||||
classification label for each element in inputs
|
||||
(0 for the negative class and 1 for the positive class).
|
||||
Returns:
|
||||
Loss tensor
|
||||
"""
|
||||
hw = inputs.shape[1]
|
||||
|
||||
pos = F.binary_cross_entropy_with_logits(inputs, torch.ones_like(inputs), reduction="none")
|
||||
neg = F.binary_cross_entropy_with_logits(inputs, torch.zeros_like(inputs), reduction="none")
|
||||
|
||||
loss = torch.einsum("nc,mc->nm", pos, targets) + torch.einsum("nc,mc->nm", neg, (1 - targets))
|
||||
|
||||
return loss / hw
|
||||
|
||||
|
||||
batch_sigmoid_ce_loss_jit = torch.jit.script(batch_sigmoid_ce_loss) # type: torch.jit.ScriptModule
|
||||
@@ -0,0 +1,55 @@
|
||||
# ------------------------------------------------------------------------
|
||||
# RF-DETR
|
||||
# Copyright (c) 2025 Roboflow. All Rights Reserved.
|
||||
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
|
||||
# ------------------------------------------------------------------------
|
||||
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
from typing import Optional
|
||||
|
||||
|
||||
def get_logger(name: str = "rf-detr", level: Optional[int] = None) -> logging.Logger:
|
||||
"""Creates and configures a logger with stdout and stderr handlers.
|
||||
|
||||
This function creates a logger that sends INFO and DEBUG level logs to stdout,
|
||||
and WARNING, ERROR, and CRITICAL level logs to stderr. If the logger already
|
||||
has handlers, it returns the existing logger without adding new handlers.
|
||||
|
||||
The log level can be specified directly or through the LOG_LEVEL environment
|
||||
variable.
|
||||
|
||||
Args:
|
||||
name: The name of the logger. Defaults to "rf-detr".
|
||||
level: The logging level to set. If None, uses the LOG_LEVEL environment
|
||||
variable, defaulting to INFO if not set.
|
||||
|
||||
Returns:
|
||||
A configured logging.Logger instance.
|
||||
"""
|
||||
if level is None:
|
||||
level = getattr(logging, os.getenv("LOG_LEVEL", "INFO").upper(), logging.INFO)
|
||||
|
||||
logger = logging.getLogger(name)
|
||||
logger.setLevel(level)
|
||||
|
||||
if not logger.handlers:
|
||||
formatter = logging.Formatter(
|
||||
"[%(asctime)s] [%(levelname)s] %(name)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S"
|
||||
)
|
||||
|
||||
stdout_handler = logging.StreamHandler(sys.stdout)
|
||||
stdout_handler.setLevel(logging.DEBUG)
|
||||
stdout_handler.addFilter(lambda r: r.levelno <= logging.INFO)
|
||||
stdout_handler.setFormatter(formatter)
|
||||
|
||||
stderr_handler = logging.StreamHandler(sys.stderr)
|
||||
stderr_handler.setLevel(logging.WARNING)
|
||||
stderr_handler.setFormatter(formatter)
|
||||
|
||||
logger.addHandler(stdout_handler)
|
||||
logger.addHandler(stderr_handler)
|
||||
logger.propagate = False
|
||||
|
||||
return logger
|
||||
@@ -0,0 +1,531 @@
|
||||
# ------------------------------------------------------------------------
|
||||
# RF-DETR
|
||||
# Copyright (c) 2025 Roboflow. All Rights Reserved.
|
||||
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
|
||||
# ------------------------------------------------------------------------
|
||||
# Copied and modified from LW-DETR (https://github.com/Atten4Vis/LW-DETR)
|
||||
# Copyright (c) 2024 Baidu. All Rights Reserved.
|
||||
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
|
||||
# ------------------------------------------------------------------------
|
||||
# Conditional DETR
|
||||
# Copyright (c) 2021 Microsoft. All Rights Reserved.
|
||||
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
|
||||
# ------------------------------------------------------------------------
|
||||
# Copied from DETR (https://github.com/facebookresearch/detr)
|
||||
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
|
||||
# ------------------------------------------------------------------------
|
||||
|
||||
"""
|
||||
Misc functions, including distributed helpers.
|
||||
|
||||
Mostly copy-paste from torchvision references.
|
||||
"""
|
||||
|
||||
import datetime
|
||||
import os
|
||||
import pickle
|
||||
import subprocess
|
||||
import tempfile
|
||||
import time
|
||||
from collections import defaultdict, deque
|
||||
from typing import Any, Dict, Generator, Iterable, List, Optional, Tuple
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
# needed due to empty tensor bug in pytorch and torchvision 0.5
|
||||
import torchvision
|
||||
from torch import Tensor
|
||||
|
||||
from surya.common.rfdetr.util.logger import get_logger
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
if float(torchvision.__version__.split(".")[1]) < 7.0:
|
||||
from torchvision.ops import _new_empty_tensor
|
||||
from torchvision.ops.misc import _output_size
|
||||
|
||||
|
||||
class SmoothedValue(object):
|
||||
"""Track a series of values and provide access to smoothed values over a
|
||||
window or the global series average.
|
||||
"""
|
||||
|
||||
def __init__(self, window_size: int = 20, fmt: Optional[str] = None) -> None:
|
||||
if fmt is None:
|
||||
fmt = "{median:.4f} ({global_avg:.4f})"
|
||||
self.deque = deque(maxlen=window_size)
|
||||
self.total = 0.0
|
||||
self.count = 0
|
||||
self.fmt = fmt
|
||||
|
||||
def update(self, value: float, n: int = 1) -> None:
|
||||
self.deque.append(value)
|
||||
self.count += n
|
||||
self.total += value * n
|
||||
|
||||
def synchronize_between_processes(self) -> None:
|
||||
"""
|
||||
Warning: does not synchronize the deque!
|
||||
"""
|
||||
if not is_dist_avail_and_initialized():
|
||||
return
|
||||
t = torch.tensor([self.count, self.total], dtype=torch.float64, device="cuda")
|
||||
dist.barrier()
|
||||
dist.all_reduce(t)
|
||||
t = t.tolist()
|
||||
self.count = int(t[0])
|
||||
self.total = t[1]
|
||||
|
||||
@property
|
||||
def median(self) -> float:
|
||||
d = torch.tensor(list(self.deque))
|
||||
return d.median().item()
|
||||
|
||||
@property
|
||||
def avg(self) -> float:
|
||||
d = torch.tensor(list(self.deque), dtype=torch.float32)
|
||||
return d.mean().item()
|
||||
|
||||
@property
|
||||
def global_avg(self) -> float:
|
||||
return self.total / self.count
|
||||
|
||||
@property
|
||||
def max(self) -> float:
|
||||
return max(self.deque)
|
||||
|
||||
@property
|
||||
def value(self) -> float:
|
||||
return self.deque[-1]
|
||||
|
||||
def __str__(self) -> str:
|
||||
return self.fmt.format(
|
||||
median=self.median, avg=self.avg, global_avg=self.global_avg, max=self.max, value=self.value
|
||||
)
|
||||
|
||||
|
||||
def all_gather(data: Any) -> List[Any]:
|
||||
"""
|
||||
Run all_gather on arbitrary picklable data (not necessarily tensors)
|
||||
Args:
|
||||
data: any picklable object
|
||||
Returns:
|
||||
list of data gathered from each rank
|
||||
"""
|
||||
world_size = get_world_size()
|
||||
if world_size == 1:
|
||||
return [data]
|
||||
|
||||
# serialized to a Tensor
|
||||
buffer = pickle.dumps(data)
|
||||
storage = torch.ByteStorage.from_buffer(buffer)
|
||||
tensor = torch.ByteTensor(storage).to("cuda")
|
||||
|
||||
# obtain Tensor size of each rank
|
||||
local_size = torch.tensor([tensor.numel()], device="cuda")
|
||||
size_list = [torch.tensor([0], device="cuda") for _ in range(world_size)]
|
||||
dist.all_gather(size_list, local_size)
|
||||
size_list = [int(size.item()) for size in size_list]
|
||||
max_size = max(size_list)
|
||||
|
||||
# receiving Tensor from all ranks
|
||||
# we pad the tensor because torch all_gather does not support
|
||||
# gathering tensors of different shapes
|
||||
tensor_list = []
|
||||
for _ in size_list:
|
||||
tensor_list.append(torch.empty((max_size,), dtype=torch.uint8, device="cuda"))
|
||||
if local_size != max_size:
|
||||
padding = torch.empty(size=(max_size - local_size,), dtype=torch.uint8, device="cuda")
|
||||
tensor = torch.cat((tensor, padding), dim=0)
|
||||
dist.all_gather(tensor_list, tensor)
|
||||
|
||||
data_list = []
|
||||
for size, tensor in zip(size_list, tensor_list):
|
||||
buffer = tensor.cpu().numpy().tobytes()[:size]
|
||||
data_list.append(pickle.loads(buffer))
|
||||
|
||||
return data_list
|
||||
|
||||
|
||||
def reduce_dict(input_dict: Dict[str, torch.Tensor], average: bool = True) -> Dict[str, torch.Tensor]:
|
||||
"""
|
||||
Args:
|
||||
input_dict (dict): all the values will be reduced
|
||||
average (bool): whether to do average or sum
|
||||
Reduce the values in the dictionary from all processes so that all processes
|
||||
have the averaged results. Returns a dict with the same fields as
|
||||
input_dict, after reduction.
|
||||
"""
|
||||
world_size = get_world_size()
|
||||
if world_size < 2:
|
||||
return input_dict
|
||||
with torch.no_grad():
|
||||
names = []
|
||||
values = []
|
||||
# sort the keys so that they are consistent across processes
|
||||
for k in sorted(input_dict.keys()):
|
||||
names.append(k)
|
||||
values.append(input_dict[k])
|
||||
values = torch.stack(values, dim=0)
|
||||
dist.all_reduce(values)
|
||||
if average:
|
||||
values /= world_size
|
||||
reduced_dict = {k: v for k, v in zip(names, values)}
|
||||
return reduced_dict
|
||||
|
||||
|
||||
class MetricLogger(object):
|
||||
def __init__(self, delimiter: str = "\t", wandb_logging: bool = False) -> None:
|
||||
self.meters = defaultdict(SmoothedValue)
|
||||
self.delimiter = delimiter
|
||||
if wandb_logging:
|
||||
import wandb
|
||||
|
||||
self.wandb = wandb
|
||||
else:
|
||||
self.wandb = None
|
||||
|
||||
def update(self, **kwargs: Any) -> None:
|
||||
for k, v in kwargs.items():
|
||||
if isinstance(v, torch.Tensor):
|
||||
v = v.item()
|
||||
assert isinstance(v, (float, int))
|
||||
self.meters[k].update(v)
|
||||
|
||||
def __getattr__(self, attr: str) -> SmoothedValue:
|
||||
if attr in self.meters:
|
||||
return self.meters[attr]
|
||||
if attr in self.__dict__:
|
||||
return self.__dict__[attr]
|
||||
raise AttributeError("'{}' object has no attribute '{}'".format(type(self).__name__, attr))
|
||||
|
||||
def __str__(self) -> str:
|
||||
loss_str = []
|
||||
for name, meter in self.meters.items():
|
||||
loss_str.append("{}: {}".format(name, str(meter)))
|
||||
return self.delimiter.join(loss_str)
|
||||
|
||||
def synchronize_between_processes(self) -> None:
|
||||
for meter in self.meters.values():
|
||||
meter.synchronize_between_processes()
|
||||
|
||||
def add_meter(self, name: str, meter: SmoothedValue) -> None:
|
||||
self.meters[name] = meter
|
||||
|
||||
def log_every(
|
||||
self, iterable: Iterable[Any], print_freq: int, header: Optional[str] = None
|
||||
) -> Generator[Any, None, None]:
|
||||
i = 0
|
||||
if not header:
|
||||
header = ""
|
||||
start_time = time.time()
|
||||
end = time.time()
|
||||
iter_time = SmoothedValue(fmt="{avg:.4f}")
|
||||
data_time = SmoothedValue(fmt="{avg:.4f}")
|
||||
space_fmt = ":" + str(len(str(len(iterable)))) + "d"
|
||||
if torch.cuda.is_available():
|
||||
log_msg = self.delimiter.join(
|
||||
[
|
||||
header,
|
||||
"[{0" + space_fmt + "}/{1}]",
|
||||
"eta: {eta}",
|
||||
"{meters}",
|
||||
"time: {time}",
|
||||
"data: {data}",
|
||||
"max mem: {memory:.0f}",
|
||||
]
|
||||
)
|
||||
else:
|
||||
log_msg = self.delimiter.join(
|
||||
[header, "[{0" + space_fmt + "}/{1}]", "eta: {eta}", "{meters}", "time: {time}", "data: {data}"]
|
||||
)
|
||||
MB = 1024.0 * 1024.0
|
||||
for obj in iterable:
|
||||
data_time.update(time.time() - end)
|
||||
yield obj
|
||||
iter_time.update(time.time() - end)
|
||||
if i % print_freq == 0 or i == len(iterable) - 1:
|
||||
eta_seconds = iter_time.global_avg * (len(iterable) - i)
|
||||
eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
|
||||
if self.wandb:
|
||||
if is_main_process():
|
||||
log_dict = {k: v.value for k, v in self.meters.items()}
|
||||
self.wandb.log(log_dict)
|
||||
if torch.cuda.is_available():
|
||||
logger.info(
|
||||
log_msg.format(
|
||||
i,
|
||||
len(iterable),
|
||||
eta=eta_string,
|
||||
meters=str(self),
|
||||
time=str(iter_time),
|
||||
data=str(data_time),
|
||||
memory=torch.cuda.max_memory_allocated() / MB,
|
||||
)
|
||||
)
|
||||
else:
|
||||
logger.info(
|
||||
log_msg.format(
|
||||
i, len(iterable), eta=eta_string, meters=str(self), time=str(iter_time), data=str(data_time)
|
||||
)
|
||||
)
|
||||
i += 1
|
||||
end = time.time()
|
||||
total_time = time.time() - start_time
|
||||
total_time_str = str(datetime.timedelta(seconds=int(total_time)))
|
||||
logger.info("{} Total time: {} ({:.4f} s / it)".format(header, total_time_str, total_time / len(iterable)))
|
||||
|
||||
|
||||
def get_sha() -> str:
|
||||
"""Return a short status string for the current git repo, or 'unknown' if unavailable."""
|
||||
cwd = os.path.dirname(os.path.abspath(__file__))
|
||||
|
||||
def _run(command: List[str]) -> str:
|
||||
return subprocess.check_output(command, cwd=cwd).decode("ascii").strip()
|
||||
|
||||
try:
|
||||
sha = _run(["git", "rev-parse", "HEAD"])
|
||||
has_diff = bool(_run(["git", "diff-index", "HEAD"]))
|
||||
status = "has uncommitted changes" if has_diff else "clean"
|
||||
branch = _run(["git", "rev-parse", "--abbrev-ref", "HEAD"])
|
||||
return f"sha: {sha}, status: {status}, branch: {branch}"
|
||||
except (subprocess.CalledProcessError, FileNotFoundError):
|
||||
return "unknown"
|
||||
|
||||
|
||||
def collate_fn(batch: List[Tuple[Any, ...]]) -> Tuple[Any, ...]:
|
||||
batch = list(zip(*batch))
|
||||
batch[0] = nested_tensor_from_tensor_list(batch[0])
|
||||
return tuple(batch)
|
||||
|
||||
|
||||
def _max_by_axis(the_list: List[List[int]]) -> List[int]:
|
||||
maxes = the_list[0]
|
||||
for sublist in the_list[1:]:
|
||||
for index, item in enumerate(sublist):
|
||||
maxes[index] = max(maxes[index], item)
|
||||
return maxes
|
||||
|
||||
|
||||
class NestedTensor(object):
|
||||
def __init__(self, tensors: Tensor, mask: Optional[Tensor]) -> None:
|
||||
self.tensors = tensors
|
||||
self.mask = mask
|
||||
|
||||
def to(self, device: torch.device) -> "NestedTensor":
|
||||
cast_tensor = self.tensors.to(device)
|
||||
mask = self.mask
|
||||
if mask is not None:
|
||||
assert mask is not None
|
||||
cast_mask = mask.to(device)
|
||||
else:
|
||||
cast_mask = None
|
||||
return NestedTensor(cast_tensor, cast_mask)
|
||||
|
||||
def decompose(self) -> Tuple[Tensor, Optional[Tensor]]:
|
||||
return self.tensors, self.mask
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return str(self.tensors)
|
||||
|
||||
|
||||
def nested_tensor_from_tensor_list(tensor_list: List[Tensor]) -> NestedTensor:
|
||||
# TODO make this more general
|
||||
if tensor_list[0].ndim == 3:
|
||||
if torchvision._is_tracing():
|
||||
# nested_tensor_from_tensor_list() does not export well to ONNX
|
||||
# call _onnx_nested_tensor_from_tensor_list() instead
|
||||
return _onnx_nested_tensor_from_tensor_list(tensor_list)
|
||||
|
||||
# TODO make it support different-sized images
|
||||
max_size = _max_by_axis([list(img.shape) for img in tensor_list])
|
||||
# min_size = tuple(min(s) for s in zip(*[img.shape for img in tensor_list]))
|
||||
batch_shape = [len(tensor_list)] + max_size
|
||||
b, c, h, w = batch_shape
|
||||
dtype = tensor_list[0].dtype
|
||||
device = tensor_list[0].device
|
||||
tensor = torch.zeros(batch_shape, dtype=dtype, device=device)
|
||||
mask = torch.ones((b, h, w), dtype=torch.bool, device=device)
|
||||
for img, pad_img, m in zip(tensor_list, tensor, mask):
|
||||
pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img)
|
||||
m[: img.shape[1], : img.shape[2]] = False
|
||||
else:
|
||||
raise ValueError("not supported")
|
||||
return NestedTensor(tensor, mask)
|
||||
|
||||
|
||||
# _onnx_nested_tensor_from_tensor_list() is an implementation of
|
||||
# nested_tensor_from_tensor_list() that is supported by ONNX tracing.
|
||||
@torch.jit.unused
|
||||
def _onnx_nested_tensor_from_tensor_list(tensor_list: List[Tensor]) -> NestedTensor:
|
||||
max_size = []
|
||||
for i in range(tensor_list[0].dim()):
|
||||
max_size_i = torch.max(torch.stack([img.shape[i] for img in tensor_list]).to(torch.float32)).to(torch.int64)
|
||||
max_size.append(max_size_i)
|
||||
max_size = tuple(max_size)
|
||||
|
||||
# work around for
|
||||
# pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img)
|
||||
# m[: img.shape[1], :img.shape[2]] = False
|
||||
# which is not yet supported in onnx
|
||||
padded_imgs = []
|
||||
padded_masks = []
|
||||
for img in tensor_list:
|
||||
padding = [(s1 - s2) for s1, s2 in zip(max_size, tuple(img.shape))]
|
||||
padded_img = torch.nn.functional.pad(img, (0, padding[2], 0, padding[1], 0, padding[0]))
|
||||
padded_imgs.append(padded_img)
|
||||
|
||||
m = torch.zeros_like(img[0], dtype=torch.int, device=img.device)
|
||||
padded_mask = torch.nn.functional.pad(m, (0, padding[2], 0, padding[1]), "constant", 1)
|
||||
padded_masks.append(padded_mask.to(torch.bool))
|
||||
|
||||
tensor = torch.stack(padded_imgs)
|
||||
mask = torch.stack(padded_masks)
|
||||
|
||||
return NestedTensor(tensor, mask=mask)
|
||||
|
||||
|
||||
def setup_for_distributed(is_master: bool) -> None:
|
||||
"""
|
||||
This function disables printing when not in master process
|
||||
"""
|
||||
import builtins as __builtin__
|
||||
import logging
|
||||
|
||||
builtin_print = __builtin__.print
|
||||
|
||||
def print(*args, **kwargs) -> None:
|
||||
force = kwargs.pop("force", False)
|
||||
if is_master or force:
|
||||
builtin_print(*args, **kwargs)
|
||||
|
||||
__builtin__.print = print
|
||||
|
||||
if not is_master:
|
||||
logging.getLogger("rf-detr").setLevel(logging.ERROR)
|
||||
|
||||
|
||||
def is_dist_avail_and_initialized():
|
||||
if not dist.is_available():
|
||||
return False
|
||||
if not dist.is_initialized():
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def get_world_size():
|
||||
if not is_dist_avail_and_initialized():
|
||||
return 1
|
||||
return dist.get_world_size()
|
||||
|
||||
|
||||
def get_rank():
|
||||
if not is_dist_avail_and_initialized():
|
||||
return 0
|
||||
return dist.get_rank()
|
||||
|
||||
|
||||
def is_main_process():
|
||||
return get_rank() == 0
|
||||
|
||||
|
||||
def save_on_master(obj, f, *args, **kwargs):
|
||||
"""
|
||||
Safely save objects, removing any callbacks that can't be pickled
|
||||
"""
|
||||
if is_main_process():
|
||||
torch.save(obj, f, *args, **kwargs)
|
||||
|
||||
|
||||
def init_distributed_mode(args: Any) -> None:
|
||||
if "RANK" in os.environ and "WORLD_SIZE" in os.environ:
|
||||
args.rank = int(os.environ["RANK"])
|
||||
args.world_size = int(os.environ["WORLD_SIZE"])
|
||||
args.gpu = int(os.environ["LOCAL_RANK"])
|
||||
elif "SLURM_PROCID" in os.environ:
|
||||
args.rank = int(os.environ["SLURM_PROCID"])
|
||||
args.gpu = args.rank % torch.cuda.device_count()
|
||||
else:
|
||||
logger.info("Not using distributed mode")
|
||||
args.distributed = False
|
||||
return
|
||||
|
||||
args.distributed = True
|
||||
|
||||
torch.cuda.set_device(args.gpu)
|
||||
args.dist_backend = "nccl"
|
||||
logger.info("| distributed init (rank {}): {}".format(args.rank, args.dist_url))
|
||||
torch.distributed.init_process_group(
|
||||
backend=args.dist_backend, init_method=args.dist_url, world_size=args.world_size, rank=args.rank
|
||||
)
|
||||
torch.distributed.barrier()
|
||||
setup_for_distributed(args.rank == 0)
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def accuracy(output: torch.Tensor, target: torch.Tensor, topk: Tuple[int, ...] = (1,)) -> List[torch.Tensor]:
|
||||
"""Computes the precision@k for the specified values of k"""
|
||||
if target.numel() == 0:
|
||||
return [torch.zeros([], device=output.device)]
|
||||
maxk = max(topk)
|
||||
batch_size = target.size(0)
|
||||
|
||||
_, pred = output.topk(maxk, 1, True, True)
|
||||
pred = pred.t()
|
||||
correct = pred.eq(target.view(1, -1).expand_as(pred))
|
||||
|
||||
res = []
|
||||
for k in topk:
|
||||
correct_k = correct[:k].view(-1).float().sum(0)
|
||||
res.append(correct_k.mul_(100.0 / batch_size))
|
||||
return res
|
||||
|
||||
|
||||
def interpolate(
|
||||
input: Tensor,
|
||||
size: Optional[List[int]] = None,
|
||||
scale_factor: Optional[float] = None,
|
||||
mode: str = "nearest",
|
||||
align_corners: Optional[bool] = None,
|
||||
) -> Tensor:
|
||||
"""
|
||||
Equivalent to nn.functional.interpolate, but with support for empty batch sizes.
|
||||
This will eventually be supported natively by PyTorch, and this
|
||||
class can go away.
|
||||
"""
|
||||
if float(torchvision.__version__.split(".")[1]) < 7.0:
|
||||
if input.numel() > 0:
|
||||
return torch.nn.functional.interpolate(input, size, scale_factor, mode, align_corners)
|
||||
|
||||
output_shape = _output_size(2, input, size, scale_factor)
|
||||
output_shape = list(input.shape[:-2]) + list(output_shape)
|
||||
return _new_empty_tensor(input, output_shape)
|
||||
else:
|
||||
return torchvision.ops.misc.interpolate(input, size, scale_factor, mode, align_corners)
|
||||
|
||||
|
||||
def inverse_sigmoid(x: torch.Tensor, eps: float = 1e-5) -> torch.Tensor:
|
||||
x = x.clamp(min=0, max=1)
|
||||
x1 = x.clamp(min=eps)
|
||||
x2 = (1 - x).clamp(min=eps)
|
||||
return torch.log(x1 / x2)
|
||||
|
||||
|
||||
def strip_checkpoint(checkpoint: str | os.PathLike[str]) -> None:
|
||||
state_dict = torch.load(checkpoint, map_location="cpu", weights_only=False)
|
||||
new_state_dict = {
|
||||
"model": state_dict["model"],
|
||||
"args": state_dict["args"],
|
||||
}
|
||||
# Create the temp file in the destination directory so os.replace stays on the same filesystem (atomic).
|
||||
checkpoint_dir = os.path.dirname(os.path.abspath(os.fspath(checkpoint)))
|
||||
with tempfile.NamedTemporaryFile(dir=checkpoint_dir, delete=False) as tmp_file:
|
||||
tmp_path = tmp_file.name
|
||||
try:
|
||||
torch.save(new_state_dict, tmp_path)
|
||||
# Atomic replace avoids leaving a partially written checkpoint on save failures/interruption.
|
||||
os.replace(tmp_path, checkpoint)
|
||||
finally:
|
||||
if os.path.exists(tmp_path):
|
||||
os.remove(tmp_path)
|
||||
@@ -0,0 +1,203 @@
|
||||
"""RfDetrTorch — rf-detr (Roboflow) detector via the vendored model copy (no rfdetr package).
|
||||
|
||||
Backs ``fast_layout``. Inference goes through the slimmed, detection-only model
|
||||
definition vendored under ``surya.common.rfdetr`` (validated byte-for-byte against the
|
||||
upstream rfdetr package). Pure PyTorch — runs on cpu/mps/cuda.
|
||||
|
||||
Model dir layout (downloaded from the Hub):
|
||||
rfdetr_<task>.pth the fine-tuned rf-detr weights
|
||||
config.json {"arch": "rf-detr-large", "categories": [{"id", "name"}, ...], ...}
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import glob
|
||||
import json
|
||||
import os
|
||||
from typing import List, Optional
|
||||
|
||||
from PIL import Image
|
||||
|
||||
|
||||
class _DetList(list):
|
||||
"""A list of detections that can also carry the page's encoder feature map (.features)."""
|
||||
|
||||
features = None
|
||||
|
||||
|
||||
_MPS_OP_OK: Optional[bool] = None
|
||||
|
||||
|
||||
def _mps_op_supported() -> bool:
|
||||
"""Probe (once) whether the DINOv2 pos-embed op the rf-detr backbone needs runs on
|
||||
MPS. It uses antialiased bicubic (aten::_upsample_bicubic2d_aa), which has no MPS
|
||||
kernel — it only works if PYTORCH_ENABLE_MPS_FALLBACK=1 was set before `import torch`
|
||||
(surya/__init__.py sets it, but that's too late if torch was already imported).
|
||||
Probing at runtime is the only reliable signal, since the env var reads "1" even when
|
||||
torch ignored it."""
|
||||
global _MPS_OP_OK
|
||||
if _MPS_OP_OK is None:
|
||||
import warnings
|
||||
|
||||
import torch
|
||||
|
||||
try:
|
||||
with warnings.catch_warnings():
|
||||
warnings.simplefilter("ignore")
|
||||
torch.nn.functional.interpolate(
|
||||
torch.zeros(1, 1, 4, 4, device="mps"),
|
||||
size=(6, 6),
|
||||
mode="bicubic",
|
||||
antialias=True,
|
||||
)
|
||||
_MPS_OP_OK = True
|
||||
except Exception:
|
||||
_MPS_OP_OK = False
|
||||
return _MPS_OP_OK
|
||||
|
||||
|
||||
def _pick_device(device: Optional[str]) -> str:
|
||||
import torch
|
||||
|
||||
if device:
|
||||
return device
|
||||
if torch.cuda.is_available():
|
||||
return "cuda"
|
||||
# Use MPS only if the CPU-fallback for the unsupported bicubic op is actually live
|
||||
# (see _mps_op_supported); otherwise auto-select CPU (rf-detr is ~0.3s/page on CPU)
|
||||
# so Apple Silicon never crashes regardless of import order.
|
||||
if torch.backends.mps.is_available() and _mps_op_supported():
|
||||
return "mps"
|
||||
return "cpu"
|
||||
|
||||
|
||||
class RfDetrTorch:
|
||||
def __init__(
|
||||
self,
|
||||
model_dir: str,
|
||||
num_threads: Optional[int] = None,
|
||||
device: Optional[str] = None,
|
||||
):
|
||||
import torch
|
||||
|
||||
if num_threads:
|
||||
torch.set_num_threads(int(num_threads))
|
||||
|
||||
self.device = _pick_device(device)
|
||||
if self.device == "mps" and not _mps_op_supported():
|
||||
# Only reached when the user explicitly forces mps (auto-select already
|
||||
# probes and drops to CPU — see _pick_device). The DINOv2 pos-embed op has
|
||||
# no MPS kernel and will raise unless PYTORCH_ENABLE_MPS_FALLBACK=1 was set
|
||||
# BEFORE torch was imported (surya/__init__.py sets it, but that's too late
|
||||
# if torch was already loaded).
|
||||
from surya.logging import get_logger
|
||||
|
||||
get_logger().warning(
|
||||
"FAST_DETECTOR_DEVICE=mps but the MPS bicubic fallback isn't active "
|
||||
"(torch was likely imported before surya); the rf-detr detector will "
|
||||
"crash. Export PYTORCH_ENABLE_MPS_FALLBACK=1 before importing torch, or use cpu/cuda."
|
||||
)
|
||||
|
||||
with open(os.path.join(model_dir, "config.json")) as f:
|
||||
cfg = json.load(f)
|
||||
|
||||
# rf-detr's predict() returns 0-indexed class ids that line up with the COCO
|
||||
# categories sorted by id (Row=0, Col=1; layout: Caption=0 ... Text=15).
|
||||
cats = sorted(cfg["categories"], key=lambda c: c["id"])
|
||||
self.id2label = {i: c["name"] for i, c in enumerate(cats)}
|
||||
|
||||
weights = cfg.get("weights")
|
||||
weights = os.path.join(model_dir, weights) if weights else None
|
||||
if not weights or not os.path.exists(weights):
|
||||
pths = sorted(glob.glob(os.path.join(model_dir, "*.pth")))
|
||||
if not pths:
|
||||
raise FileNotFoundError(f"no rf-detr .pth weights found in {model_dir}")
|
||||
weights = pths[0]
|
||||
|
||||
arch = (cfg.get("arch") or "rf-detr-large").lower()
|
||||
if "base" in arch:
|
||||
raise ValueError(
|
||||
"vendored rf-detr copy is rf-detr-large only; got arch=%r" % arch
|
||||
)
|
||||
from surya.common.rfdetr import RFDetrDetector
|
||||
|
||||
# Honor resolution / PE overrides from config.json so reduced-resolution
|
||||
# fine-tunes (e.g. the 448 layout model) run at their trained size; absent
|
||||
# these keys the predictor falls back to LARGE_ARGS (704), so older configs
|
||||
# are unaffected.
|
||||
arch_args = {}
|
||||
if cfg.get("resolution"):
|
||||
arch_args["resolution"] = int(cfg["resolution"])
|
||||
if cfg.get("positional_encoding_size"):
|
||||
arch_args["positional_encoding_size"] = int(cfg["positional_encoding_size"])
|
||||
self.model = RFDetrDetector(
|
||||
weights_path=weights, device=self.device, arch_args=arch_args or None
|
||||
)
|
||||
|
||||
def detect(
|
||||
self,
|
||||
images: List[Image.Image],
|
||||
threshold: float = 0.4,
|
||||
batch_size: int = 8,
|
||||
return_features: bool = False,
|
||||
) -> List[List[dict]]:
|
||||
"""Returns, per image, a list of {label, label_id, score, bbox:[x0,y0,x1,y1] pixels}.
|
||||
When return_features=True, each per-image list carries the encoder feature map on a
|
||||
``.features`` attribute ([C,F,F] tensor) for the reading-order head."""
|
||||
out: List = []
|
||||
for s in range(0, len(images), batch_size):
|
||||
chunk = [im.convert("RGB") for im in images[s : s + batch_size]]
|
||||
for det in self.model.predict(
|
||||
chunk, threshold=threshold, return_features=return_features
|
||||
):
|
||||
boxes, scores, labels = det["boxes"], det["scores"], det["labels"]
|
||||
dets: List[dict] = _DetList()
|
||||
for i in range(len(scores)):
|
||||
cid = int(labels[i])
|
||||
x0, y0, x1, y1 = (float(v) for v in boxes[i].tolist())
|
||||
dets.append(
|
||||
{
|
||||
"label": self.id2label.get(cid, str(cid)),
|
||||
"label_id": cid,
|
||||
"score": float(scores[i]),
|
||||
"bbox": [x0, y0, x1, y1],
|
||||
}
|
||||
)
|
||||
if return_features:
|
||||
dets.features = det.get("features")
|
||||
out.append(dets)
|
||||
return out
|
||||
|
||||
|
||||
def load_detector(
|
||||
model_dir: str, num_threads: Optional[int] = None, device: Optional[str] = None
|
||||
):
|
||||
"""Build the rf-detr torch detector from a model dir containing ``.pth`` weights
|
||||
+ ``config.json``. On CPU the torch rf-detr runs ~0.3s/page, so it remains the
|
||||
fast-mode path."""
|
||||
return RfDetrTorch(model_dir, num_threads=num_threads, device=device)
|
||||
|
||||
|
||||
def resolve_model_dir(checkpoint: str) -> str:
|
||||
"""Resolve a fast-model checkpoint to a local dir. Supports a plain local path, an
|
||||
``hf://<repo>/<subfolder>`` ref (downloaded from the Hub), or an ``s3://`` path."""
|
||||
if checkpoint and checkpoint.startswith("hf://"):
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
parts = checkpoint[len("hf://") :].split("/")
|
||||
repo_id = "/".join(parts[:2])
|
||||
subfolder = "/".join(parts[2:])
|
||||
local = snapshot_download(
|
||||
repo_id,
|
||||
allow_patterns=[f"{subfolder}/*"] if subfolder else None,
|
||||
)
|
||||
return os.path.join(local, subfolder) if subfolder else local
|
||||
if checkpoint and os.path.isdir(checkpoint):
|
||||
return checkpoint
|
||||
if checkpoint and checkpoint.startswith("s3://"):
|
||||
from surya.common.s3 import download_directory # type: ignore
|
||||
|
||||
return download_directory(checkpoint)
|
||||
raise FileNotFoundError(
|
||||
f"fast-model checkpoint not found as a local dir: {checkpoint!r}"
|
||||
)
|
||||
@@ -0,0 +1,182 @@
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import tempfile
|
||||
import time
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from pathlib import Path
|
||||
|
||||
import requests
|
||||
from tqdm import tqdm
|
||||
|
||||
from surya.logging import get_logger
|
||||
from surya.settings import settings
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
# Lock file expiration time in seconds (10 minutes)
|
||||
LOCK_EXPIRATION = 600
|
||||
|
||||
|
||||
def join_urls(url1: str, url2: str):
|
||||
url1 = url1.rstrip("/")
|
||||
url2 = url2.lstrip("/")
|
||||
return f"{url1}/{url2}"
|
||||
|
||||
|
||||
def get_model_name(pretrained_model_name_or_path: str):
|
||||
return pretrained_model_name_or_path.split("/")[0]
|
||||
|
||||
|
||||
def download_file(remote_path: str, local_path: str, chunk_size: int = 1024 * 1024):
|
||||
local_path = Path(local_path)
|
||||
try:
|
||||
response = requests.get(remote_path, stream=True, allow_redirects=True)
|
||||
response.raise_for_status() # Raise an exception for bad status codes
|
||||
|
||||
# Get file size from headers for progress bar
|
||||
total_size = int(response.headers.get('content-length', 0))
|
||||
|
||||
# Create progress bar with file name and size info
|
||||
filename = local_path.name
|
||||
pbar = tqdm(
|
||||
total=total_size,
|
||||
unit='B',
|
||||
unit_scale=True,
|
||||
unit_divisor=1024,
|
||||
desc=f"Downloading {filename}",
|
||||
miniters=1
|
||||
)
|
||||
|
||||
with open(local_path, "wb") as f:
|
||||
downloaded = 0
|
||||
for chunk in response.iter_content(chunk_size=chunk_size):
|
||||
if chunk:
|
||||
f.write(chunk)
|
||||
downloaded += len(chunk)
|
||||
pbar.update(len(chunk))
|
||||
|
||||
pbar.close()
|
||||
return local_path
|
||||
except Exception as e:
|
||||
if local_path.exists():
|
||||
local_path.unlink()
|
||||
logger.error(f"Download error for file {remote_path}: {str(e)}")
|
||||
raise
|
||||
|
||||
|
||||
def check_manifest(local_dir: str):
|
||||
local_dir = Path(local_dir)
|
||||
manifest_path = local_dir / "manifest.json"
|
||||
if not os.path.exists(manifest_path):
|
||||
return False
|
||||
|
||||
try:
|
||||
with open(manifest_path, "r") as f:
|
||||
manifest = json.load(f)
|
||||
for file in manifest["files"]:
|
||||
if not os.path.exists(local_dir / file):
|
||||
return False
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
def download_directory(remote_path: str, local_dir: str):
|
||||
model_name = get_model_name(remote_path)
|
||||
s3_url = join_urls(settings.S3_BASE_URL, remote_path)
|
||||
# Check to see if it's already downloaded
|
||||
model_exists = check_manifest(local_dir)
|
||||
if model_exists:
|
||||
return
|
||||
|
||||
# Use tempfile.TemporaryDirectory to automatically clean up
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
# Download the manifest file
|
||||
manifest_file = join_urls(s3_url, "manifest.json")
|
||||
manifest_path = os.path.join(temp_dir, "manifest.json")
|
||||
download_file(manifest_file, manifest_path)
|
||||
|
||||
# List and download all files
|
||||
with open(manifest_path, "r") as f:
|
||||
manifest = json.load(f)
|
||||
|
||||
pbar = tqdm(
|
||||
desc=f"Downloading {model_name} model to {local_dir}",
|
||||
total=len(manifest["files"]),
|
||||
)
|
||||
|
||||
with ThreadPoolExecutor(
|
||||
max_workers=settings.PARALLEL_DOWNLOAD_WORKERS
|
||||
) as executor:
|
||||
futures = []
|
||||
for file in manifest["files"]:
|
||||
remote_file = join_urls(s3_url, file)
|
||||
local_file = os.path.join(temp_dir, file)
|
||||
futures.append(executor.submit(download_file, remote_file, local_file))
|
||||
|
||||
for future in futures:
|
||||
future.result()
|
||||
pbar.update(1)
|
||||
|
||||
pbar.close()
|
||||
|
||||
# Move all files to new directory
|
||||
for file in os.listdir(temp_dir):
|
||||
shutil.move(os.path.join(temp_dir, file), local_dir)
|
||||
|
||||
|
||||
class S3DownloaderMixin:
|
||||
s3_prefix = "s3://"
|
||||
|
||||
@classmethod
|
||||
def get_local_path(cls, pretrained_model_name_or_path) -> str:
|
||||
if pretrained_model_name_or_path.startswith(cls.s3_prefix):
|
||||
pretrained_model_name_or_path = pretrained_model_name_or_path.replace(
|
||||
cls.s3_prefix, ""
|
||||
)
|
||||
cache_dir = settings.MODEL_CACHE_DIR
|
||||
local_path = os.path.join(cache_dir, pretrained_model_name_or_path)
|
||||
os.makedirs(local_path, exist_ok=True)
|
||||
else:
|
||||
local_path = ""
|
||||
return local_path
|
||||
|
||||
@classmethod
|
||||
def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
|
||||
# Allow loading models directly from the hub, or using s3
|
||||
if not pretrained_model_name_or_path.startswith(cls.s3_prefix):
|
||||
return super().from_pretrained(
|
||||
pretrained_model_name_or_path, *args, **kwargs
|
||||
)
|
||||
|
||||
local_path = cls.get_local_path(pretrained_model_name_or_path)
|
||||
pretrained_model_name_or_path = pretrained_model_name_or_path.replace(
|
||||
cls.s3_prefix, ""
|
||||
)
|
||||
|
||||
# Retry logic for downloading the model folder
|
||||
retries = 3
|
||||
delay = 5
|
||||
attempt = 0
|
||||
success = False
|
||||
while not success and attempt < retries:
|
||||
try:
|
||||
download_directory(pretrained_model_name_or_path, local_path)
|
||||
success = True # If download succeeded
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Error downloading model from {pretrained_model_name_or_path}. Attempt {attempt + 1} of {retries}. Error: {e}"
|
||||
)
|
||||
attempt += 1
|
||||
if attempt < retries:
|
||||
logger.info(f"Retrying in {delay} seconds...")
|
||||
time.sleep(delay) # Wait before retrying
|
||||
else:
|
||||
logger.error(
|
||||
f"Failed to download {pretrained_model_name_or_path} after {retries} attempts."
|
||||
)
|
||||
raise e # Reraise exception after max retries
|
||||
|
||||
return super().from_pretrained(local_path, *args, **kwargs)
|
||||
@@ -0,0 +1,44 @@
|
||||
from typing import List
|
||||
|
||||
from surya.common.polygon import PolygonBox
|
||||
|
||||
|
||||
def clean_boxes(boxes: List[PolygonBox]) -> List[PolygonBox]:
|
||||
new_boxes = []
|
||||
for box_obj in boxes:
|
||||
xs = [point[0] for point in box_obj.polygon]
|
||||
ys = [point[1] for point in box_obj.polygon]
|
||||
if max(xs) == min(xs) or max(ys) == min(ys):
|
||||
continue
|
||||
|
||||
box = box_obj.bbox
|
||||
contained = False
|
||||
for other_box_obj in boxes:
|
||||
if other_box_obj.polygon == box_obj.polygon:
|
||||
continue
|
||||
|
||||
other_box = other_box_obj.bbox
|
||||
if box == other_box:
|
||||
continue
|
||||
if (
|
||||
box[0] >= other_box[0]
|
||||
and box[1] >= other_box[1]
|
||||
and box[2] <= other_box[2]
|
||||
and box[3] <= other_box[3]
|
||||
):
|
||||
contained = True
|
||||
break
|
||||
if not contained:
|
||||
new_boxes.append(box_obj)
|
||||
return new_boxes
|
||||
|
||||
|
||||
def expand_bbox(bbox, expansion_factor=0.01):
|
||||
expansion_low = 1 - expansion_factor
|
||||
expansion_high = 1 + expansion_factor
|
||||
return [
|
||||
bbox[0] * expansion_low,
|
||||
bbox[1] * expansion_low,
|
||||
bbox[2] * expansion_high,
|
||||
bbox[3] * expansion_high,
|
||||
]
|
||||
@@ -0,0 +1,66 @@
|
||||
from PIL import ImageDraw, ImageFont
|
||||
|
||||
from surya.debug.fonts import get_font_path
|
||||
from surya.debug.text import get_text_size
|
||||
|
||||
|
||||
def draw_bboxes_on_image(
|
||||
bboxes, image, labels=None, label_font_size=10, color: str | list = "red"
|
||||
):
|
||||
polys = []
|
||||
for bb in bboxes:
|
||||
# Clockwise polygon
|
||||
poly = [[bb[0], bb[1]], [bb[2], bb[1]], [bb[2], bb[3]], [bb[0], bb[3]]]
|
||||
polys.append(poly)
|
||||
|
||||
return draw_polys_on_image(
|
||||
polys, image, labels, label_font_size=label_font_size, color=color
|
||||
)
|
||||
|
||||
|
||||
def draw_polys_on_image(
|
||||
corners,
|
||||
image,
|
||||
labels=None,
|
||||
box_padding=-1,
|
||||
label_offset=1,
|
||||
label_font_size=10,
|
||||
color: str | list = "red",
|
||||
):
|
||||
draw = ImageDraw.Draw(image)
|
||||
font_path = get_font_path()
|
||||
label_font = ImageFont.truetype(font_path, label_font_size)
|
||||
|
||||
for i in range(len(corners)):
|
||||
poly = corners[i]
|
||||
poly = [(int(p[0]), int(p[1])) for p in poly]
|
||||
draw.polygon(
|
||||
poly, outline=color[i] if isinstance(color, list) else color, width=1
|
||||
)
|
||||
|
||||
if labels is not None:
|
||||
label = labels[i]
|
||||
text_position = (
|
||||
min([p[0] for p in poly]) + label_offset,
|
||||
min([p[1] for p in poly]) + label_offset,
|
||||
)
|
||||
text_size = get_text_size(label, label_font)
|
||||
box_position = (
|
||||
text_position[0] - box_padding + label_offset,
|
||||
text_position[1] - box_padding + label_offset,
|
||||
text_position[0] + text_size[0] + box_padding + label_offset,
|
||||
text_position[1] + text_size[1] + box_padding + label_offset,
|
||||
)
|
||||
try:
|
||||
draw.rectangle(box_position, fill="white")
|
||||
except Exception as e:
|
||||
print(f"Error drawing rectangle at {box_position}: {e}")
|
||||
continue
|
||||
draw.text(
|
||||
text_position,
|
||||
label,
|
||||
fill=color[i] if isinstance(color, list) else color,
|
||||
font=label_font,
|
||||
)
|
||||
|
||||
return image
|
||||
@@ -0,0 +1,24 @@
|
||||
from typing import List, Optional
|
||||
import os
|
||||
import requests
|
||||
|
||||
from surya.settings import settings
|
||||
|
||||
|
||||
def get_font_path(langs: Optional[List[str]] = None) -> str:
|
||||
font_path = settings.RECOGNITION_RENDER_FONTS["all"]
|
||||
if langs is not None:
|
||||
for k in settings.RECOGNITION_RENDER_FONTS:
|
||||
if k in langs and len(langs) == 1:
|
||||
font_path = settings.RECOGNITION_RENDER_FONTS[k]
|
||||
break
|
||||
|
||||
if not os.path.exists(font_path):
|
||||
os.makedirs(os.path.dirname(font_path), exist_ok=True)
|
||||
font_dl_path = f"{settings.RECOGNITION_FONT_DL_BASE}/{os.path.basename(font_path)}"
|
||||
with requests.get(font_dl_path, stream=True) as r, open(font_path, 'wb') as f:
|
||||
r.raise_for_status()
|
||||
for chunk in r.iter_content(chunk_size=8192):
|
||||
f.write(chunk)
|
||||
|
||||
return font_path
|
||||
@@ -0,0 +1,64 @@
|
||||
<style>
|
||||
.katex-display-container {
|
||||
display: inline-block;
|
||||
max-width: 100%;
|
||||
overflow-x: auto;
|
||||
max-height: 100%;
|
||||
}
|
||||
|
||||
.katex-inline-container {
|
||||
display: inline-block;
|
||||
max-width: 100%;
|
||||
overflow-x: auto;
|
||||
max-height: 100%;
|
||||
}
|
||||
</style>
|
||||
<script src="https://cdn.jsdelivr.net/npm/katex@0.16.21/dist/katex.min.js" onload="setTimeout(function() {renderMath()})" async></script>
|
||||
<link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.21/dist/katex.min.css">
|
||||
<script>
|
||||
function htmlUnescape(escapedText) {
|
||||
const htmlEntities = {
|
||||
'&': '&',
|
||||
'<': '<',
|
||||
'>': '>',
|
||||
'"': '"',
|
||||
''': "'",
|
||||
' ': ' '
|
||||
};
|
||||
|
||||
return escapedText.replace(/&|<|>|"|'| /g, match => htmlEntities[match]);
|
||||
}
|
||||
|
||||
const renderMath = (function() {
|
||||
try {
|
||||
const mathElements = document.querySelectorAll('math');
|
||||
|
||||
mathElements.forEach(function(element) {
|
||||
let mathContent = element.innerHTML.trim();
|
||||
mathContent = htmlUnescape(mathContent);
|
||||
const isDisplay = element.getAttribute('display') === 'block';
|
||||
|
||||
const container = document.createElement('span');
|
||||
container.className = isDisplay ? 'katex-display-container' : 'katex-inline-container';
|
||||
element.parentNode.insertBefore(container, element);
|
||||
|
||||
try {
|
||||
katex.render(mathContent, container, {
|
||||
displayMode: isDisplay,
|
||||
throwOnError: false
|
||||
});
|
||||
|
||||
} catch (err) {
|
||||
console.error('KaTeX rendering error:', err);
|
||||
container.textContent = mathContent; // Fallback to raw text
|
||||
}
|
||||
|
||||
element.parentNode.removeChild(element);
|
||||
});
|
||||
|
||||
console.log('Math rendering complete with', mathElements.length, 'expressions');
|
||||
} catch (err) {
|
||||
console.error('Error in renderMath function:', err);
|
||||
}
|
||||
});
|
||||
</script>
|
||||
@@ -0,0 +1,90 @@
|
||||
import html as htmllib
|
||||
import os.path
|
||||
import re
|
||||
|
||||
filepath = os.path.abspath(__file__)
|
||||
|
||||
def render_text_as_html(
|
||||
bboxes: list[list[int]],
|
||||
texts: list[str],
|
||||
image_size: tuple[int, int],
|
||||
base_font_size: int = 16,
|
||||
scaler: int = 2
|
||||
):
|
||||
katex_path = os.path.join(os.path.dirname(filepath), "katex.js")
|
||||
with open(katex_path, "r") as f:
|
||||
katex_script = f.read()
|
||||
|
||||
html_content = []
|
||||
image_size = tuple([int(s * scaler) for s in image_size])
|
||||
width, height = image_size
|
||||
|
||||
|
||||
html_content.append(f"""
|
||||
<!DOCTYPE html>
|
||||
<html>
|
||||
<head>
|
||||
<style>
|
||||
body {{
|
||||
margin: 0;
|
||||
padding: 0;
|
||||
width: {width}px;
|
||||
height: {height}px;
|
||||
position: relative;
|
||||
overflow: hidden;
|
||||
background: white;
|
||||
color: black;
|
||||
}}
|
||||
.text-box {{
|
||||
position: absolute;
|
||||
overflow: hidden;
|
||||
display: flex;
|
||||
justify-content: left;
|
||||
font-family: Arial, sans-serif;
|
||||
white-space: pre-wrap;
|
||||
}}
|
||||
.vertical-text {{
|
||||
writing-mode: vertical-rl; /* Top to bottom, right to left */
|
||||
}}
|
||||
</style>
|
||||
{katex_script}
|
||||
</head>
|
||||
<body>
|
||||
""")
|
||||
|
||||
for i, (bbox, text) in enumerate(zip(bboxes, texts)):
|
||||
bbox = bbox.copy()
|
||||
bbox = [int(bb * scaler) for bb in bbox]
|
||||
x1, y1, x2, y2 = bbox
|
||||
width = x2 - x1
|
||||
height = y2 - y1
|
||||
min_dim = min(width, height)
|
||||
|
||||
# Scale font size based on box height
|
||||
font_size = min(int(min_dim * 0.75), base_font_size)
|
||||
|
||||
# Create div with absolute positioning
|
||||
div_style = (
|
||||
f"left: {x1}px; "
|
||||
f"top: {y1}px; "
|
||||
f"width: {width}px; "
|
||||
f"height: {height}px; "
|
||||
f"font-size: {font_size}px;"
|
||||
)
|
||||
|
||||
class_ = "text-box"
|
||||
if height > width * 2:
|
||||
class_ += " vertical-text"
|
||||
|
||||
# Determine if content is HTML/MathML or plain text
|
||||
if "<" in text and ">" in text and re.search(r"<(html|math|div|sub|sup|i|u|mark|small|del|b|br|code)\b", text.lower()):
|
||||
# Content is already HTML/MathML, include as-is
|
||||
html_content.append(f'<span class="{class_}" id="box-{i}" style="{div_style}">{text}</span>')
|
||||
else:
|
||||
# Plain text, escape it
|
||||
escaped_text = htmllib.escape(text)
|
||||
html_content.append(f'<span class="{class_}" id="box-{i}" style="{div_style}">{escaped_text}</span>')
|
||||
|
||||
html_content.append("</body></html>")
|
||||
|
||||
return "\n".join(html_content), image_size
|
||||
@@ -0,0 +1,8 @@
|
||||
from PIL import Image, ImageDraw
|
||||
|
||||
|
||||
def get_text_size(text, font):
|
||||
im = Image.new(mode="P", size=(0, 0))
|
||||
draw = ImageDraw.Draw(im)
|
||||
_, _, width, height = draw.textbbox((0, 0), text=text, font=font)
|
||||
return width, height
|
||||
@@ -0,0 +1,147 @@
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from typing import List, Generator, Tuple
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from PIL import Image
|
||||
from tqdm import tqdm
|
||||
|
||||
from surya.common.predictor import BasePredictor
|
||||
|
||||
from surya.detection.loader import DetectionModelLoader
|
||||
from surya.detection.parallel import FakeExecutor
|
||||
from surya.detection.util import get_total_splits, split_image
|
||||
from surya.detection.schema import TextDetectionResult
|
||||
from surya.settings import settings
|
||||
from surya.detection.heatmap import parallel_get_boxes
|
||||
|
||||
|
||||
class DetectionPredictor(BasePredictor):
|
||||
model_loader_cls = DetectionModelLoader
|
||||
batch_size = settings.DETECTOR_BATCH_SIZE
|
||||
default_batch_sizes = {"cpu": 8, "mps": 8, "cuda": 36}
|
||||
|
||||
def __call__(
|
||||
self, images: List[Image.Image], batch_size=None, include_maps=False
|
||||
) -> List[TextDetectionResult]:
|
||||
detection_generator = self.batch_detection(images, batch_size=batch_size)
|
||||
|
||||
postprocessing_futures = []
|
||||
max_workers = min(settings.DETECTOR_POSTPROCESSING_CPU_WORKERS, len(images))
|
||||
parallelize = (
|
||||
not settings.IN_STREAMLIT
|
||||
and len(images) >= settings.DETECTOR_MIN_PARALLEL_THRESH
|
||||
)
|
||||
executor = ThreadPoolExecutor if parallelize else FakeExecutor
|
||||
with executor(max_workers=max_workers) as e:
|
||||
for preds, orig_sizes in detection_generator:
|
||||
for pred, orig_size in zip(preds, orig_sizes):
|
||||
postprocessing_futures.append(
|
||||
e.submit(parallel_get_boxes, pred, orig_size, include_maps)
|
||||
)
|
||||
|
||||
return [future.result() for future in postprocessing_futures]
|
||||
|
||||
def prepare_image(self, img):
|
||||
new_size = (self.processor.size["width"], self.processor.size["height"])
|
||||
|
||||
# This double resize actually necessary for downstream accuracy
|
||||
img.thumbnail(new_size, Image.Resampling.LANCZOS)
|
||||
img = img.resize(
|
||||
new_size, Image.Resampling.LANCZOS
|
||||
) # Stretch smaller dimension to fit new size
|
||||
|
||||
img = np.asarray(img, dtype=np.uint8)
|
||||
img = self.processor(img)["pixel_values"][0]
|
||||
img = torch.from_numpy(img)
|
||||
return img
|
||||
|
||||
def batch_detection(
|
||||
self, images: List, batch_size=None
|
||||
) -> Generator[Tuple[List[List[np.ndarray]], List[Tuple[int, int]]], None, None]:
|
||||
assert all([isinstance(image, Image.Image) for image in images])
|
||||
if batch_size is None:
|
||||
batch_size = self.get_batch_size()
|
||||
heatmap_count = self.model.config.num_labels
|
||||
|
||||
orig_sizes = [image.size for image in images]
|
||||
splits_per_image = [
|
||||
get_total_splits(size, self.processor.size["height"]) for size in orig_sizes
|
||||
]
|
||||
|
||||
batches = []
|
||||
current_batch_size = 0
|
||||
current_batch = []
|
||||
for i in range(len(images)):
|
||||
if current_batch_size + splits_per_image[i] > batch_size:
|
||||
if len(current_batch) > 0:
|
||||
batches.append(current_batch)
|
||||
current_batch = []
|
||||
current_batch_size = 0
|
||||
current_batch.append(i)
|
||||
current_batch_size += splits_per_image[i]
|
||||
|
||||
if len(current_batch) > 0:
|
||||
batches.append(current_batch)
|
||||
|
||||
for batch_idx in tqdm(
|
||||
range(len(batches)), desc="Detecting bboxes", disable=self.disable_tqdm
|
||||
):
|
||||
batch_image_idxs = batches[batch_idx]
|
||||
batch_images = [images[j].convert("RGB") for j in batch_image_idxs]
|
||||
|
||||
split_index = []
|
||||
split_heights = []
|
||||
image_splits = []
|
||||
for image_idx, image in enumerate(batch_images):
|
||||
image_parts, split_height = split_image(
|
||||
image, self.processor.size["height"]
|
||||
)
|
||||
image_splits.extend(image_parts)
|
||||
split_index.extend([image_idx] * len(image_parts))
|
||||
split_heights.extend(split_height)
|
||||
|
||||
image_splits = [self.prepare_image(image) for image in image_splits]
|
||||
# Batch images in dim 0
|
||||
batch = torch.stack(image_splits, dim=0).to(self.model.dtype)
|
||||
|
||||
with settings.INFERENCE_MODE():
|
||||
pred = self.model(pixel_values=batch.to(self.model.device))
|
||||
|
||||
logits = pred.logits
|
||||
correct_shape = [
|
||||
self.processor.size["height"],
|
||||
self.processor.size["width"],
|
||||
]
|
||||
current_shape = list(logits.shape[2:])
|
||||
if current_shape != correct_shape:
|
||||
logits = F.interpolate(
|
||||
logits, size=correct_shape, mode="bilinear", align_corners=False
|
||||
)
|
||||
|
||||
logits = logits.to(torch.float32).cpu().numpy()
|
||||
preds = []
|
||||
for i, (idx, height) in enumerate(zip(split_index, split_heights)):
|
||||
# If our current prediction length is below the image idx, that means we have a new image
|
||||
# Otherwise, we need to add to the current image
|
||||
if len(preds) <= idx:
|
||||
preds.append([logits[i][k] for k in range(heatmap_count)])
|
||||
else:
|
||||
heatmaps = preds[idx]
|
||||
pred_heatmaps = [logits[i][k] for k in range(heatmap_count)]
|
||||
|
||||
if height < self.processor.size["height"]:
|
||||
# Cut off padding to get original height
|
||||
pred_heatmaps = [
|
||||
pred_heatmap[:height, :] for pred_heatmap in pred_heatmaps
|
||||
]
|
||||
|
||||
for k in range(heatmap_count):
|
||||
heatmaps[k] = np.vstack([heatmaps[k], pred_heatmaps[k]])
|
||||
preds[idx] = heatmaps
|
||||
|
||||
yield preds, [orig_sizes[j] for j in batch_image_idxs]
|
||||
|
||||
torch.cuda.empty_cache()
|
||||
@@ -0,0 +1,165 @@
|
||||
from typing import List
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
from surya.common.util import clean_boxes
|
||||
from surya.detection import TextDetectionResult
|
||||
from surya.common.polygon import PolygonBox
|
||||
from surya.settings import settings
|
||||
|
||||
|
||||
def get_dynamic_thresholds(linemap, text_threshold, low_text, typical_top10_avg=0.7):
|
||||
# Find average intensity of top 10% pixels
|
||||
flat_map = linemap.ravel()
|
||||
top_10_count = int(len(flat_map) * 0.9)
|
||||
avg_intensity = np.mean(np.partition(flat_map, top_10_count)[top_10_count:])
|
||||
scaling_factor = np.clip(avg_intensity / typical_top10_avg, 0, 1) ** (1 / 2)
|
||||
|
||||
low_text = np.clip(low_text * scaling_factor, 0.1, 0.6)
|
||||
text_threshold = np.clip(text_threshold * scaling_factor, 0.15, 0.8)
|
||||
|
||||
return text_threshold, low_text
|
||||
|
||||
|
||||
def detect_boxes(linemap, text_threshold, low_text):
|
||||
# From CRAFT - https://github.com/clovaai/CRAFT-pytorch
|
||||
# Modified to return boxes and for speed, accuracy
|
||||
img_h, img_w = linemap.shape
|
||||
|
||||
text_threshold, low_text = get_dynamic_thresholds(linemap, text_threshold, low_text)
|
||||
|
||||
text_score_comb = (linemap > low_text).astype(np.uint8)
|
||||
label_count, labels, stats, centroids = cv2.connectedComponentsWithStats(
|
||||
text_score_comb, connectivity=4
|
||||
)
|
||||
|
||||
det = []
|
||||
confidences = []
|
||||
max_confidence = 0
|
||||
|
||||
for k in range(1, label_count):
|
||||
# size filtering
|
||||
size = stats[k, cv2.CC_STAT_AREA]
|
||||
if size < 10:
|
||||
continue
|
||||
|
||||
# make segmentation map
|
||||
x, y, w, h = stats[
|
||||
k,
|
||||
[cv2.CC_STAT_LEFT, cv2.CC_STAT_TOP, cv2.CC_STAT_WIDTH, cv2.CC_STAT_HEIGHT],
|
||||
]
|
||||
|
||||
try:
|
||||
niter = int(np.sqrt(min(w, h)))
|
||||
except ValueError:
|
||||
niter = 0
|
||||
|
||||
buffer = 1
|
||||
sx, sy = max(0, x - niter - buffer), max(0, y - niter - buffer)
|
||||
ex, ey = min(img_w, x + w + niter + buffer), min(img_h, y + h + niter + buffer)
|
||||
|
||||
mask = labels[sy:ey, sx:ex] == k
|
||||
selected_linemap = linemap[sy:ey, sx:ex][mask]
|
||||
if selected_linemap.size == 0:
|
||||
continue
|
||||
|
||||
line_max = np.max(selected_linemap)
|
||||
|
||||
# thresholding
|
||||
if line_max < text_threshold:
|
||||
continue
|
||||
|
||||
segmap = mask.astype(np.uint8)
|
||||
|
||||
ksize = buffer + niter
|
||||
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (ksize, ksize))
|
||||
selected_segmap = cv2.dilate(segmap, kernel)
|
||||
|
||||
# make box
|
||||
y_inds, x_inds = np.nonzero(selected_segmap)
|
||||
x_inds += sx
|
||||
y_inds += sy
|
||||
np_contours = np.column_stack((x_inds, y_inds))
|
||||
rectangle = cv2.minAreaRect(np_contours)
|
||||
box = cv2.boxPoints(rectangle)
|
||||
|
||||
# align diamond-shape
|
||||
w, h = np.linalg.norm(box[0] - box[1]), np.linalg.norm(box[1] - box[2])
|
||||
box_ratio = max(w, h) / (min(w, h) + 1e-5)
|
||||
if abs(1 - box_ratio) <= 0.1:
|
||||
left, right = np_contours[:, 0].min(), np_contours[:, 0].max()
|
||||
top, bottom = np_contours[:, 1].min(), np_contours[:, 1].max()
|
||||
box = np.array(
|
||||
[[left, top], [right, top], [right, bottom], [left, bottom]],
|
||||
dtype=np.float32,
|
||||
)
|
||||
|
||||
# make clock-wise order
|
||||
startidx = box.sum(axis=1).argmin()
|
||||
box = np.roll(box, 4 - startidx, 0)
|
||||
|
||||
max_confidence = max(max_confidence, line_max)
|
||||
|
||||
confidences.append(line_max)
|
||||
det.append(box)
|
||||
|
||||
if max_confidence > 0:
|
||||
confidences = [c / max_confidence for c in confidences]
|
||||
return det, confidences
|
||||
|
||||
|
||||
def get_detected_boxes(textmap, text_threshold=None, low_text=None) -> List[PolygonBox]:
|
||||
if text_threshold is None:
|
||||
text_threshold = settings.DETECTOR_TEXT_THRESHOLD
|
||||
if low_text is None:
|
||||
low_text = settings.DETECTOR_BLANK_THRESHOLD
|
||||
|
||||
if textmap.dtype != np.float32:
|
||||
textmap = textmap.astype(np.float32)
|
||||
|
||||
boxes, confidences = detect_boxes(textmap, text_threshold, low_text)
|
||||
# From point form to box form
|
||||
return [
|
||||
PolygonBox(polygon=box, confidence=confidence)
|
||||
for box, confidence in zip(boxes, confidences)
|
||||
]
|
||||
|
||||
|
||||
def get_and_clean_boxes(
|
||||
textmap, processor_size, image_size, text_threshold=None, low_text=None
|
||||
) -> List[PolygonBox]:
|
||||
bboxes = get_detected_boxes(textmap, text_threshold, low_text)
|
||||
for bbox in bboxes:
|
||||
bbox.rescale(processor_size, image_size)
|
||||
bbox.fit_to_bounds([0, 0, image_size[0], image_size[1]])
|
||||
|
||||
bboxes = clean_boxes(bboxes)
|
||||
return bboxes
|
||||
|
||||
|
||||
def parallel_get_boxes(preds, orig_sizes, include_maps=False):
|
||||
heatmap, affinity_map = preds
|
||||
heat_img, aff_img = None, None
|
||||
|
||||
if include_maps:
|
||||
heat_img = Image.fromarray((heatmap * 255).astype(np.uint8))
|
||||
aff_img = Image.fromarray((affinity_map * 255).astype(np.uint8))
|
||||
heatmap_size = list(reversed(heatmap.shape))
|
||||
bboxes = get_and_clean_boxes(heatmap, heatmap_size, orig_sizes)
|
||||
for box in bboxes:
|
||||
# Skip for vertical boxes
|
||||
if box.height < 3 * box.width:
|
||||
box.expand(x_margin=0, y_margin=settings.DETECTOR_BOX_Y_EXPAND_MARGIN)
|
||||
box.fit_to_bounds(
|
||||
[0, 0, orig_sizes[0], orig_sizes[1]]
|
||||
) # Fix any bad expands
|
||||
|
||||
result = TextDetectionResult(
|
||||
bboxes=bboxes,
|
||||
heatmap=heat_img,
|
||||
affinity_map=aff_img,
|
||||
image_bbox=[0, 0, orig_sizes[0], orig_sizes[1]],
|
||||
)
|
||||
return result
|
||||
@@ -0,0 +1,53 @@
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
|
||||
from surya.common.load import ModelLoader
|
||||
from surya.detection.processor import SegformerImageProcessor
|
||||
|
||||
from surya.detection.model.config import EfficientViTConfig
|
||||
from surya.detection.model.encoderdecoder import EfficientViTForSemanticSegmentation
|
||||
from surya.logging import get_logger
|
||||
from surya.settings import settings
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
|
||||
class DetectionModelLoader(ModelLoader):
|
||||
def __init__(self, checkpoint: Optional[str] = None):
|
||||
super().__init__(checkpoint)
|
||||
|
||||
if self.checkpoint is None:
|
||||
self.checkpoint = settings.DETECTOR_MODEL_CHECKPOINT
|
||||
|
||||
def model(
|
||||
self,
|
||||
device: Optional[torch.device | str] = None,
|
||||
dtype: Optional[torch.dtype | str] = None,
|
||||
attention_implementation: Optional[str] = None,
|
||||
) -> EfficientViTForSemanticSegmentation:
|
||||
if device is None:
|
||||
device = settings.TORCH_DEVICE_MODEL
|
||||
if dtype is None:
|
||||
dtype = settings.MODEL_DTYPE
|
||||
|
||||
config = EfficientViTConfig.from_pretrained(self.checkpoint)
|
||||
model = EfficientViTForSemanticSegmentation.from_pretrained(
|
||||
self.checkpoint,
|
||||
dtype=dtype,
|
||||
config=config,
|
||||
)
|
||||
model = model.to(device)
|
||||
model = model.eval()
|
||||
|
||||
logger.debug(
|
||||
f"Loaded detection model {self.checkpoint} from {EfficientViTForSemanticSegmentation.get_local_path(self.checkpoint)} onto device {device} with dtype {dtype}"
|
||||
)
|
||||
return model
|
||||
|
||||
def processor(
|
||||
self,
|
||||
device: Optional[torch.device | str] = None,
|
||||
dtype: Optional[torch.dtype | str] = None,
|
||||
) -> SegformerImageProcessor:
|
||||
return SegformerImageProcessor.from_pretrained(self.checkpoint)
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user