Spaces:
Running
Running
Add model files
Browse files- README.md +162 -12
- __init__.py +87 -0
- __pycache__/__init__.cpython-310.pyc +0 -0
- __pycache__/layout.cpython-310.pyc +0 -0
- __pycache__/logging_config.cpython-310.pyc +0 -0
- __pycache__/ocr_backends.cpython-310.pyc +0 -0
- __pycache__/pipeline.cpython-310.pyc +0 -0
- __pycache__/table_extraction.cpython-310.pyc +0 -0
- layout.py +672 -0
- logging_config.py +44 -0
- ocr_backends.py +112 -0
- pipeline.py +505 -0
- table_extraction.py +565 -0
README.md
CHANGED
|
@@ -1,12 +1,162 @@
|
|
| 1 |
-
|
| 2 |
-
|
| 3 |
-
|
| 4 |
-
|
| 5 |
-
|
| 6 |
-
|
| 7 |
-
|
| 8 |
-
|
| 9 |
-
|
| 10 |
-
-
|
| 11 |
-
|
| 12 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# pdf_pipeline
|
| 2 |
+
|
| 3 |
+
A PDF β Markdown extraction pipeline implementing this flow:
|
| 4 |
+
|
| 5 |
+
```
|
| 6 |
+
pdf
|
| 7 |
+
ββ pdfplumber βββββββββββββββ¬βββββββββββββββββββββββββββββββ
|
| 8 |
+
β words+bboxes(xyxy) β image β
|
| 9 |
+
βΌ βΌ β
|
| 10 |
+
empty words? pp-doclayout β
|
| 11 |
+
yes β no (DocLayoutV3) β
|
| 12 |
+
β ββββββββββββββ β layout_class+bbox+order β
|
| 13 |
+
βΌ βΌ βΌ
|
| 14 |
+
OCR(pytesseract/ Plumber words + pp-doclayout matching
|
| 15 |
+
paddleocr/rapidocr) (align_words_to_layout)
|
| 16 |
+
β β
|
| 17 |
+
βΌ βΌ
|
| 18 |
+
Matching ocr bboxes (class_name, bbox, reading_order, text)
|
| 19 |
+
and layout β
|
| 20 |
+
(align_ocr_to_layout) βΌ
|
| 21 |
+
β class_name == "table"?
|
| 22 |
+
β yes β no
|
| 23 |
+
βββββββββββββ¬ββββββββββββ ββ formatted by class
|
| 24 |
+
βΌ
|
| 25 |
+
TableFormerONNX β OTSL β markdown
|
| 26 |
+
```
|
| 27 |
+
|
| 28 |
+
Both matching branches converge on the same `RegionText` shape and the same
|
| 29 |
+
downstream per-region rendering (`region_to_markdown`), including the
|
| 30 |
+
table β TableFormerONNX β OTSL β markdown branch β the diagram only draws
|
| 31 |
+
that branch coming off the pdfplumber-matching box, but a "table" layout
|
| 32 |
+
region can equally appear on a scanned (OCR'd) page, so it's handled the
|
| 33 |
+
same way in both paths.
|
| 34 |
+
|
| 35 |
+
## Modules
|
| 36 |
+
|
| 37 |
+
| File | Diagram piece(s) |
|
| 38 |
+
|------------------------|------------------|
|
| 39 |
+
| `layout.py` | `pp-doclayout` (DocLayoutV3 ONNX wrapper) + both "matching" boxes (`align_words_to_layout`, `align_ocr_to_layout`, sharing the generic `align_tokens_to_layout` core) |
|
| 40 |
+
| `ocr_backends.py` | `Ocr(pytesseract, paddleocr, rapidocr)` β also reused for table-cell OCR |
|
| 41 |
+
| `table_extraction.py` | `TableFormerONNX` β OTSL β markdown |
|
| 42 |
+
| `pipeline.py` | The `if the words are empty` branch, the `if classname is table` branch, and final markdown assembly (`process_pdf_page`, `process_pdf`) |
|
| 43 |
+
|
| 44 |
+
## What's new vs. the two source notebooks
|
| 45 |
+
|
| 46 |
+
The two original notebooks (`pdfplumber_.ipynb`, `docling_.ipynb`) already
|
| 47 |
+
had working code for: layout detection, the pdfplumberβlayout matcher, OCR
|
| 48 |
+
backends, and TableFormerβOTSL. What was missing, per the diagram, was:
|
| 49 |
+
|
| 50 |
+
- **The OCR-fallback matching box.** `align_words_to_layout` only knew how
|
| 51 |
+
to consume pdfplumber's `extract_words()` output. `align_tokens_to_layout`
|
| 52 |
+
generalizes the matching algorithm (containment + nearest-centroid
|
| 53 |
+
fallback + line re-assembly) to take any `text + bbox` token list, in
|
| 54 |
+
image-pixel space. `align_words_to_layout` and `align_ocr_to_layout` are
|
| 55 |
+
now both adapters onto that one function.
|
| 56 |
+
- **The `if classname is table` branch wired into the page loop.** The
|
| 57 |
+
notebook's per-region markdown loop only ever embedded a cropped image for
|
| 58 |
+
`table`/`image`/`chart` regions; TableFormer was only run manually,
|
| 59 |
+
separately, against one standalone table image. `region_to_markdown` now
|
| 60 |
+
checks `class_name == "table"` and runs `table_image_to_otsl` +
|
| 61 |
+
`otsl_to_markdown` automatically, falling back to an embedded image if no
|
| 62 |
+
`TableFormerONNX` runner was supplied (or if extraction throws).
|
| 63 |
+
- **Table cells prefer pdfplumber's own words over OCR.** When the PDF has a
|
| 64 |
+
real text layer, the words pdfplumber already extracted for a table region
|
| 65 |
+
are reused directly as TableFormer's cell-text source (`tokens=` on
|
| 66 |
+
`ocr_anchor_cells`/`table_image_to_otsl`) instead of re-OCR-ing the table
|
| 67 |
+
crop β cheaper and avoids OCR mistakes on text that's already exact.
|
| 68 |
+
`table_ocr_backend` is now only invoked as a fallback when a table region
|
| 69 |
+
has no underlying words (e.g. the table is itself a scanned image, or the
|
| 70 |
+
whole page is scanned and went through the page-level OCR branch).
|
| 71 |
+
- **Standalone image input.** `process_image_page`/`process_images` run the
|
| 72 |
+
same diagram on plain images (.jpg/.png/etc. β a photographed or scanned
|
| 73 |
+
page with no PDF structure at all), always via the OCR branch since
|
| 74 |
+
there's no native text layer to check. `process_document` is a single
|
| 75 |
+
dispatcher that picks the PDF path or the image path based on the input,
|
| 76 |
+
so callers don't need to branch on file type themselves. Both share the
|
| 77 |
+
same per-image core (`_process_rendered_page`) as the PDF path, so a
|
| 78 |
+
scanned PDF page and a standalone photo of a page are handled identically
|
| 79 |
+
once OCR kicks in.
|
| 80 |
+
- **Fixed a latent color-channel bug along the way.** `DocLayoutV3._preprocess`
|
| 81 |
+
always applies a BGRβRGB swap, correct only if its input really is BGR
|
| 82 |
+
(as `cv2.imread` produces). The original page-rendering code fed it a
|
| 83 |
+
PIL-derived array (RGB) directly, which silently double-swapped the
|
| 84 |
+
channels and degraded detection on every PDF page processed through this
|
| 85 |
+
pipeline. Fixed via a shared `_to_bgr_array` helper used by both the PDF
|
| 86 |
+
and image entry points.
|
| 87 |
+
- **`otsl_to_markdown`.** The OTSLβDocTagsDocumentβDoclingDocumentβmarkdown
|
| 88 |
+
steps were ad hoc, later notebook cells, applied to one hardcoded `otsl`
|
| 89 |
+
variable. Factored into one function so it can run once per detected
|
| 90 |
+
table inside a page loop.
|
| 91 |
+
|
| 92 |
+
## Requirements
|
| 93 |
+
|
| 94 |
+
```
|
| 95 |
+
pip install pdfplumber numpy opencv-python onnxruntime pydantic pillow docling-core
|
| 96 |
+
# plus whichever OCR backend(s) you use:
|
| 97 |
+
pip install pytesseract # needs the tesseract binary too
|
| 98 |
+
pip install rapidocr-onnxruntime
|
| 99 |
+
pip install paddleocr paddlepaddle
|
| 100 |
+
```
|
| 101 |
+
|
| 102 |
+
You'll also need the model artifacts referenced in the original notebooks:
|
| 103 |
+
- `PP-DocLayout/PP-DocLayoutV3.onnx`
|
| 104 |
+
- `tableformerv1/onnx/<variant>/tableformer_<variant>_{encoder,decoder_step,bbox_decoder}.onnx`
|
| 105 |
+
+ `tableformerv1/tm_config.json`
|
| 106 |
+
|
| 107 |
+
## Usage
|
| 108 |
+
|
| 109 |
+
```python
|
| 110 |
+
from pdf_pipeline import DocLayoutV3, TableFormerONNX, get_ocr_backend, process_document
|
| 111 |
+
|
| 112 |
+
layout_detector = DocLayoutV3("PP-DocLayout/PP-DocLayoutV3.onnx")
|
| 113 |
+
page_ocr_backend = get_ocr_backend("rapidocr") # used on scanned PDF pages AND any standalone image
|
| 114 |
+
table_runner = TableFormerONNX(artifact_root="tableformerv1")
|
| 115 |
+
table_ocr_backend = get_ocr_backend("paddleocr") # fallback for tables with no underlying text
|
| 116 |
+
|
| 117 |
+
# Works the same regardless of input type:
|
| 118 |
+
markdown_doc = process_document(
|
| 119 |
+
"document.pdf", # or "scan.jpg", or ["page1.png", "page2.png"]
|
| 120 |
+
layout_detector,
|
| 121 |
+
page_ocr_backend=page_ocr_backend,
|
| 122 |
+
table_runner=table_runner,
|
| 123 |
+
table_ocr_backend=table_ocr_backend,
|
| 124 |
+
)
|
| 125 |
+
|
| 126 |
+
with open("document.md", "w") as f:
|
| 127 |
+
f.write(markdown_doc)
|
| 128 |
+
```
|
| 129 |
+
|
| 130 |
+
`process_document` dispatches on what you pass it: a `.pdf` path goes
|
| 131 |
+
through `process_pdf` (pdfplumber, may have a native text layer); a single
|
| 132 |
+
image path (`.jpg`/`.png`/etc.) or a list of image paths goes through
|
| 133 |
+
`process_image_page`/`process_images` (always OCR, since standalone images
|
| 134 |
+
never have a text layer). Call `process_pdf`/`process_image_page` directly
|
| 135 |
+
if you want the type-specific kwargs (e.g. `pages=` or `resolution=` only
|
| 136 |
+
make sense for PDFs) or the full per-page debug dict instead of just the
|
| 137 |
+
markdown string.
|
| 138 |
+
|
| 139 |
+
For single-page debugging (mirrors the original notebook's main cell), use
|
| 140 |
+
`process_pdf_page(page, layout_detector, doc_path, ...)` directly β it
|
| 141 |
+
returns a dict with `markdown`, `regions`, `layout_result`, `image`, and
|
| 142 |
+
`used_ocr`, so you can inspect intermediate state (e.g. call
|
| 143 |
+
`layout_detector.visualize(result["layout_result"], "debug.png", result["image"])`)
|
| 144 |
+
before trusting the final markdown.
|
| 145 |
+
|
| 146 |
+
Cropped images (for `image`/`chart`/`table`-without-TableFormer regions) are
|
| 147 |
+
written to `<doc_stem>/images/<n>.png`, same convention as the original
|
| 148 |
+
notebook's `crop_and_save_image`.
|
| 149 |
+
|
| 150 |
+
## Notes / things to tune for your documents
|
| 151 |
+
|
| 152 |
+
- `containment_threshold` / `line_tol_ratio` (passed through `process_pdf_page`
|
| 153 |
+
β `align_words_to_layout` / `align_ocr_to_layout` β `align_tokens_to_layout`)
|
| 154 |
+
control how aggressively tokens are assigned to boxes and how lines are
|
| 155 |
+
re-grouped; the defaults (0.5 / 0.5) come from the original notebook.
|
| 156 |
+
- `SKIP_CLASSES` and `IMAGE_LIKE_CLASSES` in `pipeline.py` control which
|
| 157 |
+
layout classes are dropped vs. rendered as a markdown image link vs.
|
| 158 |
+
rendered as text β extend these for your document types (e.g. treat
|
| 159 |
+
`display_formula`/`inline_formula` specially instead of falling through
|
| 160 |
+
to plain text).
|
| 161 |
+
- If `table_runner`/`table_ocr_backend` are omitted, table regions degrade
|
| 162 |
+
gracefully to an embedded cropped image instead of raising.
|
__init__.py
ADDED
|
@@ -0,0 +1,87 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
pdf_pipeline
|
| 3 |
+
============
|
| 4 |
+
|
| 5 |
+
PDF -> markdown extraction pipeline:
|
| 6 |
+
|
| 7 |
+
pdfplumber (words+bboxes, image)
|
| 8 |
+
-> [empty words? -> OCR] -> match against pp-doclayout boxes
|
| 9 |
+
-> per-region markdown, with table regions routed through TableFormerONNX -> OTSL -> markdown
|
| 10 |
+
|
| 11 |
+
See README.md for the full architecture diagram and a usage walkthrough.
|
| 12 |
+
"""
|
| 13 |
+
|
| 14 |
+
from .layout import (
|
| 15 |
+
BBox,
|
| 16 |
+
DocLayoutResult,
|
| 17 |
+
DocLayoutV3,
|
| 18 |
+
LayoutBox,
|
| 19 |
+
PlacedWord,
|
| 20 |
+
RegionText,
|
| 21 |
+
align_ocr_to_layout,
|
| 22 |
+
align_tokens_to_layout,
|
| 23 |
+
align_words_to_layout,
|
| 24 |
+
log_layout_result,
|
| 25 |
+
pdfplumber_tokens_in_image_space,
|
| 26 |
+
print_layout_result,
|
| 27 |
+
)
|
| 28 |
+
from .logging_config import setup_pipeline_logging
|
| 29 |
+
from .ocr_backends import (
|
| 30 |
+
OCRBackend,
|
| 31 |
+
PytesseractBackend,
|
| 32 |
+
RapidOCRBackend,
|
| 33 |
+
get_ocr_backend,
|
| 34 |
+
)
|
| 35 |
+
from .pipeline import (
|
| 36 |
+
IMAGE_EXTENSIONS,
|
| 37 |
+
crop_and_save_image,
|
| 38 |
+
process_document,
|
| 39 |
+
process_images,
|
| 40 |
+
process_image_page,
|
| 41 |
+
process_pdf,
|
| 42 |
+
process_pdf_page,
|
| 43 |
+
region_to_markdown,
|
| 44 |
+
regions_to_markdown,
|
| 45 |
+
)
|
| 46 |
+
from .table_extraction import (
|
| 47 |
+
TableFormerONNX,
|
| 48 |
+
ocr_anchor_cells,
|
| 49 |
+
otsl_to_markdown,
|
| 50 |
+
seq_to_otsl,
|
| 51 |
+
table_image_to_otsl,
|
| 52 |
+
)
|
| 53 |
+
|
| 54 |
+
__all__ = [
|
| 55 |
+
"BBox",
|
| 56 |
+
"DocLayoutResult",
|
| 57 |
+
"DocLayoutV3",
|
| 58 |
+
"LayoutBox",
|
| 59 |
+
"PlacedWord",
|
| 60 |
+
"RegionText",
|
| 61 |
+
"align_ocr_to_layout",
|
| 62 |
+
"align_tokens_to_layout",
|
| 63 |
+
"align_words_to_layout",
|
| 64 |
+
"pdfplumber_tokens_in_image_space",
|
| 65 |
+
"log_layout_result",
|
| 66 |
+
"print_layout_result",
|
| 67 |
+
"OCRBackend",
|
| 68 |
+
"PaddleOCRBackend",
|
| 69 |
+
"PytesseractBackend",
|
| 70 |
+
"RapidOCRBackend",
|
| 71 |
+
"get_ocr_backend",
|
| 72 |
+
"crop_and_save_image",
|
| 73 |
+
"IMAGE_EXTENSIONS",
|
| 74 |
+
"process_document",
|
| 75 |
+
"process_images",
|
| 76 |
+
"process_image_page",
|
| 77 |
+
"process_pdf",
|
| 78 |
+
"process_pdf_page",
|
| 79 |
+
"region_to_markdown",
|
| 80 |
+
"regions_to_markdown",
|
| 81 |
+
"TableFormerONNX",
|
| 82 |
+
"ocr_anchor_cells",
|
| 83 |
+
"otsl_to_markdown",
|
| 84 |
+
"seq_to_otsl",
|
| 85 |
+
"table_image_to_otsl",
|
| 86 |
+
"setup_pipeline_logging",
|
| 87 |
+
]
|
__pycache__/__init__.cpython-310.pyc
ADDED
|
Binary file (1.7 kB). View file
|
|
|
__pycache__/layout.cpython-310.pyc
ADDED
|
Binary file (20.5 kB). View file
|
|
|
__pycache__/logging_config.cpython-310.pyc
ADDED
|
Binary file (1.52 kB). View file
|
|
|
__pycache__/ocr_backends.cpython-310.pyc
ADDED
|
Binary file (4.61 kB). View file
|
|
|
__pycache__/pipeline.cpython-310.pyc
ADDED
|
Binary file (15.6 kB). View file
|
|
|
__pycache__/table_extraction.cpython-310.pyc
ADDED
|
Binary file (21.1 kB). View file
|
|
|
layout.py
ADDED
|
@@ -0,0 +1,672 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
layout.py
|
| 3 |
+
=========
|
| 4 |
+
|
| 5 |
+
Two responsibilities, matching the top half of the pipeline diagram:
|
| 6 |
+
|
| 7 |
+
1. `DocLayoutV3` β thin ONNXRuntime wrapper around PP-DocLayoutV3. Given a
|
| 8 |
+
page image it returns layout boxes: (class_name, bbox[xyxy], reading_order).
|
| 9 |
+
This is unchanged from the original pdfplumber_.ipynb.
|
| 10 |
+
|
| 11 |
+
2. Matching: turn a flat list of "tokens" (word/line-level text + bbox) into
|
| 12 |
+
per-layout-region text, by containment + nearest-centroid fallback, then
|
| 13 |
+
re-assembling lines within each region.
|
| 14 |
+
|
| 15 |
+
The original notebook only had `align_words_to_layout`, which is hard-wired
|
| 16 |
+
to pdfplumber's `page.extract_words()` output (point-space, needs scaling
|
| 17 |
+
to image pixels). The diagram has a *second*, parallel matching box
|
| 18 |
+
("Matching ocr bboxes and layout") for the branch where pdfplumber finds no
|
| 19 |
+
extractable words and an OCR backend is used instead. OCR tokens are
|
| 20 |
+
already in image-pixel space (OCR runs on the same rendered image that was
|
| 21 |
+
fed to PP-DocLayout), so no scaling is needed there.
|
| 22 |
+
|
| 23 |
+
Rather than duplicating the matching logic, `align_tokens_to_layout` is the
|
| 24 |
+
single generic implementation. `align_words_to_layout` and
|
| 25 |
+
`align_ocr_to_layout` are thin adapters that normalize their respective
|
| 26 |
+
inputs into the same token shape and call it.
|
| 27 |
+
"""
|
| 28 |
+
|
| 29 |
+
from __future__ import annotations
|
| 30 |
+
|
| 31 |
+
import logging
|
| 32 |
+
from dataclasses import dataclass
|
| 33 |
+
from pathlib import Path
|
| 34 |
+
from typing import Dict, List, Optional, Sequence, Tuple, Union
|
| 35 |
+
|
| 36 |
+
logger = logging.getLogger(__name__)
|
| 37 |
+
|
| 38 |
+
import cv2
|
| 39 |
+
import numpy as np
|
| 40 |
+
from pydantic import BaseModel, Field
|
| 41 |
+
|
| 42 |
+
try:
|
| 43 |
+
import onnxruntime as ort
|
| 44 |
+
except ImportError: # pragma: no cover - allows importing this module for
|
| 45 |
+
ort = None # static inspection without onnxruntime installed.
|
| 46 |
+
|
| 47 |
+
|
| 48 |
+
LAYOUT_CLASSES = [
|
| 49 |
+
"abstract",
|
| 50 |
+
"algorithm",
|
| 51 |
+
"aside_text",
|
| 52 |
+
"chart",
|
| 53 |
+
"content",
|
| 54 |
+
"display_formula",
|
| 55 |
+
"doc_title",
|
| 56 |
+
"figure_title",
|
| 57 |
+
"footer",
|
| 58 |
+
"footer_image",
|
| 59 |
+
"footnote",
|
| 60 |
+
"formula_number",
|
| 61 |
+
"header",
|
| 62 |
+
"header_image",
|
| 63 |
+
"image",
|
| 64 |
+
"inline_formula",
|
| 65 |
+
"number",
|
| 66 |
+
"paragraph_title",
|
| 67 |
+
"reference",
|
| 68 |
+
"reference_content",
|
| 69 |
+
"seal",
|
| 70 |
+
"table",
|
| 71 |
+
"text",
|
| 72 |
+
"vertical_text",
|
| 73 |
+
"vision_footnote",
|
| 74 |
+
]
|
| 75 |
+
|
| 76 |
+
|
| 77 |
+
# --------------------------------------------------------------------------- #
|
| 78 |
+
# Layout detection result types
|
| 79 |
+
# --------------------------------------------------------------------------- #
|
| 80 |
+
class BBox(BaseModel):
|
| 81 |
+
"""Axis-aligned bounding box in original image pixel coordinates."""
|
| 82 |
+
|
| 83 |
+
xmin: float
|
| 84 |
+
ymin: float
|
| 85 |
+
xmax: float
|
| 86 |
+
ymax: float
|
| 87 |
+
|
| 88 |
+
def as_tuple(self) -> Tuple[float, float, float, float]:
|
| 89 |
+
return (self.xmin, self.ymin, self.xmax, self.ymax)
|
| 90 |
+
|
| 91 |
+
|
| 92 |
+
class LayoutBox(BaseModel):
|
| 93 |
+
"""A single detected layout element."""
|
| 94 |
+
|
| 95 |
+
class_id: int
|
| 96 |
+
class_name: str
|
| 97 |
+
score: float
|
| 98 |
+
bbox: BBox
|
| 99 |
+
reading_order: int
|
| 100 |
+
|
| 101 |
+
model_config = {"frozen": True}
|
| 102 |
+
|
| 103 |
+
|
| 104 |
+
class DocLayoutResult(BaseModel):
|
| 105 |
+
"""Full result for one image."""
|
| 106 |
+
|
| 107 |
+
image_path: Optional[str] = None
|
| 108 |
+
image_width: int
|
| 109 |
+
image_height: int
|
| 110 |
+
boxes: List[LayoutBox] = Field(default_factory=list)
|
| 111 |
+
|
| 112 |
+
def sorted_by_reading_order(self) -> List[LayoutBox]:
|
| 113 |
+
return sorted(self.boxes, key=lambda b: b.reading_order)
|
| 114 |
+
|
| 115 |
+
|
| 116 |
+
# --------------------------------------------------------------------------- #
|
| 117 |
+
# PP-DocLayoutV3 ONNX wrapper
|
| 118 |
+
# --------------------------------------------------------------------------- #
|
| 119 |
+
|
| 120 |
+
class DocLayoutV3:
|
| 121 |
+
"""
|
| 122 |
+
Thin OO wrapper around the PP-DocLayoutV3 ONNX model.
|
| 123 |
+
|
| 124 |
+
Usage:
|
| 125 |
+
detector = DocLayoutV3("PP-DocLayout/PP-DocLayoutV3.onnx")
|
| 126 |
+
result = detector.predict("page.png")
|
| 127 |
+
for box in result.sorted_by_reading_order():
|
| 128 |
+
print(box.class_name, box.bbox, box.reading_order)
|
| 129 |
+
|
| 130 |
+
# optional: draw + save a visualization
|
| 131 |
+
detector.visualize(result, "pp_doclayout.png")
|
| 132 |
+
"""
|
| 133 |
+
|
| 134 |
+
def __init__(
|
| 135 |
+
self,
|
| 136 |
+
model_path: Union[str, Path],
|
| 137 |
+
target_input_size: Tuple[int, int] = (800, 800),
|
| 138 |
+
score_threshold: float = 0.5,
|
| 139 |
+
class_names: Optional[List[str]] = None,
|
| 140 |
+
providers: Optional[List[str]] = None,
|
| 141 |
+
intra_op_num_threads: Optional[int] = None,
|
| 142 |
+
inter_op_num_threads: Optional[int] = None,
|
| 143 |
+
enable_mem_pattern: bool = True,
|
| 144 |
+
graph_optimization_level: "ort.GraphOptimizationLevel" = None,
|
| 145 |
+
):
|
| 146 |
+
if ort is None:
|
| 147 |
+
raise ImportError("onnxruntime is required for DocLayoutV3")
|
| 148 |
+
|
| 149 |
+
graph_optimization_level = (
|
| 150 |
+
graph_optimization_level or ort.GraphOptimizationLevel.ORT_ENABLE_ALL
|
| 151 |
+
)
|
| 152 |
+
|
| 153 |
+
self.model_path = str(model_path)
|
| 154 |
+
self.target_input_size = target_input_size
|
| 155 |
+
self.score_threshold = score_threshold
|
| 156 |
+
self.class_names = class_names or LAYOUT_CLASSES
|
| 157 |
+
|
| 158 |
+
self.session = self._build_session(
|
| 159 |
+
providers=providers,
|
| 160 |
+
intra_op_num_threads=intra_op_num_threads,
|
| 161 |
+
inter_op_num_threads=inter_op_num_threads,
|
| 162 |
+
enable_mem_pattern=enable_mem_pattern,
|
| 163 |
+
graph_optimization_level=graph_optimization_level,
|
| 164 |
+
)
|
| 165 |
+
self.providers_used = self.session.get_providers()
|
| 166 |
+
|
| 167 |
+
self.input_names = [i.name for i in self.session.get_inputs()]
|
| 168 |
+
self.output_names = [o.name for o in self.session.get_outputs()]
|
| 169 |
+
|
| 170 |
+
self._mean = np.array([0.0, 0.0, 0.0], dtype=np.float32)
|
| 171 |
+
self._std = std = np.array([1.0, 1.0, 1.0], dtype=np.float32)
|
| 172 |
+
|
| 173 |
+
|
| 174 |
+
# ------------------------------------------------------------------ #
|
| 175 |
+
def _build_session(
|
| 176 |
+
self,
|
| 177 |
+
providers: Optional[List[str]],
|
| 178 |
+
intra_op_num_threads: Optional[int],
|
| 179 |
+
inter_op_num_threads: Optional[int],
|
| 180 |
+
enable_mem_pattern: bool,
|
| 181 |
+
graph_optimization_level,
|
| 182 |
+
) -> "ort.InferenceSession":
|
| 183 |
+
available = ort.get_available_providers()
|
| 184 |
+
|
| 185 |
+
if providers is not None:
|
| 186 |
+
resolved = [p for p in providers if p in available]
|
| 187 |
+
if not resolved:
|
| 188 |
+
raise ValueError(
|
| 189 |
+
f"None of the requested providers {providers} are "
|
| 190 |
+
f"available. Available providers: {available}"
|
| 191 |
+
)
|
| 192 |
+
else:
|
| 193 |
+
preferred_order = ["CUDAExecutionProvider", "CPUExecutionProvider"]
|
| 194 |
+
resolved = [p for p in preferred_order if p in available]
|
| 195 |
+
if not resolved:
|
| 196 |
+
resolved = available or ["CPUExecutionProvider"]
|
| 197 |
+
|
| 198 |
+
sess_options = ort.SessionOptions()
|
| 199 |
+
sess_options.graph_optimization_level = graph_optimization_level
|
| 200 |
+
sess_options.enable_mem_pattern = enable_mem_pattern
|
| 201 |
+
sess_options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
|
| 202 |
+
|
| 203 |
+
if intra_op_num_threads is not None:
|
| 204 |
+
sess_options.intra_op_num_threads = intra_op_num_threads
|
| 205 |
+
if inter_op_num_threads is not None:
|
| 206 |
+
sess_options.inter_op_num_threads = inter_op_num_threads
|
| 207 |
+
|
| 208 |
+
session = ort.InferenceSession(
|
| 209 |
+
self.model_path,
|
| 210 |
+
sess_options=sess_options,
|
| 211 |
+
providers=resolved,
|
| 212 |
+
)
|
| 213 |
+
|
| 214 |
+
actual = session.get_providers()
|
| 215 |
+
if resolved[0] not in actual:
|
| 216 |
+
import warnings
|
| 217 |
+
|
| 218 |
+
warnings.warn(
|
| 219 |
+
f"Requested provider order {resolved} but session is "
|
| 220 |
+
f"actually running on {actual}.",
|
| 221 |
+
RuntimeWarning,
|
| 222 |
+
stacklevel=2,
|
| 223 |
+
)
|
| 224 |
+
|
| 225 |
+
return session
|
| 226 |
+
|
| 227 |
+
#--------------------------------------------------------------------#
|
| 228 |
+
def add_border_to_image(self, image: np.ndarray) -> np.ndarray:
|
| 229 |
+
"""
|
| 230 |
+
Add a border to the image to ensure it is square.
|
| 231 |
+
"""
|
| 232 |
+
orig_h, orig_w = image.shape[:2]
|
| 233 |
+
if orig_h > orig_w:
|
| 234 |
+
pad_h = (orig_h - orig_w) // 2
|
| 235 |
+
pad_v = (orig_h - orig_w) // 2
|
| 236 |
+
else:
|
| 237 |
+
pad_h = (orig_w - orig_h) // 2
|
| 238 |
+
pad_v = (orig_w - orig_h) // 2
|
| 239 |
+
return cv2.copyMakeBorder(
|
| 240 |
+
image,
|
| 241 |
+
pad_v,
|
| 242 |
+
pad_v,
|
| 243 |
+
pad_h,
|
| 244 |
+
pad_h,
|
| 245 |
+
cv2.BORDER_CONSTANT,
|
| 246 |
+
value=(0, 0, 0),
|
| 247 |
+
)
|
| 248 |
+
height, width = image.shape[:2]
|
| 249 |
+
max_dim = max(height, width)
|
| 250 |
+
border = (max_dim - height) // 2, (max_dim - width) // 2
|
| 251 |
+
image = cv2.copyMakeBorder(
|
| 252 |
+
image,
|
| 253 |
+
border[0],
|
| 254 |
+
border[0],
|
| 255 |
+
border[1],
|
| 256 |
+
border[1],
|
| 257 |
+
cv2.BORDER_CONSTANT,
|
| 258 |
+
value=[255, 255, 255],
|
| 259 |
+
)
|
| 260 |
+
return image
|
| 261 |
+
# ------------------------------------------------------------------ #
|
| 262 |
+
def median_filter(self, image:np.ndarray):
|
| 263 |
+
"""
|
| 264 |
+
Apply median filter to the image.
|
| 265 |
+
"""
|
| 266 |
+
return cv2.medianBlur(image, 3)
|
| 267 |
+
def _preprocess(self, image: np.ndarray) -> Tuple[np.ndarray, float, float]:
|
| 268 |
+
orig_h, orig_w = image.shape[:2]
|
| 269 |
+
target_h, target_w = self.target_input_size
|
| 270 |
+
#image = self.add_border_to_image(image)
|
| 271 |
+
|
| 272 |
+
|
| 273 |
+
scale_h = target_h / orig_h
|
| 274 |
+
scale_w = target_w / orig_w
|
| 275 |
+
|
| 276 |
+
|
| 277 |
+
resized = cv2.resize(image, (target_w, target_h), interpolation=cv2.INTER_CUBIC)
|
| 278 |
+
|
| 279 |
+
rgb = cv2.cvtColor(resized, cv2.COLOR_BGR2RGB)
|
| 280 |
+
blob = rgb.astype(np.float32) / 255.0
|
| 281 |
+
blob = (blob - self._mean) / self._std
|
| 282 |
+
|
| 283 |
+
blob = blob.transpose(2, 0, 1)[np.newaxis, ...]
|
| 284 |
+
return blob, scale_h, scale_w
|
| 285 |
+
|
| 286 |
+
def _infer(self, image: np.ndarray) -> Tuple[np.ndarray, int, int]:
|
| 287 |
+
orig_h, orig_w = image.shape[:2]
|
| 288 |
+
input_blob, scale_h, scale_w = self._preprocess(image)
|
| 289 |
+
|
| 290 |
+
target_h, target_w = self.target_input_size
|
| 291 |
+
preprocess_shape = [np.array([target_h, target_w], dtype=np.float32)]
|
| 292 |
+
|
| 293 |
+
input_feed = {
|
| 294 |
+
self.input_names[0]: preprocess_shape,
|
| 295 |
+
self.input_names[1]: input_blob,
|
| 296 |
+
self.input_names[2]: [[scale_h, scale_w]],
|
| 297 |
+
}
|
| 298 |
+
|
| 299 |
+
# shape=(N, 7): [label_index, score, xmin, ymin, xmax, ymax, read_order]
|
| 300 |
+
output = self.session.run(self.output_names, input_feed)[0]
|
| 301 |
+
return output, orig_h, orig_w
|
| 302 |
+
|
| 303 |
+
def _to_result(
|
| 304 |
+
self,
|
| 305 |
+
raw_boxes: np.ndarray,
|
| 306 |
+
orig_h: int,
|
| 307 |
+
orig_w: int,
|
| 308 |
+
image_path: Optional[str],
|
| 309 |
+
) -> DocLayoutResult:
|
| 310 |
+
logger.info(f"raw_boxes: {raw_boxes}")
|
| 311 |
+
|
| 312 |
+
filtered = raw_boxes[raw_boxes[:, 1] > self.score_threshold]
|
| 313 |
+
filtered = filtered[np.argsort(filtered[:, 5])]
|
| 314 |
+
|
| 315 |
+
|
| 316 |
+
boxes: List[LayoutBox] = []
|
| 317 |
+
for row in filtered:
|
| 318 |
+
cls_id = int(row[0])
|
| 319 |
+
boxes.append(
|
| 320 |
+
LayoutBox(
|
| 321 |
+
class_id=cls_id,
|
| 322 |
+
class_name=self.class_names[cls_id]
|
| 323 |
+
if 0 <= cls_id < len(self.class_names)
|
| 324 |
+
else str(cls_id),
|
| 325 |
+
score=float(row[1]),
|
| 326 |
+
bbox=BBox(
|
| 327 |
+
xmin=float(row[2]),
|
| 328 |
+
ymin=float(row[3]),
|
| 329 |
+
xmax=float(row[4]),
|
| 330 |
+
ymax=float(row[5]),
|
| 331 |
+
),
|
| 332 |
+
reading_order=int(row[6]),
|
| 333 |
+
)
|
| 334 |
+
)
|
| 335 |
+
|
| 336 |
+
result = DocLayoutResult(
|
| 337 |
+
image_path=image_path,
|
| 338 |
+
image_width=orig_w,
|
| 339 |
+
image_height=orig_h,
|
| 340 |
+
boxes=boxes,
|
| 341 |
+
)
|
| 342 |
+
logger.info(
|
| 343 |
+
"DocLayoutV3: %dx%d image -> %d boxes (threshold=%.2f)",
|
| 344 |
+
orig_w,
|
| 345 |
+
orig_h,
|
| 346 |
+
len(boxes),
|
| 347 |
+
self.score_threshold,
|
| 348 |
+
)
|
| 349 |
+
return result
|
| 350 |
+
|
| 351 |
+
# ------------------------------------------------------------------ #
|
| 352 |
+
# Public API
|
| 353 |
+
# ------------------------------------------------------------------ #
|
| 354 |
+
def predict(self, image: Union[str, Path, np.ndarray]) -> DocLayoutResult:
|
| 355 |
+
image_path: Optional[str] = None
|
| 356 |
+
if isinstance(image, (str, Path)):
|
| 357 |
+
image_path = str(image)
|
| 358 |
+
img_array = cv2.imread(image_path)
|
| 359 |
+
if img_array is None:
|
| 360 |
+
raise FileNotFoundError(f"Could not read image: {image_path}")
|
| 361 |
+
else:
|
| 362 |
+
img_array = image
|
| 363 |
+
|
| 364 |
+
raw_boxes, orig_h, orig_w = self._infer(img_array)
|
| 365 |
+
return self._to_result(raw_boxes, orig_h, orig_w, image_path)
|
| 366 |
+
|
| 367 |
+
def predict_batch(
|
| 368 |
+
self, images: List[Union[str, Path, np.ndarray]]
|
| 369 |
+
) -> List[DocLayoutResult]:
|
| 370 |
+
return [self.predict(img) for img in images]
|
| 371 |
+
|
| 372 |
+
def visualize(
|
| 373 |
+
self,
|
| 374 |
+
result: DocLayoutResult,
|
| 375 |
+
output_path: Union[str, Path],
|
| 376 |
+
source_image: Optional[Union[str, Path, np.ndarray]] = None,
|
| 377 |
+
) -> str:
|
| 378 |
+
src = source_image if source_image is not None else result.image_path
|
| 379 |
+
if src is None:
|
| 380 |
+
raise ValueError(
|
| 381 |
+
"No source image available to draw on; pass source_image=..."
|
| 382 |
+
)
|
| 383 |
+
|
| 384 |
+
if isinstance(src, (str, Path)):
|
| 385 |
+
img = cv2.imread(str(src))
|
| 386 |
+
if img is None:
|
| 387 |
+
raise FileNotFoundError(f"Could not read image: {src}")
|
| 388 |
+
else:
|
| 389 |
+
img = src.copy()
|
| 390 |
+
|
| 391 |
+
for box in result.boxes:
|
| 392 |
+
x0, y0, x1, y1 = (int(round(v)) for v in box.bbox.as_tuple())
|
| 393 |
+
cv2.rectangle(img, (x0, y0), (x1, y1), (0, 0, 255), 2)
|
| 394 |
+
label = f"{box.reading_order}|{box.class_name}"
|
| 395 |
+
text_y = max(y0 - 10, 0)
|
| 396 |
+
cv2.putText(
|
| 397 |
+
img,
|
| 398 |
+
label,
|
| 399 |
+
(x0, text_y),
|
| 400 |
+
cv2.FONT_HERSHEY_SIMPLEX,
|
| 401 |
+
0.5,
|
| 402 |
+
(255, 0, 0),
|
| 403 |
+
1,
|
| 404 |
+
cv2.LINE_AA,
|
| 405 |
+
)
|
| 406 |
+
|
| 407 |
+
output_path = str(output_path)
|
| 408 |
+
cv2.imwrite(output_path, img)
|
| 409 |
+
return output_path
|
| 410 |
+
|
| 411 |
+
|
| 412 |
+
def log_layout_result(result: DocLayoutResult) -> None:
|
| 413 |
+
"""Log every detected layout box at INFO level."""
|
| 414 |
+
for box in result.sorted_by_reading_order():
|
| 415 |
+
b = box.bbox
|
| 416 |
+
logger.info(
|
| 417 |
+
" layout [%d] %s score=%.3f bbox=(%.0f, %.0f, %.0f, %.0f)",
|
| 418 |
+
box.reading_order,
|
| 419 |
+
box.class_name,
|
| 420 |
+
box.score,
|
| 421 |
+
b.xmin,
|
| 422 |
+
b.ymin,
|
| 423 |
+
b.xmax,
|
| 424 |
+
b.ymax,
|
| 425 |
+
)
|
| 426 |
+
|
| 427 |
+
|
| 428 |
+
def print_layout_result(result: DocLayoutResult) -> None:
|
| 429 |
+
"""Print layout boxes to stdout (legacy helper; prefer ``log_layout_result``)."""
|
| 430 |
+
log_layout_result(result)
|
| 431 |
+
|
| 432 |
+
|
| 433 |
+
# --------------------------------------------------------------------------- #
|
| 434 |
+
# Token <-> layout matching (the two "matching" boxes in the diagram)
|
| 435 |
+
# --------------------------------------------------------------------------- #
|
| 436 |
+
@dataclass(frozen=True)
|
| 437 |
+
class PlacedWord:
|
| 438 |
+
text: str
|
| 439 |
+
x0: float
|
| 440 |
+
x1: float
|
| 441 |
+
top: float
|
| 442 |
+
bottom: float
|
| 443 |
+
|
| 444 |
+
|
| 445 |
+
@dataclass
|
| 446 |
+
class RegionText:
|
| 447 |
+
reading_order: int
|
| 448 |
+
class_name: str
|
| 449 |
+
bbox: tuple # (xmin, ymin, xmax, ymax) in image-pixel space
|
| 450 |
+
text: str
|
| 451 |
+
words: List[PlacedWord]
|
| 452 |
+
|
| 453 |
+
|
| 454 |
+
def _line_cluster_fast(words: List[PlacedWord], tol: float) -> List[List[PlacedWord]]:
|
| 455 |
+
"""Group words into lines by `top` proximity, then sort each line left-to-right.
|
| 456 |
+
`tol` is computed once globally by the caller instead of per-region."""
|
| 457 |
+
if not words:
|
| 458 |
+
return []
|
| 459 |
+
ordered = sorted(words, key=lambda w: w.top)
|
| 460 |
+
lines: List[List[PlacedWord]] = []
|
| 461 |
+
current = [ordered[0]]
|
| 462 |
+
current_top = ordered[0].top
|
| 463 |
+
for w in ordered[1:]:
|
| 464 |
+
if abs(w.top - current_top) <= tol:
|
| 465 |
+
current.append(w)
|
| 466 |
+
else:
|
| 467 |
+
lines.append(current)
|
| 468 |
+
current = [w]
|
| 469 |
+
current_top = w.top
|
| 470 |
+
lines.append(current)
|
| 471 |
+
for line in lines:
|
| 472 |
+
line.sort(key=lambda w: w.x0)
|
| 473 |
+
return lines
|
| 474 |
+
|
| 475 |
+
|
| 476 |
+
def align_tokens_to_layout(
|
| 477 |
+
tokens: Sequence[PlacedWord],
|
| 478 |
+
layout_result: DocLayoutResult,
|
| 479 |
+
containment_threshold: float = 0.5,
|
| 480 |
+
line_tol_ratio: float = 0.5,
|
| 481 |
+
) -> List[RegionText]:
|
| 482 |
+
"""
|
| 483 |
+
Generic matcher: given tokens already expressed in image-pixel space
|
| 484 |
+
(same coordinate system as `layout_result`), assign each token to its
|
| 485 |
+
best-matching layout box (smallest box containing it above
|
| 486 |
+
`containment_threshold`, falling back to nearest centroid for tokens
|
| 487 |
+
that don't sit inside any box), then re-assemble lines of text per box.
|
| 488 |
+
|
| 489 |
+
This is the shared core behind both `align_words_to_layout` (pdfplumber
|
| 490 |
+
path) and `align_ocr_to_layout` (OCR fallback path) in the diagram.
|
| 491 |
+
"""
|
| 492 |
+
boxes = layout_result.boxes
|
| 493 |
+
if not boxes or not tokens:
|
| 494 |
+
logger.debug(
|
| 495 |
+
"align_tokens_to_layout: skipped (boxes=%d, tokens=%d)",
|
| 496 |
+
len(boxes) if boxes else 0,
|
| 497 |
+
len(tokens) if tokens else 0,
|
| 498 |
+
)
|
| 499 |
+
return []
|
| 500 |
+
|
| 501 |
+
logger.debug(
|
| 502 |
+
"align_tokens_to_layout: %d tokens -> %d layout boxes (containment=%.2f)",
|
| 503 |
+
len(tokens),
|
| 504 |
+
len(boxes),
|
| 505 |
+
containment_threshold,
|
| 506 |
+
)
|
| 507 |
+
|
| 508 |
+
wx0 = np.fromiter((t.x0 for t in tokens), dtype=np.float64, count=len(tokens))
|
| 509 |
+
wx1 = np.fromiter((t.x1 for t in tokens), dtype=np.float64, count=len(tokens))
|
| 510 |
+
wtop = np.fromiter((t.top for t in tokens), dtype=np.float64, count=len(tokens))
|
| 511 |
+
wbot = np.fromiter((t.bottom for t in tokens), dtype=np.float64, count=len(tokens))
|
| 512 |
+
word_area = np.clip((wx1 - wx0) * (wbot - wtop), 1e-6, None)
|
| 513 |
+
|
| 514 |
+
bx0 = np.array([b.bbox.xmin for b in boxes])
|
| 515 |
+
by0 = np.array([b.bbox.ymin for b in boxes])
|
| 516 |
+
bx1 = np.array([b.bbox.xmax for b in boxes])
|
| 517 |
+
by1 = np.array([b.bbox.ymax for b in boxes])
|
| 518 |
+
box_area = np.clip((bx1 - bx0) * (by1 - by0), 1e-6, None)
|
| 519 |
+
|
| 520 |
+
ix0 = np.maximum(wx0[:, None], bx0[None, :])
|
| 521 |
+
iy0 = np.maximum(wtop[:, None], by0[None, :])
|
| 522 |
+
ix1 = np.minimum(wx1[:, None], bx1[None, :])
|
| 523 |
+
iy1 = np.minimum(wbot[:, None], by1[None, :])
|
| 524 |
+
inter = np.clip(ix1 - ix0, 0, None) * np.clip(iy1 - iy0, 0, None)
|
| 525 |
+
containment = inter / word_area[:, None]
|
| 526 |
+
|
| 527 |
+
masked_area = np.where(containment >= containment_threshold, box_area[None, :], np.inf)
|
| 528 |
+
best = np.argmin(masked_area, axis=1)
|
| 529 |
+
matched = np.isfinite(masked_area[np.arange(len(tokens)), best])
|
| 530 |
+
if not matched.all():
|
| 531 |
+
bcx, bcy = (bx0 + bx1) / 2.0, (by0 + by1) / 2.0
|
| 532 |
+
wcx, wcy = (wx0 + wx1) / 2.0, (wtop + wbot) / 2.0
|
| 533 |
+
dist = (wcx[:, None] - bcx[None, :]) ** 2 + (wcy[:, None] - bcy[None, :]) ** 2
|
| 534 |
+
best = np.where(matched, best, np.argmin(dist, axis=1))
|
| 535 |
+
|
| 536 |
+
global_tol = float(np.median(wbot - wtop)) * line_tol_ratio
|
| 537 |
+
|
| 538 |
+
order = np.argsort(best, kind="stable")
|
| 539 |
+
sorted_idx = best[order]
|
| 540 |
+
split_points = np.searchsorted(sorted_idx, np.arange(len(boxes)))
|
| 541 |
+
split_points = np.append(split_points, len(order))
|
| 542 |
+
|
| 543 |
+
results = []
|
| 544 |
+
for box_idx, box in enumerate(boxes):
|
| 545 |
+
idxs = order[split_points[box_idx]:split_points[box_idx + 1]]
|
| 546 |
+
box_words = [tokens[i] for i in idxs]
|
| 547 |
+
lines = _line_cluster_fast(box_words, global_tol)
|
| 548 |
+
text = "\n".join(" ".join(w.text for w in line) for line in lines)
|
| 549 |
+
results.append(
|
| 550 |
+
RegionText(box.reading_order, box.class_name, box.bbox.as_tuple(), text, box_words)
|
| 551 |
+
)
|
| 552 |
+
logger.debug(
|
| 553 |
+
" matched box [%d] %s: %d words, text_len=%d",
|
| 554 |
+
box.reading_order,
|
| 555 |
+
box.class_name,
|
| 556 |
+
len(box_words),
|
| 557 |
+
len(text),
|
| 558 |
+
)
|
| 559 |
+
results.sort(key=lambda r: r.reading_order)
|
| 560 |
+
return results
|
| 561 |
+
|
| 562 |
+
|
| 563 |
+
def pdfplumber_tokens_in_image_space(
|
| 564 |
+
page, image_width: int, image_height: int
|
| 565 |
+
) -> List[PlacedWord]:
|
| 566 |
+
"""
|
| 567 |
+
Extract a pdfplumber page's words and scale them from PDF point-space
|
| 568 |
+
into the pixel space of a page image rendered at
|
| 569 |
+
(image_width, image_height) β e.g. via `page.to_image(resolution=...)`,
|
| 570 |
+
the same image PP-DocLayout ran on.
|
| 571 |
+
|
| 572 |
+
Returns `[]` if `extract_words()` finds nothing (a scanned page with no
|
| 573 |
+
text layer). Factored out of `align_words_to_layout` so callers that need
|
| 574 |
+
these same scaled tokens for something else (e.g. pulling exact words for
|
| 575 |
+
a table region instead of OCR-ing it) don't have to re-derive them.
|
| 576 |
+
"""
|
| 577 |
+
raw_words = page.extract_words(
|
| 578 |
+
x_tolerance=2,
|
| 579 |
+
y_tolerance=3,
|
| 580 |
+
keep_blank_chars=False,
|
| 581 |
+
use_text_flow=False,
|
| 582 |
+
)
|
| 583 |
+
if not raw_words:
|
| 584 |
+
logger.debug("pdfplumber extract_words: 0 words (scanned or image-only page)")
|
| 585 |
+
return []
|
| 586 |
+
|
| 587 |
+
scale_x = image_width / page.width
|
| 588 |
+
scale_y = image_height / page.height
|
| 589 |
+
logger.debug(
|
| 590 |
+
"pdfplumber extract_words: %d words scaled to %dx%d (scale %.3f, %.3f)",
|
| 591 |
+
len(raw_words),
|
| 592 |
+
image_width,
|
| 593 |
+
image_height,
|
| 594 |
+
scale_x,
|
| 595 |
+
scale_y,
|
| 596 |
+
)
|
| 597 |
+
return [
|
| 598 |
+
PlacedWord(
|
| 599 |
+
text=w["text"],
|
| 600 |
+
x0=w["x0"] * scale_x,
|
| 601 |
+
x1=w["x1"] * scale_x,
|
| 602 |
+
top=w["top"] * scale_y,
|
| 603 |
+
bottom=w["bottom"] * scale_y,
|
| 604 |
+
)
|
| 605 |
+
for w in raw_words
|
| 606 |
+
]
|
| 607 |
+
|
| 608 |
+
|
| 609 |
+
def align_words_to_layout(
|
| 610 |
+
layout_result: DocLayoutResult,
|
| 611 |
+
page,
|
| 612 |
+
containment_threshold: float = 0.5,
|
| 613 |
+
line_tol_ratio: float = 0.5,
|
| 614 |
+
) -> List[RegionText]:
|
| 615 |
+
"""
|
| 616 |
+
Plumber-path matcher ("Plumber_bboxes_words and pp-doclayout matching" box).
|
| 617 |
+
|
| 618 |
+
`page` is a pdfplumber Page. Its `extract_words()` output is in PDF point
|
| 619 |
+
space, so it's scaled up to the layout image's pixel space before matching.
|
| 620 |
+
"""
|
| 621 |
+
if not layout_result.boxes:
|
| 622 |
+
return []
|
| 623 |
+
|
| 624 |
+
tokens = pdfplumber_tokens_in_image_space(
|
| 625 |
+
page, layout_result.image_width, layout_result.image_height
|
| 626 |
+
)
|
| 627 |
+
if not tokens:
|
| 628 |
+
return []
|
| 629 |
+
return align_tokens_to_layout(tokens, layout_result, containment_threshold, line_tol_ratio)
|
| 630 |
+
|
| 631 |
+
|
| 632 |
+
def align_ocr_to_layout(
|
| 633 |
+
layout_result: DocLayoutResult,
|
| 634 |
+
ocr_tokens: Sequence[Dict],
|
| 635 |
+
containment_threshold: float = 0.5,
|
| 636 |
+
line_tol_ratio: float = 0.5,
|
| 637 |
+
) -> List[RegionText]:
|
| 638 |
+
"""
|
| 639 |
+
OCR-fallback matcher ("Matching ocr bboxes and layout" box) β used when
|
| 640 |
+
pdfplumber's `extract_words()` returns nothing (e.g. a scanned page).
|
| 641 |
+
|
| 642 |
+
`ocr_tokens` is the output of an `OCRBackend.get_text_boxes(image)` call
|
| 643 |
+
(see ocr_backends.py): a list of {"text": str, "bbox": [x1, y1, x2, y2]}
|
| 644 |
+
dicts already in the rendered image's pixel space β the same space the
|
| 645 |
+
layout image was detected in β so no scaling is required here.
|
| 646 |
+
"""
|
| 647 |
+
if not layout_result.boxes or not ocr_tokens:
|
| 648 |
+
logger.debug(
|
| 649 |
+
"align_ocr_to_layout: skipped (boxes=%d, ocr_tokens=%d)",
|
| 650 |
+
len(layout_result.boxes),
|
| 651 |
+
len(ocr_tokens),
|
| 652 |
+
)
|
| 653 |
+
return []
|
| 654 |
+
|
| 655 |
+
logger.info(
|
| 656 |
+
"OCR matching: %d OCR tokens -> %d layout boxes",
|
| 657 |
+
len(ocr_tokens),
|
| 658 |
+
len(layout_result.boxes),
|
| 659 |
+
)
|
| 660 |
+
|
| 661 |
+
tokens = [
|
| 662 |
+
PlacedWord(
|
| 663 |
+
text=t["text"],
|
| 664 |
+
x0=float(t["bbox"][0]),
|
| 665 |
+
x1=float(t["bbox"][2]),
|
| 666 |
+
top=float(t["bbox"][1]),
|
| 667 |
+
bottom=float(t["bbox"][3]),
|
| 668 |
+
)
|
| 669 |
+
for t in ocr_tokens
|
| 670 |
+
if t.get("text")
|
| 671 |
+
]
|
| 672 |
+
return align_tokens_to_layout(tokens, layout_result, containment_threshold, line_tol_ratio)
|
logging_config.py
ADDED
|
@@ -0,0 +1,44 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
logging_config.py
|
| 3 |
+
=================
|
| 4 |
+
|
| 5 |
+
Central logging setup for the pdf_pipeline package. Call ``setup_pipeline_logging``
|
| 6 |
+
once at process start (e.g. from ``example_usage.py``) to enable step-by-step
|
| 7 |
+
logs from pipeline, layout, OCR, and table extraction modules.
|
| 8 |
+
"""
|
| 9 |
+
|
| 10 |
+
from __future__ import annotations
|
| 11 |
+
|
| 12 |
+
import logging
|
| 13 |
+
from typing import Optional, Union
|
| 14 |
+
|
| 15 |
+
|
| 16 |
+
def setup_pipeline_logging(
|
| 17 |
+
level: Union[int, str] = logging.INFO,
|
| 18 |
+
log_file: Optional[str] = None,
|
| 19 |
+
) -> None:
|
| 20 |
+
"""
|
| 21 |
+
Configure root logging for the pipeline.
|
| 22 |
+
|
| 23 |
+
Parameters
|
| 24 |
+
----------
|
| 25 |
+
level:
|
| 26 |
+
Logging level (``logging.DEBUG``, ``logging.INFO``, etc.) or a level name
|
| 27 |
+
such as ``"DEBUG"``.
|
| 28 |
+
log_file:
|
| 29 |
+
Optional path to also write logs to a file.
|
| 30 |
+
"""
|
| 31 |
+
if isinstance(level, str):
|
| 32 |
+
level = getattr(logging, level.upper(), logging.INFO)
|
| 33 |
+
|
| 34 |
+
fmt = "%(asctime)s | %(levelname)-8s | %(name)s | %(message)s"
|
| 35 |
+
handlers: list[logging.Handler] = [logging.StreamHandler()]
|
| 36 |
+
if log_file:
|
| 37 |
+
handlers.append(logging.FileHandler(log_file, encoding="utf-8"))
|
| 38 |
+
|
| 39 |
+
logging.basicConfig(level=level, format=fmt, handlers=handlers, force=True)
|
| 40 |
+
|
| 41 |
+
# Keep noisy third-party libraries quieter unless DEBUG is requested.
|
| 42 |
+
if level > logging.DEBUG:
|
| 43 |
+
for name in ("PIL", "onnxruntime", "urllib3"):
|
| 44 |
+
logging.getLogger(name).setLevel(logging.WARNING)
|
ocr_backends.py
ADDED
|
@@ -0,0 +1,112 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
ocr_backends.py
|
| 3 |
+
================
|
| 4 |
+
|
| 5 |
+
OCR backends with a common interface: `get_text_boxes(image) -> list[dict]`,
|
| 6 |
+
each dict being `{"text": str, "bbox": [x1, y1, x2, y2]}` in the input
|
| 7 |
+
image's pixel space.
|
| 8 |
+
|
| 9 |
+
These backends are used in *two* places in the pipeline (diagram):
|
| 10 |
+
|
| 11 |
+
1. "Ocr(pytesseract, paddleocr, rapidocr)" β page-level OCR fallback when
|
| 12 |
+
pdfplumber finds no extractable words (scanned pages).
|
| 13 |
+
2. Table-cell OCR inside TableFormerONNX's `ocr_anchor_cells` step (see
|
| 14 |
+
table_extraction.py) β same backends, reused as-is.
|
| 15 |
+
|
| 16 |
+
Unchanged in substance from docling_.ipynb, just split out into its own
|
| 17 |
+
module and with the legacy crop-based `.read()` method dropped (it was a
|
| 18 |
+
no-op stub for every backend already β `get_text_boxes` was the only one
|
| 19 |
+
actually implemented and used).
|
| 20 |
+
"""
|
| 21 |
+
|
| 22 |
+
from __future__ import annotations
|
| 23 |
+
|
| 24 |
+
import logging
|
| 25 |
+
from abc import ABC, abstractmethod
|
| 26 |
+
|
| 27 |
+
import cv2
|
| 28 |
+
import numpy as np
|
| 29 |
+
from PIL import Image
|
| 30 |
+
import os
|
| 31 |
+
logger = logging.getLogger(__name__)
|
| 32 |
+
|
| 33 |
+
|
| 34 |
+
|
| 35 |
+
|
| 36 |
+
|
| 37 |
+
|
| 38 |
+
def _log_ocr_tokens(backend_name: str, image: Image.Image, tokens: list[dict]) -> None:
|
| 39 |
+
logger.debug("OCR (%s) on image %dx%d", backend_name, image.width, image.height)
|
| 40 |
+
logger.info("OCR (%s) returned %d text boxes", backend_name, len(tokens))
|
| 41 |
+
if logger.isEnabledFor(logging.DEBUG):
|
| 42 |
+
for t in tokens[:10]:
|
| 43 |
+
logger.debug(" %r @ %s", t["text"], t["bbox"])
|
| 44 |
+
|
| 45 |
+
|
| 46 |
+
class OCRBackend(ABC):
|
| 47 |
+
@abstractmethod
|
| 48 |
+
def get_text_boxes(self, image: Image.Image) -> list[dict]:
|
| 49 |
+
"""Run OCR on the full image once; return all text tokens with
|
| 50 |
+
their global pixel-space bounding boxes [x1, y1, x2, y2]."""
|
| 51 |
+
raise NotImplementedError
|
| 52 |
+
|
| 53 |
+
|
| 54 |
+
class RapidOCRBackend(OCRBackend):
|
| 55 |
+
def __init__(self,det_model_path,rec_model_path,rec_keys_path):
|
| 56 |
+
from rapidocr_onnxruntime import RapidOCR
|
| 57 |
+
|
| 58 |
+
self._engine = RapidOCR( det_model_path=det_model_path,
|
| 59 |
+
rec_model_path=rec_model_path,
|
| 60 |
+
rec_keys_path=rec_keys_path)
|
| 61 |
+
|
| 62 |
+
@staticmethod
|
| 63 |
+
def _to_bgr(image: Image.Image) -> np.ndarray:
|
| 64 |
+
rgb = np.array(image.convert("RGB"))
|
| 65 |
+
return cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)
|
| 66 |
+
|
| 67 |
+
def get_text_boxes(self, image: Image.Image) -> list[dict]:
|
| 68 |
+
result, _ = self._engine(self._to_bgr(image), use_det=True, use_cls=False)
|
| 69 |
+
tokens = []
|
| 70 |
+
if result:
|
| 71 |
+
for item in result:
|
| 72 |
+
pts = np.array(item[0]) # 4 corners
|
| 73 |
+
x1, y1 = pts.min(axis=0)
|
| 74 |
+
x2, y2 = pts.max(axis=0)
|
| 75 |
+
tokens.append(
|
| 76 |
+
{"text": str(item[1]).strip(), "bbox": [float(x1), float(y1), float(x2), float(y2)]}
|
| 77 |
+
)
|
| 78 |
+
_log_ocr_tokens("rapidocr", image, tokens)
|
| 79 |
+
return tokens
|
| 80 |
+
|
| 81 |
+
|
| 82 |
+
class PytesseractBackend(OCRBackend):
|
| 83 |
+
def __init__(self, lang: str = "fra+eng"):
|
| 84 |
+
import pytesseract
|
| 85 |
+
|
| 86 |
+
self._pt = pytesseract
|
| 87 |
+
self.lang = lang
|
| 88 |
+
|
| 89 |
+
def get_text_boxes(self, image: Image.Image) -> list[dict]:
|
| 90 |
+
data = self._pt.image_to_data(image, lang=self.lang, output_type=self._pt.Output.DICT)
|
| 91 |
+
tokens = []
|
| 92 |
+
for i in range(len(data["text"])):
|
| 93 |
+
text = data["text"][i].strip()
|
| 94 |
+
if text:
|
| 95 |
+
x = data["left"][i]
|
| 96 |
+
y = data["top"][i]
|
| 97 |
+
w = data["width"][i]
|
| 98 |
+
h = data["height"][i]
|
| 99 |
+
tokens.append({"text": text, "bbox": [float(x), float(y), float(x + w), float(y + h)]})
|
| 100 |
+
_log_ocr_tokens("pytesseract", image, tokens)
|
| 101 |
+
return tokens
|
| 102 |
+
|
| 103 |
+
|
| 104 |
+
|
| 105 |
+
|
| 106 |
+
def get_ocr_backend(name: str, **kwargs) -> OCRBackend:
|
| 107 |
+
if name == "rapidocr":
|
| 108 |
+
return RapidOCRBackend(**kwargs)
|
| 109 |
+
if name == "pytesseract":
|
| 110 |
+
return PytesseractBackend(**kwargs)
|
| 111 |
+
|
| 112 |
+
raise ValueError(f"unknown OCR backend: {name!r} (use 'rapidocr', 'pytesseract')")
|
pipeline.py
ADDED
|
@@ -0,0 +1,505 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
pipeline.py
|
| 3 |
+
===========
|
| 4 |
+
|
| 5 |
+
Wires together layout.py, ocr_backends.py, and table_extraction.py into the
|
| 6 |
+
exact flow drawn in the diagram:
|
| 7 |
+
|
| 8 |
+
pdf -> pdfplumber -> words+bboxes, image
|
| 9 |
+
| \\
|
| 10 |
+
if words empty? pp-doclayout -> layout boxes
|
| 11 |
+
yes / \\no | |
|
| 12 |
+
OCR Plumber+layout match | OCR+layout match
|
| 13 |
+
\\__________________________________/________/
|
| 14 |
+
|
|
| 15 |
+
(class_name, bbox, reading_order, text) per region
|
| 16 |
+
|
|
| 17 |
+
if class_name == "table" -> TableFormerONNX -> OTSL -> markdown
|
| 18 |
+
else -> markdown formatted by class_name
|
| 19 |
+
|
| 20 |
+
Two entry points sit on top of a shared per-image core (`_process_rendered_page`):
|
| 21 |
+
|
| 22 |
+
- `process_pdf_page` / `process_pdf` β PDF input via pdfplumber. May have
|
| 23 |
+
native words (Plumber-path matching) or not (OCR fallback).
|
| 24 |
+
- `process_image_page` / `process_images` β standalone image input
|
| 25 |
+
(.jpg/.png/etc., e.g. a photographed or scanned page with no PDF
|
| 26 |
+
structure at all). These never have native words, so they always go
|
| 27 |
+
through the OCR branch β the same code path a scanned PDF page uses.
|
| 28 |
+
|
| 29 |
+
`process_document` is a single dispatcher that picks the right one based on
|
| 30 |
+
the input (a .pdf path, an image path, or a list of image paths treated as
|
| 31 |
+
ordered pages of one document).
|
| 32 |
+
"""
|
| 33 |
+
|
| 34 |
+
from __future__ import annotations
|
| 35 |
+
|
| 36 |
+
import logging
|
| 37 |
+
import os
|
| 38 |
+
from pathlib import Path
|
| 39 |
+
from typing import List, Optional, Sequence, Union
|
| 40 |
+
import base64
|
| 41 |
+
import io
|
| 42 |
+
import cv2
|
| 43 |
+
import numpy as np
|
| 44 |
+
from PIL import Image
|
| 45 |
+
|
| 46 |
+
from .layout import (
|
| 47 |
+
DocLayoutV3,
|
| 48 |
+
PlacedWord,
|
| 49 |
+
RegionText,
|
| 50 |
+
align_ocr_to_layout,
|
| 51 |
+
align_tokens_to_layout,
|
| 52 |
+
log_layout_result,
|
| 53 |
+
pdfplumber_tokens_in_image_space,
|
| 54 |
+
)
|
| 55 |
+
from .ocr_backends import OCRBackend
|
| 56 |
+
from .table_extraction import TableFormerONNX, otsl_to_markdown, table_image_to_otsl
|
| 57 |
+
|
| 58 |
+
logger = logging.getLogger(__name__)
|
| 59 |
+
|
| 60 |
+
# Layout classes that contribute nothing to the output markdown.
|
| 61 |
+
SKIP_CLASSES = {"header", "footer", "number", "seal"}
|
| 62 |
+
|
| 63 |
+
# Layout classes rendered as a cropped image link rather than as text.
|
| 64 |
+
IMAGE_LIKE_CLASSES = {"chart", "header_image", "footer_image", "vision_footnote"}
|
| 65 |
+
|
| 66 |
+
# Extensions routed through the standalone-image path by process_document.
|
| 67 |
+
IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff", ".webp", ".gif"}
|
| 68 |
+
|
| 69 |
+
|
| 70 |
+
def _to_bgr_array(img: Image.Image) -> np.ndarray:
|
| 71 |
+
"""
|
| 72 |
+
Convert a PIL image to a BGR ndarray.
|
| 73 |
+
|
| 74 |
+
`DocLayoutV3._preprocess` always applies a BGR->RGB channel swap, which
|
| 75 |
+
is only correct if the array it receives is actually BGR (as
|
| 76 |
+
`cv2.imread` produces). PIL images are RGB, so feeding one in directly
|
| 77 |
+
β as the original page-rendering code did β silently double-swaps the
|
| 78 |
+
channels and degrades detection. Converting explicitly here keeps every
|
| 79 |
+
caller (PDF-rendered pages and standalone images alike) consistent with
|
| 80 |
+
what `DocLayoutV3` actually expects from an in-memory array.
|
| 81 |
+
"""
|
| 82 |
+
return cv2.cvtColor(np.array(img.convert("RGB")), cv2.COLOR_RGB2BGR)
|
| 83 |
+
|
| 84 |
+
|
| 85 |
+
def crop_and_save_image(image: Image.Image, bbox: tuple, index: int, doc_path: str) -> str:
|
| 86 |
+
"""Crop `bbox` (xyxy, pixel space) out of `image` and save it under
|
| 87 |
+
`<doc_stem>/images/<index>.png`, returning that path."""
|
| 88 |
+
stem = doc_path.split(".")[0]
|
| 89 |
+
images_dir = os.path.join(stem, "images")
|
| 90 |
+
os.makedirs(images_dir, exist_ok=True)
|
| 91 |
+
cropped = image.crop(bbox)
|
| 92 |
+
out_path = os.path.join(images_dir, f"{index}.png")
|
| 93 |
+
cropped.save(out_path)
|
| 94 |
+
return out_path
|
| 95 |
+
|
| 96 |
+
|
| 97 |
+
|
| 98 |
+
def crop_to_base64(image: Image.Image, bbox: tuple, fmt: str = "PNG") -> str:
|
| 99 |
+
"""Crop `bbox` (xyxy, pixel space) out of `image` and return a base64
|
| 100 |
+
data URI, ready to embed directly in markdown β no file written to disk."""
|
| 101 |
+
cropped = image.crop(bbox)
|
| 102 |
+
buffer = io.BytesIO()
|
| 103 |
+
cropped.save(buffer, format=fmt)
|
| 104 |
+
b64_data = base64.b64encode(buffer.getvalue()).decode("ascii")
|
| 105 |
+
return f"data:image/{fmt.lower()};base64,{b64_data}"
|
| 106 |
+
|
| 107 |
+
|
| 108 |
+
def _table_tokens_from_page(
|
| 109 |
+
bbox: tuple, page_tokens: Optional[List[PlacedWord]]
|
| 110 |
+
) -> List[dict]:
|
| 111 |
+
"""
|
| 112 |
+
Pull the pdfplumber word tokens whose centroid falls inside a table
|
| 113 |
+
region's `bbox` (xyxy, full-page image pixel space β the same space
|
| 114 |
+
`page_tokens` is already expressed in) and re-express them relative to
|
| 115 |
+
the cropped table image's own origin (top-left of `bbox`), which is the
|
| 116 |
+
coordinate convention `table_image_to_otsl`/`ocr_anchor_cells` expects.
|
| 117 |
+
|
| 118 |
+
Returns `[]` if there's no text layer (`page_tokens` is None/empty) or
|
| 119 |
+
no words happen to land inside this particular region β callers should
|
| 120 |
+
fall back to OCR in that case.
|
| 121 |
+
"""
|
| 122 |
+
if not page_tokens:
|
| 123 |
+
return []
|
| 124 |
+
xmin, ymin, xmax, ymax = bbox
|
| 125 |
+
local_tokens = []
|
| 126 |
+
for t in page_tokens:
|
| 127 |
+
cx, cy = (t.x0 + t.x1) / 2.0, (t.top + t.bottom) / 2.0
|
| 128 |
+
if xmin <= cx <= xmax and ymin <= cy <= ymax:
|
| 129 |
+
local_tokens.append(
|
| 130 |
+
{"text": t.text, "bbox": [t.x0 - xmin, t.top - ymin, t.x1 - xmin, t.bottom - ymin]}
|
| 131 |
+
)
|
| 132 |
+
return local_tokens
|
| 133 |
+
|
| 134 |
+
|
| 135 |
+
def region_to_markdown(
|
| 136 |
+
region: RegionText,
|
| 137 |
+
index: int,
|
| 138 |
+
img: Image.Image,
|
| 139 |
+
doc_path: str,
|
| 140 |
+
table_runner: Optional[TableFormerONNX] = None,
|
| 141 |
+
table_ocr_backend: Optional[OCRBackend] = None,
|
| 142 |
+
page_tokens: Optional[List[PlacedWord]] = None,
|
| 143 |
+
) -> str:
|
| 144 |
+
label = region.class_name
|
| 145 |
+
logger.debug("Rendering region %d: class=%s bbox=%s", index, label, region.bbox)
|
| 146 |
+
|
| 147 |
+
if label in SKIP_CLASSES:
|
| 148 |
+
logger.debug("Skipping region %d (%s)", index, label)
|
| 149 |
+
return ""
|
| 150 |
+
|
| 151 |
+
if label == "table":
|
| 152 |
+
if table_runner is not None:
|
| 153 |
+
table_tokens = _table_tokens_from_page(region.bbox, page_tokens)
|
| 154 |
+
tmp_path = None
|
| 155 |
+
try:
|
| 156 |
+
import tempfile
|
| 157 |
+
cropped = img.crop(region.bbox)
|
| 158 |
+
with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
|
| 159 |
+
cropped.save(tmp, format="PNG")
|
| 160 |
+
tmp_path = tmp.name
|
| 161 |
+
|
| 162 |
+
if table_tokens:
|
| 163 |
+
logger.info(
|
| 164 |
+
"Table region %d: using %d pdfplumber tokens (no OCR)",
|
| 165 |
+
index, len(table_tokens),
|
| 166 |
+
)
|
| 167 |
+
otsl = table_image_to_otsl(table_runner, tmp_path, tokens=table_tokens)
|
| 168 |
+
elif table_ocr_backend is not None:
|
| 169 |
+
logger.info(
|
| 170 |
+
"Table region %d: no pdfplumber tokens β OCR via %s",
|
| 171 |
+
index, type(table_ocr_backend).__name__,
|
| 172 |
+
)
|
| 173 |
+
otsl = table_image_to_otsl(table_runner, tmp_path, backend=table_ocr_backend)
|
| 174 |
+
else:
|
| 175 |
+
raise ValueError(
|
| 176 |
+
"no pdfplumber words found in this table region and no "
|
| 177 |
+
"table_ocr_backend was provided to fall back on"
|
| 178 |
+
)
|
| 179 |
+
table_md, df = otsl_to_markdown(otsl)
|
| 180 |
+
if df is not None:
|
| 181 |
+
logger.info(
|
| 182 |
+
"Table region %d: extracted %d rows x %d cols (Excel export disabled)",
|
| 183 |
+
index, df.shape[0], df.shape[1],
|
| 184 |
+
)
|
| 185 |
+
logger.debug("Table region %d dataframe:\n%s", index, df)
|
| 186 |
+
# Excel export intentionally disabled for now.
|
| 187 |
+
|
| 188 |
+
return f"{table_md}\n\n"
|
| 189 |
+
except Exception as exc:
|
| 190 |
+
logger.warning(
|
| 191 |
+
"Table region %d: extraction failed (%s) β falling back to image embed",
|
| 192 |
+
index, exc,
|
| 193 |
+
)
|
| 194 |
+
image_b64 = crop_to_base64(img, region.bbox)
|
| 195 |
+
return f"<!-- table extraction failed: {exc} -->\n\n\n"
|
| 196 |
+
finally:
|
| 197 |
+
if tmp_path and os.path.exists(tmp_path):
|
| 198 |
+
os.remove(tmp_path)
|
| 199 |
+
logger.info("Table region %d: no TableFormer configured β embedding crop", index)
|
| 200 |
+
image_b64 = crop_to_base64(img, region.bbox)
|
| 201 |
+
return f"\n\n"
|
| 202 |
+
|
| 203 |
+
if label in IMAGE_LIKE_CLASSES:
|
| 204 |
+
image_b64 = crop_to_base64(img, region.bbox)
|
| 205 |
+
logger.info("Image-like region %d (%s): embedded as base64", index, label)
|
| 206 |
+
return f"\n\n"
|
| 207 |
+
|
| 208 |
+
text = region.text
|
| 209 |
+
if not text:
|
| 210 |
+
return ""
|
| 211 |
+
|
| 212 |
+
if label == "doc_title":
|
| 213 |
+
return f"# {text}\n\n"
|
| 214 |
+
if label == "paragraph_title":
|
| 215 |
+
return f"## {text}\n\n"
|
| 216 |
+
if label == "aside_text":
|
| 217 |
+
return f"> {text}\n\n"
|
| 218 |
+
if label == "figure_title":
|
| 219 |
+
return f"### {text}\n\n"
|
| 220 |
+
if label == "image":
|
| 221 |
+
image_b64 = crop_to_base64(img, region.bbox)
|
| 222 |
+
output = f"\n\n"
|
| 223 |
+
if text.strip():
|
| 224 |
+
output += f"{text}\n\n"
|
| 225 |
+
return output
|
| 226 |
+
|
| 227 |
+
return f"{text}\n\n"
|
| 228 |
+
|
| 229 |
+
def regions_to_markdown(
|
| 230 |
+
regions: List[RegionText],
|
| 231 |
+
img: Image.Image,
|
| 232 |
+
doc_path: str,
|
| 233 |
+
table_runner: Optional[TableFormerONNX] = None,
|
| 234 |
+
table_ocr_backend: Optional[OCRBackend] = None,
|
| 235 |
+
page_tokens: Optional[List[PlacedWord]] = None,
|
| 236 |
+
) -> str:
|
| 237 |
+
"""Render a full list of reading-order-sorted regions to one markdown string."""
|
| 238 |
+
parts = []
|
| 239 |
+
for i, region in enumerate(regions):
|
| 240 |
+
parts.append(
|
| 241 |
+
region_to_markdown(region, i, img, doc_path, table_runner, table_ocr_backend, page_tokens)
|
| 242 |
+
)
|
| 243 |
+
return "".join(parts)
|
| 244 |
+
|
| 245 |
+
|
| 246 |
+
def _process_rendered_page(
|
| 247 |
+
img: Image.Image,
|
| 248 |
+
layout_detector: DocLayoutV3,
|
| 249 |
+
doc_path: str,
|
| 250 |
+
page_index: int,
|
| 251 |
+
page_tokens: Optional[List[PlacedWord]],
|
| 252 |
+
page_ocr_backend: Optional[OCRBackend],
|
| 253 |
+
table_runner: Optional[TableFormerONNX],
|
| 254 |
+
table_ocr_backend: Optional[OCRBackend],
|
| 255 |
+
) -> dict:
|
| 256 |
+
"""
|
| 257 |
+
Shared core of the diagram's bottom half, given an already-loaded page
|
| 258 |
+
image: run layout detection, then route to either the pdfplumber-words
|
| 259 |
+
matcher or the OCR-fallback matcher, then render per-region markdown.
|
| 260 |
+
|
| 261 |
+
`page_tokens=None` (or `[]`) always forces the OCR path β this is what
|
| 262 |
+
a standalone image (no native text layer at all) and a scanned PDF page
|
| 263 |
+
(pdfplumber found no words) both end up doing identically.
|
| 264 |
+
"""
|
| 265 |
+
logger.info("Page %d: running layout detection on %s", page_index, doc_path)
|
| 266 |
+
layout_result = layout_detector.predict(_to_bgr_array(img))
|
| 267 |
+
log_layout_result(layout_result)
|
| 268 |
+
|
| 269 |
+
used_ocr = False
|
| 270 |
+
if page_tokens:
|
| 271 |
+
logger.info(
|
| 272 |
+
"Page %d: matching path pdfplumber (%d tokens)",
|
| 273 |
+
page_index,
|
| 274 |
+
len(page_tokens),
|
| 275 |
+
)
|
| 276 |
+
if logger.isEnabledFor(logging.DEBUG):
|
| 277 |
+
for t in page_tokens[:5]:
|
| 278 |
+
logger.debug(" sample token: %r @ (%.0f, %.0f)", t.text, t.x0, t.top)
|
| 279 |
+
regions = align_tokens_to_layout(page_tokens, layout_result)
|
| 280 |
+
else:
|
| 281 |
+
used_ocr = True
|
| 282 |
+
if page_ocr_backend is None:
|
| 283 |
+
raise ValueError(
|
| 284 |
+
"No native text/tokens available (scanned page or standalone "
|
| 285 |
+
"image input) and no page_ocr_backend was provided to fall back on."
|
| 286 |
+
)
|
| 287 |
+
logger.warning(
|
| 288 |
+
"Page %d: no text layer β OCR via %s",
|
| 289 |
+
page_index,
|
| 290 |
+
type(page_ocr_backend).__name__,
|
| 291 |
+
)
|
| 292 |
+
ocr_tokens = page_ocr_backend.get_text_boxes(img)
|
| 293 |
+
regions = align_ocr_to_layout(layout_result, ocr_tokens)
|
| 294 |
+
|
| 295 |
+
logger.info("Page %d: matched %d regions (used_ocr=%s)", page_index, len(regions), used_ocr)
|
| 296 |
+
for r in regions:
|
| 297 |
+
preview = (r.text[:80] + "...") if len(r.text) > 80 else r.text
|
| 298 |
+
logger.info(
|
| 299 |
+
" region [%d] %s words=%d text=%r",
|
| 300 |
+
r.reading_order,
|
| 301 |
+
r.class_name,
|
| 302 |
+
len(r.words),
|
| 303 |
+
preview,
|
| 304 |
+
)
|
| 305 |
+
|
| 306 |
+
markdown_doc = regions_to_markdown(
|
| 307 |
+
regions, img, doc_path, table_runner, table_ocr_backend, page_tokens=page_tokens or None
|
| 308 |
+
)
|
| 309 |
+
|
| 310 |
+
return {
|
| 311 |
+
"page_index": page_index,
|
| 312 |
+
"used_ocr": used_ocr,
|
| 313 |
+
"layout_result": layout_result,
|
| 314 |
+
"regions": regions,
|
| 315 |
+
"markdown": markdown_doc,
|
| 316 |
+
"image": img,
|
| 317 |
+
}
|
| 318 |
+
|
| 319 |
+
|
| 320 |
+
def process_pdf_page(
|
| 321 |
+
page,
|
| 322 |
+
layout_detector: DocLayoutV3,
|
| 323 |
+
doc_path: str,
|
| 324 |
+
page_index: int = 0,
|
| 325 |
+
resolution: int = 150,
|
| 326 |
+
page_ocr_backend: Optional[OCRBackend] = None,
|
| 327 |
+
table_runner: Optional[TableFormerONNX] = None,
|
| 328 |
+
table_ocr_backend: Optional[OCRBackend] = None,
|
| 329 |
+
) -> dict:
|
| 330 |
+
"""
|
| 331 |
+
Run the full diagram for a single pdfplumber Page.
|
| 332 |
+
|
| 333 |
+
- page_ocr_backend: used for the "Ocr(...)" box, only invoked if
|
| 334 |
+
page.extract_words() comes back empty (scanned page / no text layer).
|
| 335 |
+
- table_runner / table_ocr_backend: used for the "TableFormerONNX" box,
|
| 336 |
+
only invoked for regions whose class_name == "table", and only as a
|
| 337 |
+
fallback when no pdfplumber words land inside that region. If neither
|
| 338 |
+
is usable, table regions fall back to an embedded cropped image instead.
|
| 339 |
+
|
| 340 |
+
Returns a dict with the page's markdown plus the intermediate regions,
|
| 341 |
+
so callers can inspect/debug a single page without re-running detection.
|
| 342 |
+
"""
|
| 343 |
+
page_image = page.to_image(resolution=resolution)
|
| 344 |
+
img = page_image.original.copy()
|
| 345 |
+
|
| 346 |
+
# Computed once: reused both for matching regions (Plumber-path) and,
|
| 347 |
+
# later, for pulling exact words inside any "table" region instead of
|
| 348 |
+
# re-OCR-ing it. `[]` here means "no text layer" (e.g. a scanned page).
|
| 349 |
+
img_w, img_h = img.size
|
| 350 |
+
logger.info(
|
| 351 |
+
"Processing PDF page %d of %s (resolution=%d, image=%dx%d)",
|
| 352 |
+
page_index,
|
| 353 |
+
doc_path,
|
| 354 |
+
resolution,
|
| 355 |
+
img_w,
|
| 356 |
+
img_h,
|
| 357 |
+
)
|
| 358 |
+
page_tokens = pdfplumber_tokens_in_image_space(page, img_w, img_h)
|
| 359 |
+
logger.info("Page %d: pdfplumber extracted %d word tokens", page_index, len(page_tokens))
|
| 360 |
+
|
| 361 |
+
return _process_rendered_page(
|
| 362 |
+
img, layout_detector, doc_path, page_index, page_tokens,
|
| 363 |
+
page_ocr_backend, table_runner, table_ocr_backend,
|
| 364 |
+
)
|
| 365 |
+
|
| 366 |
+
|
| 367 |
+
def process_pdf(
|
| 368 |
+
doc_path: str,
|
| 369 |
+
layout_detector: DocLayoutV3,
|
| 370 |
+
pages: Optional[List[int]] = None,
|
| 371 |
+
resolution: int = 150,
|
| 372 |
+
page_ocr_backend: Optional[OCRBackend] = None,
|
| 373 |
+
table_runner: Optional[TableFormerONNX] = None,
|
| 374 |
+
table_ocr_backend: Optional[OCRBackend] = None,
|
| 375 |
+
) -> str:
|
| 376 |
+
"""
|
| 377 |
+
Run the full diagram over an entire PDF (or a subset of page indices)
|
| 378 |
+
and return one concatenated markdown document.
|
| 379 |
+
"""
|
| 380 |
+
import pdfplumber
|
| 381 |
+
|
| 382 |
+
logger.info("Processing PDF: %s (pages=%s)", doc_path, pages if pages is not None else "all")
|
| 383 |
+
markdown_chunks = []
|
| 384 |
+
with pdfplumber.open(doc_path) as pdf:
|
| 385 |
+
page_indices = pages if pages is not None else range(len(pdf.pages))
|
| 386 |
+
for i in page_indices:
|
| 387 |
+
page = pdf.pages[i]
|
| 388 |
+
result = process_pdf_page(
|
| 389 |
+
page,
|
| 390 |
+
layout_detector,
|
| 391 |
+
doc_path,
|
| 392 |
+
page_index=i,
|
| 393 |
+
resolution=resolution,
|
| 394 |
+
page_ocr_backend=page_ocr_backend,
|
| 395 |
+
table_runner=table_runner,
|
| 396 |
+
table_ocr_backend=table_ocr_backend,
|
| 397 |
+
)
|
| 398 |
+
markdown_chunks.append(result["markdown"])
|
| 399 |
+
|
| 400 |
+
logger.info("Finished PDF: %s (%d pages)", doc_path, len(markdown_chunks))
|
| 401 |
+
return "\n".join(markdown_chunks)
|
| 402 |
+
|
| 403 |
+
|
| 404 |
+
def process_image_page(
|
| 405 |
+
image_path: Union[str, Path],
|
| 406 |
+
layout_detector: DocLayoutV3,
|
| 407 |
+
doc_path: Optional[str] = None,
|
| 408 |
+
page_index: int = 0,
|
| 409 |
+
page_ocr_backend: Optional[OCRBackend] = None,
|
| 410 |
+
table_runner: Optional[TableFormerONNX] = None,
|
| 411 |
+
table_ocr_backend: Optional[OCRBackend] = None,
|
| 412 |
+
) -> dict:
|
| 413 |
+
"""
|
| 414 |
+
Run the diagram on a single standalone image (a photo or scan, not a
|
| 415 |
+
PDF page). There's no PDF text layer to check, so this always takes the
|
| 416 |
+
OCR branch β equivalent to a PDF page where `extract_words()` came back
|
| 417 |
+
empty. `page_ocr_backend` is therefore required, not optional, here.
|
| 418 |
+
|
| 419 |
+
`doc_path` controls where cropped table/image regions get saved
|
| 420 |
+
(`<doc_stem>/images/<n>.png`); defaults to `image_path` itself.
|
| 421 |
+
"""
|
| 422 |
+
if page_ocr_backend is None:
|
| 423 |
+
raise ValueError("process_image_page requires page_ocr_backend β standalone images have no text layer.")
|
| 424 |
+
|
| 425 |
+
logger.info("Processing standalone image page %d: %s", page_index, image_path)
|
| 426 |
+
img = Image.open(image_path).convert("RGB")
|
| 427 |
+
doc_path = doc_path or str(image_path)
|
| 428 |
+
|
| 429 |
+
return _process_rendered_page(
|
| 430 |
+
img, layout_detector, doc_path, page_index, None,
|
| 431 |
+
page_ocr_backend, table_runner, table_ocr_backend,
|
| 432 |
+
)
|
| 433 |
+
|
| 434 |
+
|
| 435 |
+
def process_images(
|
| 436 |
+
image_paths: Sequence[Union[str, Path]],
|
| 437 |
+
layout_detector: DocLayoutV3,
|
| 438 |
+
page_ocr_backend: Optional[OCRBackend] = None,
|
| 439 |
+
table_runner: Optional[TableFormerONNX] = None,
|
| 440 |
+
table_ocr_backend: Optional[OCRBackend] = None,
|
| 441 |
+
) -> str:
|
| 442 |
+
"""
|
| 443 |
+
Run the diagram over a batch of standalone images (e.g. photographed
|
| 444 |
+
pages of one physical document) in the given order, and return one
|
| 445 |
+
concatenated markdown document β the image-input analog of `process_pdf`.
|
| 446 |
+
"""
|
| 447 |
+
logger.info("Processing %d standalone images", len(image_paths))
|
| 448 |
+
chunks = []
|
| 449 |
+
for i, path in enumerate(image_paths):
|
| 450 |
+
result = process_image_page(
|
| 451 |
+
path, layout_detector, doc_path=str(path), page_index=i,
|
| 452 |
+
page_ocr_backend=page_ocr_backend, table_runner=table_runner,
|
| 453 |
+
table_ocr_backend=table_ocr_backend,
|
| 454 |
+
)
|
| 455 |
+
chunks.append(result["markdown"])
|
| 456 |
+
logger.info("Finished image batch (%d pages)", len(chunks))
|
| 457 |
+
return "\n".join(chunks)
|
| 458 |
+
|
| 459 |
+
|
| 460 |
+
def process_document(
|
| 461 |
+
path: Union[str, Path, Sequence[Union[str, Path]]],
|
| 462 |
+
layout_detector: DocLayoutV3,
|
| 463 |
+
pages: Optional[List[int]] = None,
|
| 464 |
+
resolution: int = 150,
|
| 465 |
+
page_ocr_backend: Optional[OCRBackend] = None,
|
| 466 |
+
table_runner: Optional[TableFormerONNX] = None,
|
| 467 |
+
table_ocr_backend: Optional[OCRBackend] = None,
|
| 468 |
+
) -> str:
|
| 469 |
+
"""
|
| 470 |
+
Single entry point for the whole pipeline β dispatches to the PDF path
|
| 471 |
+
or the image path based on the input, so callers don't need to care
|
| 472 |
+
which kind of file they have:
|
| 473 |
+
|
| 474 |
+
- a path to a `.pdf` -> `process_pdf`
|
| 475 |
+
- a path to a single image (.jpg/.png/etc.) -> `process_image_page`
|
| 476 |
+
- a list of image paths -> `process_images` (treated as ordered pages
|
| 477 |
+
of one document)
|
| 478 |
+
"""
|
| 479 |
+
if isinstance(path, (list, tuple)):
|
| 480 |
+
logger.info("process_document: dispatching to process_images (%d paths)", len(path))
|
| 481 |
+
return process_images(
|
| 482 |
+
path, layout_detector, page_ocr_backend=page_ocr_backend,
|
| 483 |
+
table_runner=table_runner, table_ocr_backend=table_ocr_backend,
|
| 484 |
+
)
|
| 485 |
+
|
| 486 |
+
suffix = Path(path).suffix.lower()
|
| 487 |
+
if suffix == ".pdf":
|
| 488 |
+
logger.info("process_document: dispatching to process_pdf (%s)", path)
|
| 489 |
+
return process_pdf(
|
| 490 |
+
str(path), layout_detector, pages=pages, resolution=resolution,
|
| 491 |
+
page_ocr_backend=page_ocr_backend, table_runner=table_runner,
|
| 492 |
+
table_ocr_backend=table_ocr_backend,
|
| 493 |
+
)
|
| 494 |
+
if suffix in IMAGE_EXTENSIONS:
|
| 495 |
+
logger.info("process_document: dispatching to process_image_page (%s)", path)
|
| 496 |
+
result = process_image_page(
|
| 497 |
+
path, layout_detector, page_ocr_backend=page_ocr_backend,
|
| 498 |
+
table_runner=table_runner, table_ocr_backend=table_ocr_backend,
|
| 499 |
+
)
|
| 500 |
+
return result["markdown"]
|
| 501 |
+
|
| 502 |
+
raise ValueError(
|
| 503 |
+
f"Unsupported file type {suffix!r} for {path!r}; expected a .pdf or "
|
| 504 |
+
f"an image ({sorted(IMAGE_EXTENSIONS)})."
|
| 505 |
+
)
|
table_extraction.py
ADDED
|
@@ -0,0 +1,565 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
table_extraction.py
|
| 3 |
+
====================
|
| 4 |
+
|
| 5 |
+
Right-hand branch of the diagram: when a matched region's class is "table",
|
| 6 |
+
its cropped image goes through TableFormerONNX -> OTSL string -> markdown.
|
| 7 |
+
|
| 8 |
+
`TableFormerONNX`, `ocr_anchor_cells`, `seq_to_otsl`, and `table_image_to_otsl`
|
| 9 |
+
are carried over from docling_.ipynb essentially unchanged (decoupled from
|
| 10 |
+
that notebook's specific OCR-backend import style; backends now come from
|
| 11 |
+
ocr_backends.py). `otsl_to_markdown` is new β it's the "otsl -> markdown" step
|
| 12 |
+
in the diagram, which in the original notebook was done ad hoc in later
|
| 13 |
+
cells (build a <doctag>, run DocTagsDocument/DoclingDocument). Folding it
|
| 14 |
+
into one function makes it reusable per-table inside the full-page pipeline.
|
| 15 |
+
It returns `(markdown, dataframe)` β the dataframe is built from
|
| 16 |
+
`TableItem.export_to_dataframe()` with `collapse_spanned_columns` applied,
|
| 17 |
+
since a raw rectangular DataFrame otherwise duplicates a colspan'd cell's
|
| 18 |
+
value into every column it visually spans (see that function's docstring
|
| 19 |
+
for why, and why `pd.read_html(export_to_html())` has the identical issue).
|
| 20 |
+
"""
|
| 21 |
+
|
| 22 |
+
from __future__ import annotations
|
| 23 |
+
|
| 24 |
+
import json
|
| 25 |
+
import logging
|
| 26 |
+
import time
|
| 27 |
+
from pathlib import Path
|
| 28 |
+
from typing import Optional, Union
|
| 29 |
+
from urllib.parse import urlparse
|
| 30 |
+
|
| 31 |
+
import cv2
|
| 32 |
+
import numpy as np
|
| 33 |
+
import onnxruntime as ort
|
| 34 |
+
import pandas as pd
|
| 35 |
+
from PIL import Image
|
| 36 |
+
|
| 37 |
+
from .ocr_backends import OCRBackend
|
| 38 |
+
|
| 39 |
+
logger = logging.getLogger(__name__)
|
| 40 |
+
|
| 41 |
+
ANCHOR_TAGS = {"fcel", "ecel", "ched", "rhed", "srow"} # tags that occupy their own grid slot
|
| 42 |
+
EMPTY_OTSL_PLACEHOLDER = "?" # inserted when a non-empty-tagged cell OCRs to nothing,
|
| 43 |
+
# so the docling_core parser never sees an empty content cell
|
| 44 |
+
|
| 45 |
+
MEAN = np.array([0.94247851, 0.94254675, 0.94292611], dtype=np.float32)
|
| 46 |
+
STD = np.array([0.17910956, 0.17940403, 0.17931663], dtype=np.float32)
|
| 47 |
+
IMAGE_SIZE = 448
|
| 48 |
+
|
| 49 |
+
|
| 50 |
+
# --------------------------------------------------------------------------- #
|
| 51 |
+
# bbox merge helpers (lcel chains)
|
| 52 |
+
# --------------------------------------------------------------------------- #
|
| 53 |
+
def merge_bboxes(box1: np.ndarray, box2: np.ndarray) -> np.ndarray:
|
| 54 |
+
"""Union the cxcywh boxes of an lcel chain's stub anchor and the real cell that closes it."""
|
| 55 |
+
new_w = (box2[0] + box2[2] / 2) - (box1[0] - box1[2] / 2)
|
| 56 |
+
new_h = (box2[1] + box2[3] / 2) - (box1[1] - box1[3] / 2)
|
| 57 |
+
new_left = box1[0] - box1[2] / 2
|
| 58 |
+
new_top = min((box2[1] - box2[3] / 2), (box1[1] - box1[3] / 2))
|
| 59 |
+
return np.array(
|
| 60 |
+
[new_left + new_w / 2, new_top + new_h / 2, new_w, new_h], dtype=np.float32
|
| 61 |
+
)
|
| 62 |
+
|
| 63 |
+
|
| 64 |
+
def apply_bbox_merge(coords: np.ndarray, bboxes_to_merge: dict) -> np.ndarray:
|
| 65 |
+
coords = [np.asarray(c, dtype=np.float32) for c in coords]
|
| 66 |
+
merged, skip = [], set()
|
| 67 |
+
for i, box1 in enumerate(coords):
|
| 68 |
+
if i in bboxes_to_merge:
|
| 69 |
+
j = bboxes_to_merge[i]
|
| 70 |
+
if j >= 0:
|
| 71 |
+
skip.add(j)
|
| 72 |
+
merged.append(merge_bboxes(box1, coords[j]))
|
| 73 |
+
elif i not in skip:
|
| 74 |
+
merged.append(box1)
|
| 75 |
+
return np.stack(merged) if merged else np.empty((0, 4), dtype=np.float32)
|
| 76 |
+
|
| 77 |
+
|
| 78 |
+
def make_session(path: str, threads: int = 2) -> ort.InferenceSession:
|
| 79 |
+
opts = ort.SessionOptions()
|
| 80 |
+
opts.inter_op_num_threads = 2
|
| 81 |
+
opts.intra_op_num_threads = 4
|
| 82 |
+
opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
|
| 83 |
+
opts.enable_mem_pattern = True
|
| 84 |
+
opts.enable_mem_reuse = True
|
| 85 |
+
opts.enable_profiling = False
|
| 86 |
+
return ort.InferenceSession(path, sess_options=opts, providers=["CPUExecutionProvider"])
|
| 87 |
+
|
| 88 |
+
|
| 89 |
+
# --------------------------------------------------------------------------- #
|
| 90 |
+
# TableFormer ONNX runner
|
| 91 |
+
# --------------------------------------------------------------------------- #
|
| 92 |
+
class TableFormerONNX:
|
| 93 |
+
"""
|
| 94 |
+
encoder + decoder-step (autoregressive loop, KV-cache) + bbox_decoder,
|
| 95 |
+
all ONNX Runtime. Produces a structure sequence (OTSL tags) plus a
|
| 96 |
+
normalized cxcywh bbox per anchor cell.
|
| 97 |
+
|
| 98 |
+
Usage:
|
| 99 |
+
runner = TableFormerONNX(artifact_root="tableformerv1", variant="accurate")
|
| 100 |
+
"""
|
| 101 |
+
|
| 102 |
+
def __init__(self, artifact_root: str = "tableformerv1", variant: str = "accurate", threads: int = 2):
|
| 103 |
+
root = Path(artifact_root)
|
| 104 |
+
onnx_dir = root / "onnx" / variant
|
| 105 |
+
self.enc = make_session(str(onnx_dir / f"tableformer_{variant}_encoder.onnx"), threads)
|
| 106 |
+
self.dec = make_session(str(onnx_dir / f"tableformer_{variant}_decoder_step.onnx"), threads)
|
| 107 |
+
self.bbox = make_session(str(onnx_dir / f"tableformer_{variant}_bbox_decoder.onnx"), threads)
|
| 108 |
+
|
| 109 |
+
with open(root / "tm_config.json", encoding="utf-8") as f:
|
| 110 |
+
self.tag_map = json.load(f)["dataset_wordmap"]["word_map_tag"]
|
| 111 |
+
self.rev_tag = {v: k for k, v in self.tag_map.items()}
|
| 112 |
+
|
| 113 |
+
cache_input = next(i for i in self.dec.get_inputs() if i.name == "cache")
|
| 114 |
+
self.num_layers = int(cache_input.shape[0])
|
| 115 |
+
self.embed_dim = int(cache_input.shape[3])
|
| 116 |
+
|
| 117 |
+
@staticmethod
|
| 118 |
+
def is_valid_url(url_string: str) -> bool:
|
| 119 |
+
try:
|
| 120 |
+
result = urlparse(url_string)
|
| 121 |
+
return all([result.scheme in ("http", "https"), result.netloc])
|
| 122 |
+
except ValueError:
|
| 123 |
+
return False
|
| 124 |
+
|
| 125 |
+
def image_preprocess(self, img) -> np.ndarray:
|
| 126 |
+
"""Accepts a path, URL, PIL.Image, or HWC RGB ndarray."""
|
| 127 |
+
if isinstance(img, str) and self.is_valid_url(img):
|
| 128 |
+
import requests
|
| 129 |
+
|
| 130 |
+
img = Image.open(requests.get(img, stream=True, timeout=30).raw)
|
| 131 |
+
if isinstance(img, (str, Path)):
|
| 132 |
+
img = Image.open(img)
|
| 133 |
+
if isinstance(img, Image.Image):
|
| 134 |
+
img = np.array(img.convert("RGB"))
|
| 135 |
+
elif isinstance(img, np.ndarray) and img.ndim == 3 and img.shape[2] == 3:
|
| 136 |
+
pass
|
| 137 |
+
else:
|
| 138 |
+
raise TypeError(f"Unsupported image type: {type(img)}")
|
| 139 |
+
|
| 140 |
+
img = (img.astype(np.float32) - 255.0 * MEAN) / STD
|
| 141 |
+
img = cv2.resize(img, (IMAGE_SIZE, IMAGE_SIZE), interpolation=cv2.INTER_CUBIC)
|
| 142 |
+
img = img.transpose(2, 1, 0) / 255.0
|
| 143 |
+
return np.expand_dims(img.astype(np.float32), axis=0)
|
| 144 |
+
|
| 145 |
+
def draw_bbox(self, img_path: str, bboxs: np.ndarray, out_path: str = "draw_p.png") -> None:
|
| 146 |
+
"""Debug helper β draws predicted cxcywh boxes on the ORIGINAL-resolution image."""
|
| 147 |
+
from PIL import ImageDraw
|
| 148 |
+
|
| 149 |
+
pil_image = Image.open(img_path)
|
| 150 |
+
width, height = pil_image.size
|
| 151 |
+
draw = ImageDraw.Draw(pil_image)
|
| 152 |
+
for cx, cy, w, h in np.asarray(bboxs):
|
| 153 |
+
xyxy = (
|
| 154 |
+
int((cx - 0.5 * w) * width), int((cy - 0.5 * h) * height),
|
| 155 |
+
int((cx + 0.5 * w) * width), int((cy + 0.5 * h) * height),
|
| 156 |
+
)
|
| 157 |
+
draw.rectangle(xyxy, outline="red", width=2)
|
| 158 |
+
pil_image.save(out_path)
|
| 159 |
+
|
| 160 |
+
def predict(self, img: np.ndarray, max_steps: int = 1024) -> dict:
|
| 161 |
+
"""img: float32 [1, 3, 448, 448] (preprocessed). Returns seq / outputs_class /
|
| 162 |
+
outputs_coord (cxcywh, normalized [0,1]) / bboxes_to_merge / timings."""
|
| 163 |
+
t0 = time.time()
|
| 164 |
+
enc_out, memory = self.enc.run(None, {"image": img})
|
| 165 |
+
t_enc = time.time() - t0
|
| 166 |
+
|
| 167 |
+
wm = self.tag_map
|
| 168 |
+
decoded_tags = np.array([[wm["<start>"]]], dtype=np.int64)
|
| 169 |
+
cache = np.zeros((self.num_layers, 0, 1, self.embed_dim), dtype=np.float32)
|
| 170 |
+
skip_next_tag = True
|
| 171 |
+
prev_tag_ucel = False
|
| 172 |
+
line_num = 0
|
| 173 |
+
first_lcel = True
|
| 174 |
+
bbox_ind = 0
|
| 175 |
+
cur_bbox_ind = -1
|
| 176 |
+
bboxes_to_merge: dict = {}
|
| 177 |
+
tag_H_buf = []
|
| 178 |
+
|
| 179 |
+
t1 = time.time()
|
| 180 |
+
for _ in range(max_steps):
|
| 181 |
+
logits, last_hidden, cache = self.dec.run(
|
| 182 |
+
None, {"decoded_tags": decoded_tags, "memory": memory, "cache": cache}
|
| 183 |
+
)
|
| 184 |
+
new_tag = int(np.argmax(logits, axis=1)[0])
|
| 185 |
+
|
| 186 |
+
if line_num == 0 and new_tag == wm["xcel"]:
|
| 187 |
+
new_tag = wm["lcel"]
|
| 188 |
+
if prev_tag_ucel and new_tag == wm["lcel"]:
|
| 189 |
+
new_tag = wm["fcel"]
|
| 190 |
+
|
| 191 |
+
if new_tag == wm["<end>"]:
|
| 192 |
+
decoded_tags = np.concatenate(
|
| 193 |
+
[decoded_tags, np.array([[new_tag]], dtype=np.int64)], axis=0
|
| 194 |
+
)
|
| 195 |
+
break
|
| 196 |
+
|
| 197 |
+
if not skip_next_tag:
|
| 198 |
+
if new_tag in (wm["fcel"], wm["ecel"], wm["ched"], wm["rhed"],
|
| 199 |
+
wm["srow"], wm["nl"], wm["ucel"]):
|
| 200 |
+
tag_H_buf.append(last_hidden[:, 0, :].copy())
|
| 201 |
+
if first_lcel is not True:
|
| 202 |
+
bboxes_to_merge[cur_bbox_ind] = bbox_ind
|
| 203 |
+
bbox_ind += 1
|
| 204 |
+
|
| 205 |
+
if new_tag != wm["lcel"]:
|
| 206 |
+
first_lcel = True
|
| 207 |
+
else:
|
| 208 |
+
if first_lcel:
|
| 209 |
+
tag_H_buf.append(last_hidden[:, 0, :].copy())
|
| 210 |
+
first_lcel = False
|
| 211 |
+
cur_bbox_ind = bbox_ind
|
| 212 |
+
bboxes_to_merge[cur_bbox_ind] = -1
|
| 213 |
+
bbox_ind += 1
|
| 214 |
+
|
| 215 |
+
skip_next_tag = new_tag in (wm["nl"], wm["ucel"], wm["xcel"])
|
| 216 |
+
prev_tag_ucel = new_tag == wm["ucel"]
|
| 217 |
+
if new_tag == wm["nl"]:
|
| 218 |
+
line_num += 1
|
| 219 |
+
|
| 220 |
+
decoded_tags = np.concatenate(
|
| 221 |
+
[decoded_tags, np.array([[new_tag]], dtype=np.int64)], axis=0
|
| 222 |
+
)
|
| 223 |
+
t_dec = time.time() - t1
|
| 224 |
+
|
| 225 |
+
seq = decoded_tags.squeeze().tolist()
|
| 226 |
+
|
| 227 |
+
if tag_H_buf:
|
| 228 |
+
tag_H_stacked = np.stack(
|
| 229 |
+
[h[None, ...] if h.ndim == 1 else h for h in tag_H_buf], axis=0
|
| 230 |
+
).reshape(-1, 1, self.embed_dim).astype(np.float32)
|
| 231 |
+
t2 = time.time()
|
| 232 |
+
cls_logits, coord = self.bbox.run(None, {"enc_out": enc_out, "tag_H_stacked": tag_H_stacked})
|
| 233 |
+
t_bbox = time.time() - t2
|
| 234 |
+
coord = apply_bbox_merge(coord, bboxes_to_merge)
|
| 235 |
+
else:
|
| 236 |
+
cls_logits = np.empty((0, 3), dtype=np.float32)
|
| 237 |
+
coord = np.empty((0, 4), dtype=np.float32)
|
| 238 |
+
t_bbox = 0.0
|
| 239 |
+
|
| 240 |
+
return {
|
| 241 |
+
"seq": seq,
|
| 242 |
+
"outputs_class": cls_logits,
|
| 243 |
+
"outputs_coord": coord,
|
| 244 |
+
"bboxes_to_merge": bboxes_to_merge,
|
| 245 |
+
"timings": {"encoder": t_enc, "decoder": t_dec, "bbox": t_bbox,
|
| 246 |
+
"total": t_enc + t_dec + t_bbox},
|
| 247 |
+
}
|
| 248 |
+
|
| 249 |
+
|
| 250 |
+
# --------------------------------------------------------------------------- #
|
| 251 |
+
# OTSL assembly: structure tokens + per-cell OCR -> docling-ready OTSL string
|
| 252 |
+
# --------------------------------------------------------------------------- #
|
| 253 |
+
def compute_ioa(ocr_bbox: list[float], cell_bbox: list[float]) -> float:
|
| 254 |
+
"""Intersection over OCR-Token Area (IoA): maps a token to a cell if most
|
| 255 |
+
of the token sits inside it."""
|
| 256 |
+
ox1, oy1, ox2, oy2 = ocr_bbox
|
| 257 |
+
cx1, cy1, cx2, cy2 = cell_bbox
|
| 258 |
+
|
| 259 |
+
ix1 = max(ox1, cx1)
|
| 260 |
+
iy1 = max(oy1, cy1)
|
| 261 |
+
ix2 = min(ox2, cx2)
|
| 262 |
+
iy2 = min(oy2, cy2)
|
| 263 |
+
|
| 264 |
+
if ix2 <= ix1 or iy2 <= iy1:
|
| 265 |
+
return 0.0
|
| 266 |
+
|
| 267 |
+
inter_area = (ix2 - ix1) * (iy2 - iy1)
|
| 268 |
+
ocr_area = (ox2 - ox1) * (oy2 - oy1)
|
| 269 |
+
return inter_area / max(ocr_area, 1e-6)
|
| 270 |
+
|
| 271 |
+
|
| 272 |
+
def ocr_anchor_cells(
|
| 273 |
+
orig_image_path: str,
|
| 274 |
+
seq: list,
|
| 275 |
+
box_o: np.ndarray,
|
| 276 |
+
rev_tag: dict,
|
| 277 |
+
backend: Optional[OCRBackend] = None,
|
| 278 |
+
tokens: Optional[list[dict]] = None,
|
| 279 |
+
pad: int = 4,
|
| 280 |
+
narrow_width: float = 0.08,
|
| 281 |
+
) -> list:
|
| 282 |
+
"""Maps text tokens into TableFormer cell boundaries via spatial intersection.
|
| 283 |
+
|
| 284 |
+
Either `backend` (run OCR on the table crop) or pre-extracted `tokens`
|
| 285 |
+
(e.g. pdfplumber words already expressed in the crop's local pixel space β
|
| 286 |
+
see pipeline.py) must be supplied. `tokens` takes priority: when a PDF
|
| 287 |
+
has a real text layer, reusing those exact words is both cheaper and more
|
| 288 |
+
accurate than re-OCRing a table crop pdfplumber already read correctly.
|
| 289 |
+
"""
|
| 290 |
+
if tokens is None and backend is None:
|
| 291 |
+
raise ValueError("ocr_anchor_cells needs either `tokens` or `backend`")
|
| 292 |
+
|
| 293 |
+
orig = Image.open(orig_image_path).convert("RGB")
|
| 294 |
+
orig_w, orig_h = orig.size
|
| 295 |
+
|
| 296 |
+
anchor_tags = [rev_tag[t] for t in seq[1:-1] if rev_tag[t] in ANCHOR_TAGS]
|
| 297 |
+
assert len(anchor_tags) == box_o.shape[0], (
|
| 298 |
+
f"{len(anchor_tags)} anchor tags vs {box_o.shape[0]} boxes β structural alignment issue."
|
| 299 |
+
)
|
| 300 |
+
|
| 301 |
+
ocr_tokens = tokens if tokens is not None else backend.get_text_boxes(orig)
|
| 302 |
+
source = "pdfplumber tokens" if tokens is not None else type(backend).__name__
|
| 303 |
+
logger.debug(
|
| 304 |
+
"ocr_anchor_cells: %d anchor cells, %d tokens from %s",
|
| 305 |
+
len(anchor_tags),
|
| 306 |
+
len(ocr_tokens),
|
| 307 |
+
source,
|
| 308 |
+
)
|
| 309 |
+
texts = []
|
| 310 |
+
|
| 311 |
+
for i, (tag, (cx, cy, w, h)) in enumerate(zip(anchor_tags, box_o)):
|
| 312 |
+
if tag == "ecel":
|
| 313 |
+
texts.append("")
|
| 314 |
+
continue
|
| 315 |
+
|
| 316 |
+
x1 = max(0, (cx - w / 2) * orig_w - pad)
|
| 317 |
+
y1 = max(0, (cy - h / 2) * orig_h - pad)
|
| 318 |
+
x2 = min(orig_w, (cx + w / 2) * orig_w + pad)
|
| 319 |
+
y2 = min(orig_h, (cy + h / 2) * orig_h + pad)
|
| 320 |
+
cell_bbox = [x1, y1, x2, y2]
|
| 321 |
+
|
| 322 |
+
matched_tokens = [t for t in ocr_tokens if compute_ioa(t["bbox"], cell_bbox) >= 0.45]
|
| 323 |
+
|
| 324 |
+
if not matched_tokens:
|
| 325 |
+
texts.append("")
|
| 326 |
+
logger.debug(" cell[%d] tag=%s (empty)", i, tag)
|
| 327 |
+
continue
|
| 328 |
+
|
| 329 |
+
matched_tokens.sort(key=lambda t: (int(t["bbox"][1] / 6), t["bbox"][0]))
|
| 330 |
+
cell_string = " ".join(t["text"] for t in matched_tokens)
|
| 331 |
+
texts.append(cell_string.strip())
|
| 332 |
+
logger.debug(" cell[%d] tag=%s text=%r", i, tag, cell_string.strip()[:60])
|
| 333 |
+
|
| 334 |
+
return texts
|
| 335 |
+
|
| 336 |
+
|
| 337 |
+
def _otsl_cell_text(tag: str, text: str) -> str | None:
|
| 338 |
+
if tag == "ecel":
|
| 339 |
+
return None
|
| 340 |
+
cleaned = (text or "").strip()
|
| 341 |
+
return cleaned if cleaned else EMPTY_OTSL_PLACEHOLDER
|
| 342 |
+
|
| 343 |
+
|
| 344 |
+
def seq_to_otsl(seq: list, rev_tag: dict, texts: list) -> str:
|
| 345 |
+
"""Re-walk seq, inserting OCR text right after each anchor tag."""
|
| 346 |
+
out, ti = [], 0
|
| 347 |
+
for tok in seq[1:-1]:
|
| 348 |
+
name = rev_tag[tok]
|
| 349 |
+
if name == "nl":
|
| 350 |
+
out.append("<nl>")
|
| 351 |
+
continue
|
| 352 |
+
out.append(f"<{name}>")
|
| 353 |
+
if name in ANCHOR_TAGS:
|
| 354 |
+
cell_text = _otsl_cell_text(name, texts[ti])
|
| 355 |
+
if cell_text is not None:
|
| 356 |
+
out.append(cell_text)
|
| 357 |
+
ti += 1
|
| 358 |
+
return "".join(out)
|
| 359 |
+
|
| 360 |
+
|
| 361 |
+
def table_image_to_otsl(
|
| 362 |
+
runner: TableFormerONNX,
|
| 363 |
+
orig_image_path: str,
|
| 364 |
+
backend: Optional[OCRBackend] = None,
|
| 365 |
+
tokens: Optional[list[dict]] = None,
|
| 366 |
+
) -> str:
|
| 367 |
+
"""End-to-end: image -> ONNX TableFormer -> anchors filled from `tokens`
|
| 368 |
+
(preferred, e.g. pdfplumber words) or OCR'd via `backend` -> OTSL string."""
|
| 369 |
+
logger.info("TableFormer: processing %s", orig_image_path)
|
| 370 |
+
preproc = runner.image_preprocess(orig_image_path)
|
| 371 |
+
out = runner.predict(preproc)
|
| 372 |
+
logger.debug(
|
| 373 |
+
"TableFormer: seq length=%d, anchor boxes=%d",
|
| 374 |
+
len(out["seq"]),
|
| 375 |
+
out["outputs_coord"].shape[0],
|
| 376 |
+
)
|
| 377 |
+
texts = ocr_anchor_cells(
|
| 378 |
+
orig_image_path, out["seq"], out["outputs_coord"], runner.rev_tag,
|
| 379 |
+
backend=backend, tokens=tokens,
|
| 380 |
+
)
|
| 381 |
+
otsl = seq_to_otsl(out["seq"], runner.rev_tag, texts)
|
| 382 |
+
logger.debug("TableFormer OTSL preview: %s", otsl[:200])
|
| 383 |
+
return otsl
|
| 384 |
+
|
| 385 |
+
|
| 386 |
+
def collapse_spanned_columns(df: "pd.DataFrame") -> "pd.DataFrame":
|
| 387 |
+
"""
|
| 388 |
+
Collapse adjacent columns that are exact duplicates (same header text,
|
| 389 |
+
identical value in every row) back into a single column.
|
| 390 |
+
|
| 391 |
+
A plain rectangular DataFrame has no native way to represent a cell that
|
| 392 |
+
spans multiple columns (e.g. an OTSL `lcel` chain / HTML `colspan`).
|
| 393 |
+
Both `TableItem.export_to_dataframe()` and the `export_to_html()` +
|
| 394 |
+
`pd.read_html()` round-trip handle this the same way: they repeat the
|
| 395 |
+
spanned cell's value into every column it covers, so a 2-column-wide
|
| 396 |
+
"Description" header ends up as two separate "Description" columns
|
| 397 |
+
with identical content. This collapses those back into one.
|
| 398 |
+
|
| 399 |
+
Note: this is a heuristic β if two genuinely distinct adjacent columns
|
| 400 |
+
happen to share both the same header text and identical values in
|
| 401 |
+
every row of this specific table, they'll also get collapsed. That's
|
| 402 |
+
expected to be rare in practice (real tables don't usually have two
|
| 403 |
+
different columns with the same name and identical data), but worth
|
| 404 |
+
knowing if a table doesn't collapse cleanly.
|
| 405 |
+
"""
|
| 406 |
+
cols = list(df.columns)
|
| 407 |
+
if len(cols) <= 1:
|
| 408 |
+
return df
|
| 409 |
+
|
| 410 |
+
keep = [0]
|
| 411 |
+
for i in range(1, len(cols)):
|
| 412 |
+
prev = keep[-1]
|
| 413 |
+
# pandas suffixes repeated header names with ".1", ".2", ... on
|
| 414 |
+
# read_html; export_to_dataframe doesn't, so strip defensively.
|
| 415 |
+
prev_header = str(cols[prev]).split(".")[0]
|
| 416 |
+
cur_header = str(cols[i]).split(".")[0]
|
| 417 |
+
same_header = prev_header == cur_header
|
| 418 |
+
same_values = df.iloc[:, prev].astype(str).equals(df.iloc[:, i].astype(str))
|
| 419 |
+
if same_header and same_values:
|
| 420 |
+
continue # duplicate of the column we're keeping β drop it
|
| 421 |
+
keep.append(i)
|
| 422 |
+
|
| 423 |
+
out = df.iloc[:, keep].copy()
|
| 424 |
+
out.columns = [str(c).split(".")[0] for c in out.columns]
|
| 425 |
+
return out
|
| 426 |
+
|
| 427 |
+
|
| 428 |
+
def _parse_md_table_row(line: str) -> list[str]:
|
| 429 |
+
"""Split one markdown table row into its cell strings, stripping whitespace."""
|
| 430 |
+
return [c.strip() for c in line.strip().strip("|").split("|")]
|
| 431 |
+
|
| 432 |
+
|
| 433 |
+
def _is_separator_row(cells: list[str]) -> bool:
|
| 434 |
+
"""Return True if every cell is a markdown alignment/separator marker (e.g. `---`, `:---:`)."""
|
| 435 |
+
return all(set(c.replace(":", "").replace("-", "")) == set() and len(c) > 0 for c in cells)
|
| 436 |
+
|
| 437 |
+
|
| 438 |
+
def collapse_markdown_table(md: str) -> str:
|
| 439 |
+
"""
|
| 440 |
+
Apply the same duplicate-adjacent-column collapse to every markdown table
|
| 441 |
+
inside `md`, in-place. This is necessary because markdown has no colspan
|
| 442 |
+
syntax β docling_core's `export_to_markdown()` represents a colspan-N cell
|
| 443 |
+
by repeating its value in N consecutive columns, which looks like
|
| 444 |
+
`| Description | Description | Total |` for a 2-column-wide "Description"
|
| 445 |
+
header. This function detects and removes those duplicate columns from the
|
| 446 |
+
markdown string so the rendered table matches what the source document
|
| 447 |
+
actually meant.
|
| 448 |
+
|
| 449 |
+
The collapse criterion is identical to `collapse_spanned_columns`: two
|
| 450 |
+
adjacent columns are merged when they share the same header text (after
|
| 451 |
+
stripping pandas `.1`/`.2` dedup suffixes, though docling markdown doesn't
|
| 452 |
+
add those) AND every data row has the same value in both columns.
|
| 453 |
+
|
| 454 |
+
Non-table lines (headings, paragraphs, etc.) are passed through unchanged.
|
| 455 |
+
"""
|
| 456 |
+
lines = md.splitlines()
|
| 457 |
+
out_lines: list[str] = []
|
| 458 |
+
i = 0
|
| 459 |
+
|
| 460 |
+
while i < len(lines):
|
| 461 |
+
line = lines[i]
|
| 462 |
+
|
| 463 |
+
# Detect the start of a markdown table: a line containing at least one |
|
| 464 |
+
# that isn't just punctuation.
|
| 465 |
+
if "|" not in line:
|
| 466 |
+
out_lines.append(line)
|
| 467 |
+
i += 1
|
| 468 |
+
continue
|
| 469 |
+
|
| 470 |
+
# Collect the full table block (contiguous lines containing |).
|
| 471 |
+
table_block: list[str] = []
|
| 472 |
+
while i < len(lines) and "|" in lines[i]:
|
| 473 |
+
table_block.append(lines[i])
|
| 474 |
+
i += 1
|
| 475 |
+
|
| 476 |
+
if len(table_block) < 2:
|
| 477 |
+
# Not a real table (need at least header + separator).
|
| 478 |
+
out_lines.extend(table_block)
|
| 479 |
+
continue
|
| 480 |
+
|
| 481 |
+
# Parse into header, optional separator, and data rows.
|
| 482 |
+
header_cells = _parse_md_table_row(table_block[0])
|
| 483 |
+
sep_idx = 1 if _is_separator_row(_parse_md_table_row(table_block[1])) else None
|
| 484 |
+
|
| 485 |
+
data_rows: list[list[str]] = []
|
| 486 |
+
for row_line in table_block[(2 if sep_idx is not None else 1):]:
|
| 487 |
+
data_rows.append(_parse_md_table_row(row_line))
|
| 488 |
+
|
| 489 |
+
if not data_rows:
|
| 490 |
+
out_lines.extend(table_block)
|
| 491 |
+
continue
|
| 492 |
+
|
| 493 |
+
# Pad short rows so indexing is safe.
|
| 494 |
+
n_cols = len(header_cells)
|
| 495 |
+
data_rows = [r + [""] * max(0, n_cols - len(r)) for r in data_rows]
|
| 496 |
+
|
| 497 |
+
# Determine which columns to keep β same logic as collapse_spanned_columns.
|
| 498 |
+
keep: list[int] = [0]
|
| 499 |
+
for j in range(1, n_cols):
|
| 500 |
+
prev = keep[-1]
|
| 501 |
+
prev_header = header_cells[prev].split(".")[0]
|
| 502 |
+
cur_header = header_cells[j].split(".")[0]
|
| 503 |
+
same_header = prev_header == cur_header
|
| 504 |
+
same_values = all(
|
| 505 |
+
(row[prev] if prev < len(row) else "") == (row[j] if j < len(row) else "")
|
| 506 |
+
for row in data_rows
|
| 507 |
+
)
|
| 508 |
+
if same_header and same_values:
|
| 509 |
+
continue
|
| 510 |
+
keep.append(j)
|
| 511 |
+
|
| 512 |
+
if len(keep) == n_cols:
|
| 513 |
+
# Nothing to collapse β pass through unchanged.
|
| 514 |
+
out_lines.extend(table_block)
|
| 515 |
+
continue
|
| 516 |
+
|
| 517 |
+
# Rebuild the table with only the kept columns.
|
| 518 |
+
def _fmt_row(cells: list[str], indices: list[int]) -> str:
|
| 519 |
+
return "| " + " | ".join(cells[k] for k in indices) + " |"
|
| 520 |
+
|
| 521 |
+
def _fmt_sep(indices: list[int]) -> str:
|
| 522 |
+
return "| " + " | ".join("---" for _ in indices) + " |"
|
| 523 |
+
|
| 524 |
+
out_lines.append(_fmt_row(header_cells, keep))
|
| 525 |
+
if sep_idx is not None:
|
| 526 |
+
out_lines.append(_fmt_sep(keep))
|
| 527 |
+
for row in data_rows:
|
| 528 |
+
out_lines.append(_fmt_row(row, keep))
|
| 529 |
+
|
| 530 |
+
return "\n".join(out_lines)
|
| 531 |
+
|
| 532 |
+
|
| 533 |
+
def otsl_to_markdown(
|
| 534 |
+
otsl: str, dummy_image_size: tuple[int, int] = (512, 512)
|
| 535 |
+
) -> tuple[str, "pd.DataFrame"]:
|
| 536 |
+
"""
|
| 537 |
+
Wrap a raw OTSL tag string as <doctag><otsl>...</otsl></doctag>, load it
|
| 538 |
+
through docling_core's DocTagsDocument/DoclingDocument, and export both
|
| 539 |
+
markdown and a DataFrame for the (first) detected table.
|
| 540 |
+
|
| 541 |
+
Returns `(markdown, dataframe)` β both have duplicate columns from
|
| 542 |
+
colspan cells collapsed:
|
| 543 |
+
- markdown: via `collapse_markdown_table` (markdown has no colspan syntax
|
| 544 |
+
so docling repeats the spanned cell's text into every column it covers)
|
| 545 |
+
- dataframe: via `TableItem.export_to_dataframe()` + `collapse_spanned_columns`
|
| 546 |
+
(same duplication issue at the DataFrame level, same fix)
|
| 547 |
+
"""
|
| 548 |
+
from docling_core.types.doc.document import DocTagsDocument, DoclingDocument
|
| 549 |
+
|
| 550 |
+
otsl_doc = f"<doctag><otsl>{otsl}</otsl></doctag>"
|
| 551 |
+
dummy_image = Image.new("RGB", dummy_image_size, color="black")
|
| 552 |
+
doctags_doc = DocTagsDocument.from_doctags_and_image_pairs([otsl_doc], [dummy_image])
|
| 553 |
+
doc = DoclingDocument.load_from_doctags(doctags_doc)
|
| 554 |
+
|
| 555 |
+
md = collapse_markdown_table(doc.export_to_markdown())
|
| 556 |
+
|
| 557 |
+
if doc.tables:
|
| 558 |
+
df = doc.tables[0].export_to_dataframe(doc=doc)
|
| 559 |
+
df = collapse_spanned_columns(df)
|
| 560 |
+
logger.info("OTSL -> markdown: %d rows x %d cols", df.shape[0], df.shape[1])
|
| 561 |
+
else:
|
| 562 |
+
df = pd.DataFrame()
|
| 563 |
+
logger.warning("OTSL -> markdown: no tables detected in OTSL output")
|
| 564 |
+
|
| 565 |
+
return md, df
|