BabelDOC 0.2.13__tar.gz → 0.2.15__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {babeldoc-0.2.13 → babeldoc-0.2.15}/PKG-INFO +3 -1
- {babeldoc-0.2.13 → babeldoc-0.2.15}/README.md +1 -0
- babeldoc-0.2.15/babeldoc/__init__.py +1 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/assets/assets.py +35 -1
- {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/assets/embedding_assets_metadata.py +10 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/const.py +1 -1
- {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/backend/pdf_creater.py +6 -1
- {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/midend/paragraph_finder.py +2 -0
- babeldoc-0.2.15/babeldoc/document_il/midend/table_parser.py +164 -0
- babeldoc-0.2.15/babeldoc/docvision/table_detection/rapidocr.py +316 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/high_level.py +37 -23
- {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/main.py +14 -1
- {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/pdfinterp.py +7 -7
- {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/translation_config.py +3 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/pyproject.toml +3 -2
- {babeldoc-0.2.13 → babeldoc-0.2.15}/uv.lock +100 -7
- babeldoc-0.2.13/babeldoc/__init__.py +0 -1
- {babeldoc-0.2.13 → babeldoc-0.2.15}/.cursorignore +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/.github/ISSUE_TEMPLATE/bug_report.yaml +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/.github/ISSUE_TEMPLATE/feature_request.yaml +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/.github/dependabot.yml +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/.github/labels.yml +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/.github/release-drafter.yml +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/.github/workflows/codeql.yml +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/.github/workflows/docs.yml +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/.github/workflows/labeler.yml +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/.github/workflows/publish-to-pypi.yml +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/.github/workflows/test.yml +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/.gitignore +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/.pre-commit-config.yaml +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/LICENSE +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/asynchronize/__init__.py +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/converter.py +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/__init__.py +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/babeldoc_exception/BabelDOCException.py +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/backend/__init__.py +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/frontend/__init__.py +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/frontend/il_creater.py +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/il_version_1.py +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/il_version_1.rnc +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/il_version_1.rng +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/il_version_1.xsd +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/midend/__init__.py +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/midend/add_debug_information.py +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/midend/detect_scanned_file.py +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/midend/il_translator.py +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/midend/il_translator_llm_only.py +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/midend/layout_parser.py +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/midend/remove_descent.py +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/midend/styles_and_formulas.py +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/midend/typesetting.py +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/translator/__init__.py +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/translator/cache.py +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/translator/translator.py +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/utils/fontmap.py +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/utils/layout_helper.py +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/utils/style_helper.py +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/xml_converter.py +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/docvision/README.md +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/docvision/__init__.py +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/docvision/doclayout.py +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/docvision/rpc_doclayout.py +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/progress_monitor.py +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/result_merger.py +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/split_manager.py +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/tools/generate_font_metadata.py +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/CODE_OF_CONDUCT.md +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/CONTRIBUTING.md +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/CONTRIBUTOR_REWARD.md +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/ImplementationDetails/AsyncTranslate/AsyncTranslate.md +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/ImplementationDetails/ILTranslator/ILTranslator.md +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/ImplementationDetails/PDFCreation/PDFCreation.md +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/ImplementationDetails/PDFParsing/PDFParsing.md +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/ImplementationDetails/ParagraphFinding/ParagraphFinding.md +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/ImplementationDetails/README.md +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/ImplementationDetails/StylesAndFormulas/StylesAndFormulas.md +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/ImplementationDetails/Typesetting/Typesetting.md +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/README.md +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/deploy.sh +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/images/babeldoc-banner.png +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/images/babeldoc-big-logo-darkmode-with-transparent-background.png +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/images/babeldoc-big-logo-darkmode-with-transparent-background.svg +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/images/babeldoc-big-logo-with-transparent-background.png +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/images/babeldoc-big-logo-with-transparent-background.svg +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/images/babeldoc-big-logo.png +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/images/babeldoc-contributor_reward_example.png +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/images/babeldoc-preview.gif +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/images/babeldoc-small-logo-with-transparent-background.png +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/images/babeldoc-small-logo-with-transparent-background.svg +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/images/babeldoc-small-logo.png +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/index.md +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/intro-to-pdf-object.md +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/requirements.txt +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/mkdocs.yml +0 -0
- {babeldoc-0.2.13 → babeldoc-0.2.15}/tests/test_translation_config.py +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: BabelDOC
|
|
3
|
-
Version: 0.2.
|
|
3
|
+
Version: 0.2.15
|
|
4
4
|
Summary: Yet Another Document Translator
|
|
5
5
|
Project-URL: Homepage, https://github.com/funstory-ai/BabelDOC
|
|
6
6
|
Project-URL: Issues, https://github.com/funstory-ai/BabelDOC/issues
|
|
@@ -30,6 +30,7 @@ Requires-Dist: psutil>=7.0.0
|
|
|
30
30
|
Requires-Dist: pydantic>=2.10.6
|
|
31
31
|
Requires-Dist: pymupdf>=1.25.1
|
|
32
32
|
Requires-Dist: python-levenshtein>=0.27.1
|
|
33
|
+
Requires-Dist: rapidocr-onnxruntime>=1.4.4
|
|
33
34
|
Requires-Dist: rich>=13.9.4
|
|
34
35
|
Requires-Dist: scikit-image>=0.25.2
|
|
35
36
|
Requires-Dist: tenacity>=9.0.0
|
|
@@ -181,6 +182,7 @@ uv run babeldoc --files example.pdf --files example2.pdf --openai --openai-model
|
|
|
181
182
|
- `--watermark-output-mode`: Control watermark output mode: 'watermarked' (default) adds watermark to translated PDF, 'no_watermark' doesn't add watermark, 'both' outputs both versions.
|
|
182
183
|
- `--max-pages-per-part`: Maximum number of pages per part for split translation. If not set, no splitting will be performed.
|
|
183
184
|
- `--no-watermark`: [DEPRECATED] Use --watermark-output-mode=no_watermark instead.
|
|
185
|
+
- `--translate-table-text`: Translate table text (experimental, default: False)
|
|
184
186
|
|
|
185
187
|
> [!TIP]
|
|
186
188
|
> - Both `--skip-clean` and `--dual-translate-first` may help improve compatibility with some PDF readers
|
|
@@ -136,6 +136,7 @@ uv run babeldoc --files example.pdf --files example2.pdf --openai --openai-model
|
|
|
136
136
|
- `--watermark-output-mode`: Control watermark output mode: 'watermarked' (default) adds watermark to translated PDF, 'no_watermark' doesn't add watermark, 'both' outputs both versions.
|
|
137
137
|
- `--max-pages-per-part`: Maximum number of pages per part for split translation. If not set, no splitting will be performed.
|
|
138
138
|
- `--no-watermark`: [DEPRECATED] Use --watermark-output-mode=no_watermark instead.
|
|
139
|
+
- `--translate-table-text`: Translate table text (experimental, default: False)
|
|
139
140
|
|
|
140
141
|
> [!TIP]
|
|
141
142
|
> - Both `--skip-clean` and `--dual-translate-first` may help improve compatibility with some PDF readers
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
__version__ = "0.2.15"
|
|
@@ -14,6 +14,10 @@ from babeldoc.assets.embedding_assets_metadata import (
|
|
|
14
14
|
from babeldoc.assets.embedding_assets_metadata import EMBEDDING_FONT_METADATA
|
|
15
15
|
from babeldoc.assets.embedding_assets_metadata import FONT_METADATA_URL
|
|
16
16
|
from babeldoc.assets.embedding_assets_metadata import FONT_URL_BY_UPSTREAM
|
|
17
|
+
from babeldoc.assets.embedding_assets_metadata import (
|
|
18
|
+
TABLE_DETECTION_RAPIDOCR_MODEL_SHA3_256,
|
|
19
|
+
)
|
|
20
|
+
from babeldoc.assets.embedding_assets_metadata import TABLE_DETECTION_RAPIDOCR_MODEL_URL
|
|
17
21
|
from babeldoc.assets.embedding_assets_metadata import TIKTOKEN_CACHES
|
|
18
22
|
from babeldoc.const import get_cache_file_path
|
|
19
23
|
from tenacity import retry
|
|
@@ -204,10 +208,36 @@ async def get_doclayout_onnx_model_path_async(client: httpx.AsyncClient | None =
|
|
|
204
208
|
return onnx_path
|
|
205
209
|
|
|
206
210
|
|
|
211
|
+
async def get_table_detection_rapidocr_model_path_async(
|
|
212
|
+
client: httpx.AsyncClient | None = None,
|
|
213
|
+
):
|
|
214
|
+
onnx_path = get_cache_file_path("ch_PP-OCRv4_det_infer.onnx", "models")
|
|
215
|
+
if verify_file(onnx_path, TABLE_DETECTION_RAPIDOCR_MODEL_SHA3_256):
|
|
216
|
+
return onnx_path
|
|
217
|
+
|
|
218
|
+
logger.info("table detection rapidocr model not found or corrupted, downloading...")
|
|
219
|
+
fastest_upstream, _ = await get_fastest_upstream_for_model(client)
|
|
220
|
+
if fastest_upstream is None:
|
|
221
|
+
logger.error("Failed to get fastest upstream")
|
|
222
|
+
exit(1)
|
|
223
|
+
|
|
224
|
+
url = TABLE_DETECTION_RAPIDOCR_MODEL_URL[fastest_upstream]
|
|
225
|
+
|
|
226
|
+
await download_file(client, url, onnx_path, TABLE_DETECTION_RAPIDOCR_MODEL_SHA3_256)
|
|
227
|
+
logger.info(
|
|
228
|
+
f"Download table detection rapidocr model from {fastest_upstream} success"
|
|
229
|
+
)
|
|
230
|
+
return onnx_path
|
|
231
|
+
|
|
232
|
+
|
|
207
233
|
def get_doclayout_onnx_model_path():
|
|
208
234
|
return run_coro(get_doclayout_onnx_model_path_async())
|
|
209
235
|
|
|
210
236
|
|
|
237
|
+
def get_table_detection_rapidocr_model_path():
|
|
238
|
+
return run_coro(get_table_detection_rapidocr_model_path_async())
|
|
239
|
+
|
|
240
|
+
|
|
211
241
|
def get_font_url_by_name_and_upstream(font_file_name: str, upstream: str):
|
|
212
242
|
if upstream not in FONT_URL_BY_UPSTREAM:
|
|
213
243
|
logger.critical(f"Invalid upstream: {upstream}")
|
|
@@ -329,7 +359,11 @@ def generate_all_assets_file_list():
|
|
|
329
359
|
{
|
|
330
360
|
"name": "doclayout_yolo_docstructbench_imgsz1024.onnx",
|
|
331
361
|
"sha3_256": DOCLAYOUT_YOLO_DOCSTRUCTBENCH_IMGSZ1024ONNX_SHA3_256,
|
|
332
|
-
}
|
|
362
|
+
},
|
|
363
|
+
{
|
|
364
|
+
"name": "ch_PP-OCRv4_det_infer.onnx",
|
|
365
|
+
"sha3_256": TABLE_DETECTION_RAPIDOCR_MODEL_SHA3_256,
|
|
366
|
+
},
|
|
333
367
|
)
|
|
334
368
|
return result
|
|
335
369
|
|
|
@@ -2,6 +2,10 @@ DOCLAYOUT_YOLO_DOCSTRUCTBENCH_IMGSZ1024ONNX_SHA3_256 = (
|
|
|
2
2
|
"60be061226930524958b5465c8c04af3d7c03bcb0beb66454f5da9f792e3cf2a"
|
|
3
3
|
)
|
|
4
4
|
|
|
5
|
+
TABLE_DETECTION_RAPIDOCR_MODEL_SHA3_256 = (
|
|
6
|
+
"062f4619afe91b33147c033acadecbb53f2a7b99ac703d157b96d5b10948da5e"
|
|
7
|
+
)
|
|
8
|
+
|
|
5
9
|
TIKTOKEN_CACHES = {
|
|
6
10
|
"fb374d419588a4632f3f557e76b4b70aebbca790": "cb04bcda5782cfbbe77f2f991d92c0ea785d9496ef1137c91dfc3c8c324528d6"
|
|
7
11
|
}
|
|
@@ -26,6 +30,12 @@ DOC_LAYOUT_ONNX_MODEL_URL = {
|
|
|
26
30
|
"modelscope": "https://www.modelscope.cn/models/AI-ModelScope/DocLayout-YOLO-DocStructBench-onnx/resolve/master/doclayout_yolo_docstructbench_imgsz1024.onnx",
|
|
27
31
|
}
|
|
28
32
|
|
|
33
|
+
TABLE_DETECTION_RAPIDOCR_MODEL_URL = {
|
|
34
|
+
"huggingface": "https://huggingface.co/spaces/RapidAI/RapidOCR/resolve/main/models/text_det/ch_PP-OCRv4_det_infer.onnx",
|
|
35
|
+
"hf-mirror": "https://hf-mirror.com/spaces/RapidAI/RapidOCR/resolve/main/models/text_det/ch_PP-OCRv4_det_infer.onnx",
|
|
36
|
+
"modelscope": "https://www.modelscope.cn/models/RapidAI/RapidOCR/resolve/master/onnx/PP-OCRv4/det/ch_PP-OCRv4_det_infer.onnx",
|
|
37
|
+
}
|
|
38
|
+
|
|
29
39
|
# from https://github.com/funstory-ai/BabelDOC-Assets/blob/main/font_metadata.json
|
|
30
40
|
EMBEDDING_FONT_METADATA = {
|
|
31
41
|
"GoNotoKurrent-Bold.ttf": {
|
|
@@ -730,7 +730,12 @@ class PDFCreater:
|
|
|
730
730
|
draw_op.append(b" Tj ET Q \n")
|
|
731
731
|
for xobj in page.pdf_xobject:
|
|
732
732
|
draw_op = xobj_draw_ops[xobj.xobj_id]
|
|
733
|
-
|
|
733
|
+
try:
|
|
734
|
+
pdf.update_stream(xobj.xref_id, draw_op.tobytes())
|
|
735
|
+
except Exception:
|
|
736
|
+
logger.warning(
|
|
737
|
+
f"update xref {xobj.xref_id} stream fail, continue"
|
|
738
|
+
)
|
|
734
739
|
# pdf.update_stream(xobj.xref_id, b'')
|
|
735
740
|
for rect in page.pdf_rectangle:
|
|
736
741
|
self._debug_render_rectangle(page_op, rect)
|
|
@@ -287,6 +287,7 @@ class ParagraphFinder:
|
|
|
287
287
|
"abandon",
|
|
288
288
|
"figure_caption",
|
|
289
289
|
"table_caption",
|
|
290
|
+
"table_text",
|
|
290
291
|
]
|
|
291
292
|
|
|
292
293
|
def get_layout(
|
|
@@ -313,6 +314,7 @@ class ParagraphFinder:
|
|
|
313
314
|
"table_footnote",
|
|
314
315
|
"table_caption",
|
|
315
316
|
"figure_caption",
|
|
317
|
+
"table_text",
|
|
316
318
|
"table",
|
|
317
319
|
"figure",
|
|
318
320
|
"abandon",
|
|
@@ -0,0 +1,164 @@
|
|
|
1
|
+
import logging
|
|
2
|
+
from pathlib import Path
|
|
3
|
+
|
|
4
|
+
import cv2
|
|
5
|
+
import numpy as np
|
|
6
|
+
from pymupdf import Document
|
|
7
|
+
|
|
8
|
+
from babeldoc.document_il import il_version_1
|
|
9
|
+
from babeldoc.document_il.utils.style_helper import GREEN
|
|
10
|
+
from babeldoc.translation_config import TranslationConfig
|
|
11
|
+
|
|
12
|
+
logger = logging.getLogger(__name__)
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
class TableParser:
|
|
16
|
+
stage_name = "Parse Table"
|
|
17
|
+
|
|
18
|
+
def __init__(self, translation_config: TranslationConfig):
|
|
19
|
+
self.translation_config = translation_config
|
|
20
|
+
self.model = translation_config.table_model
|
|
21
|
+
|
|
22
|
+
def _save_debug_image(self, image: np.ndarray, layouts, page_number: int):
|
|
23
|
+
"""Save debug image with drawn boxes if debug mode is enabled."""
|
|
24
|
+
if not self.translation_config.debug:
|
|
25
|
+
return
|
|
26
|
+
|
|
27
|
+
if not isinstance(layouts, list):
|
|
28
|
+
layouts = [layouts]
|
|
29
|
+
debug_dir = Path(
|
|
30
|
+
self.translation_config.get_working_file_path("table-ocr-box-image")
|
|
31
|
+
)
|
|
32
|
+
debug_dir.mkdir(parents=True, exist_ok=True)
|
|
33
|
+
|
|
34
|
+
# Draw boxes on the image
|
|
35
|
+
debug_image = image.copy()
|
|
36
|
+
for layout in layouts:
|
|
37
|
+
for box in layout.boxes:
|
|
38
|
+
x0, y0, x1, y1 = box.xyxy
|
|
39
|
+
cv2.rectangle(
|
|
40
|
+
debug_image,
|
|
41
|
+
(int(x0), int(y0)),
|
|
42
|
+
(int(x1), int(y1)),
|
|
43
|
+
(0, 255, 0),
|
|
44
|
+
2,
|
|
45
|
+
)
|
|
46
|
+
# Add text label
|
|
47
|
+
cv2.putText(
|
|
48
|
+
debug_image,
|
|
49
|
+
layout.names[box.cls],
|
|
50
|
+
(int(x0), int(y0) - 5),
|
|
51
|
+
cv2.FONT_HERSHEY_SIMPLEX,
|
|
52
|
+
0.5,
|
|
53
|
+
(0, 255, 0),
|
|
54
|
+
1,
|
|
55
|
+
)
|
|
56
|
+
|
|
57
|
+
# Save the image
|
|
58
|
+
output_path = debug_dir / f"{page_number}.jpg"
|
|
59
|
+
cv2.imwrite(str(output_path), debug_image)
|
|
60
|
+
|
|
61
|
+
def _save_debug_box_to_page(self, page: il_version_1.Page):
|
|
62
|
+
"""Save debug boxes and text labels to the PDF page."""
|
|
63
|
+
if not self.translation_config.debug:
|
|
64
|
+
return
|
|
65
|
+
|
|
66
|
+
color = GREEN
|
|
67
|
+
|
|
68
|
+
for layout in page.page_layout:
|
|
69
|
+
# Create a rectangle box
|
|
70
|
+
rect = il_version_1.PdfRectangle(
|
|
71
|
+
box=il_version_1.Box(
|
|
72
|
+
x=layout.box.x,
|
|
73
|
+
y=layout.box.y,
|
|
74
|
+
x2=layout.box.x2,
|
|
75
|
+
y2=layout.box.y2,
|
|
76
|
+
),
|
|
77
|
+
graphic_state=color,
|
|
78
|
+
debug_info=True,
|
|
79
|
+
)
|
|
80
|
+
page.pdf_rectangle.append(rect)
|
|
81
|
+
|
|
82
|
+
# Create text label at top-left corner
|
|
83
|
+
# Note: PDF coordinates are from bottom-left,
|
|
84
|
+
# so we use y2 for top position
|
|
85
|
+
style = il_version_1.PdfStyle(
|
|
86
|
+
font_id="china-ss",
|
|
87
|
+
font_size=4,
|
|
88
|
+
graphic_state=color,
|
|
89
|
+
)
|
|
90
|
+
page.pdf_paragraph.append(
|
|
91
|
+
il_version_1.PdfParagraph(
|
|
92
|
+
first_line_indent=False,
|
|
93
|
+
box=il_version_1.Box(
|
|
94
|
+
x=layout.box.x,
|
|
95
|
+
y=layout.box.y2,
|
|
96
|
+
x2=layout.box.x2,
|
|
97
|
+
y2=layout.box.y2 + 5,
|
|
98
|
+
),
|
|
99
|
+
vertical=False,
|
|
100
|
+
pdf_style=style,
|
|
101
|
+
unicode=layout.class_name,
|
|
102
|
+
pdf_paragraph_composition=[
|
|
103
|
+
il_version_1.PdfParagraphComposition(
|
|
104
|
+
pdf_same_style_unicode_characters=il_version_1.PdfSameStyleUnicodeCharacters(
|
|
105
|
+
unicode=layout.class_name,
|
|
106
|
+
pdf_style=style,
|
|
107
|
+
debug_info=True,
|
|
108
|
+
),
|
|
109
|
+
),
|
|
110
|
+
],
|
|
111
|
+
xobj_id=-1,
|
|
112
|
+
),
|
|
113
|
+
)
|
|
114
|
+
|
|
115
|
+
def process(self, docs: il_version_1.Document, mupdf_doc: Document):
|
|
116
|
+
"""Generate layouts for all pages that need to be translated."""
|
|
117
|
+
# Get pages that need to be translated
|
|
118
|
+
have_table_pages = {}
|
|
119
|
+
for page in docs.page:
|
|
120
|
+
for layout in page.page_layout:
|
|
121
|
+
if layout.class_name == "table":
|
|
122
|
+
have_table_pages[page.page_number] = page
|
|
123
|
+
with self.translation_config.progress_monitor.stage_start(
|
|
124
|
+
self.stage_name,
|
|
125
|
+
len(have_table_pages),
|
|
126
|
+
) as progress:
|
|
127
|
+
# Process predictions for each page
|
|
128
|
+
for page, layouts in self.model.handle_document(
|
|
129
|
+
have_table_pages.values(),
|
|
130
|
+
mupdf_doc,
|
|
131
|
+
self.translation_config,
|
|
132
|
+
self._save_debug_image,
|
|
133
|
+
):
|
|
134
|
+
page_layouts = []
|
|
135
|
+
for layout in layouts.boxes:
|
|
136
|
+
# Convert coordinate system from picture to il
|
|
137
|
+
# system to the il coordinate system
|
|
138
|
+
x0, y0, x1, y1 = layout.xyxy
|
|
139
|
+
pix = mupdf_doc[page.page_number].get_pixmap()
|
|
140
|
+
h, w = pix.height, pix.width
|
|
141
|
+
x0, y0, x1, y1 = (
|
|
142
|
+
np.clip(int(x0 - 1), 0, w - 1),
|
|
143
|
+
np.clip(int(h - y1 - 1), 0, h - 1),
|
|
144
|
+
np.clip(int(x1 + 1), 0, w - 1),
|
|
145
|
+
np.clip(int(h - y0 + 1), 0, h - 1),
|
|
146
|
+
)
|
|
147
|
+
page_layout = il_version_1.PageLayout(
|
|
148
|
+
id=len(page_layouts) + 1,
|
|
149
|
+
box=il_version_1.Box(
|
|
150
|
+
x0.item(),
|
|
151
|
+
y0.item(),
|
|
152
|
+
x1.item(),
|
|
153
|
+
y1.item(),
|
|
154
|
+
),
|
|
155
|
+
conf=layout.conf.item(),
|
|
156
|
+
class_name=layouts.names[layout.cls],
|
|
157
|
+
)
|
|
158
|
+
page_layouts.append(page_layout)
|
|
159
|
+
|
|
160
|
+
page.page_layout.extend(page_layouts)
|
|
161
|
+
self._save_debug_box_to_page(page)
|
|
162
|
+
progress.advance(1)
|
|
163
|
+
|
|
164
|
+
return docs
|
|
@@ -0,0 +1,316 @@
|
|
|
1
|
+
import logging
|
|
2
|
+
import re
|
|
3
|
+
from collections.abc import Generator
|
|
4
|
+
|
|
5
|
+
import cv2
|
|
6
|
+
import numpy as np
|
|
7
|
+
from babeldoc.assets.assets import get_table_detection_rapidocr_model_path
|
|
8
|
+
from babeldoc.docvision.doclayout import YoloBox
|
|
9
|
+
from babeldoc.docvision.doclayout import YoloResult
|
|
10
|
+
from rapidocr_onnxruntime import RapidOCR
|
|
11
|
+
|
|
12
|
+
try:
|
|
13
|
+
import onnxruntime
|
|
14
|
+
except ImportError as e:
|
|
15
|
+
if "DLL load failed" in str(e):
|
|
16
|
+
raise OSError(
|
|
17
|
+
"Microsoft Visual C++ Redistributable is not installed. "
|
|
18
|
+
"Download it at https://aka.ms/vs/17/release/vc_redist.x64.exe"
|
|
19
|
+
) from e
|
|
20
|
+
raise
|
|
21
|
+
import babeldoc.document_il.il_version_1
|
|
22
|
+
import pymupdf
|
|
23
|
+
|
|
24
|
+
logger = logging.getLogger(__name__)
|
|
25
|
+
|
|
26
|
+
|
|
27
|
+
def convert_to_yolo_result(predictions):
|
|
28
|
+
"""
|
|
29
|
+
Convert RapidOCR predictions to YoloResult format.
|
|
30
|
+
|
|
31
|
+
Args:
|
|
32
|
+
predictions (list): List of predictions, where each prediction is a list of coordinates
|
|
33
|
+
in format [[x1, y1], [x2, y2], [x3, y3], [x4, y4], (text, confidence)]
|
|
34
|
+
or a numpy array of format [x1, y1, x2, y2, ...]
|
|
35
|
+
|
|
36
|
+
Returns:
|
|
37
|
+
YoloResult: Converted predictions in YoloResult format
|
|
38
|
+
"""
|
|
39
|
+
boxes = []
|
|
40
|
+
|
|
41
|
+
for pred in predictions:
|
|
42
|
+
# Check if the prediction is in the format of 4 corner points
|
|
43
|
+
if isinstance(pred, list) and len(pred) >= 5 and isinstance(pred[0], list):
|
|
44
|
+
# Convert 4 corner points to xyxy format (min x, min y, max x, max y)
|
|
45
|
+
points = np.array(pred[:4])
|
|
46
|
+
x1, y1 = points[:, 0].min(), points[:, 1].min()
|
|
47
|
+
x2, y2 = points[:, 0].max(), points[:, 1].max()
|
|
48
|
+
xyxy = [x1, y1, x2, y2]
|
|
49
|
+
box = YoloBox(xyxy=xyxy, conf=1.0, cls="text")
|
|
50
|
+
# Check if the prediction is already in xyxy format
|
|
51
|
+
elif isinstance(pred, list | np.ndarray) and len(pred) >= 4:
|
|
52
|
+
if isinstance(pred, np.ndarray):
|
|
53
|
+
pred = pred.tolist()
|
|
54
|
+
xyxy = pred[:4]
|
|
55
|
+
box = YoloBox(xyxy=xyxy, conf=1.0, cls="text")
|
|
56
|
+
else:
|
|
57
|
+
continue
|
|
58
|
+
|
|
59
|
+
boxes.append(box)
|
|
60
|
+
|
|
61
|
+
return YoloResult(names=["text"], boxes=boxes)
|
|
62
|
+
|
|
63
|
+
|
|
64
|
+
def create_yolo_result_from_nested_coords(nested_coords: np.ndarray, names: dict):
|
|
65
|
+
boxes = []
|
|
66
|
+
|
|
67
|
+
for quad in nested_coords.tolist():
|
|
68
|
+
if len(quad) != 4:
|
|
69
|
+
continue
|
|
70
|
+
|
|
71
|
+
# Convert quad coordinates to xyxy format (min x, min y, max x, max y)
|
|
72
|
+
x1, y1, x2, y2 = quad
|
|
73
|
+
|
|
74
|
+
# Create YoloBox with confidence 1.0 and class 'text'
|
|
75
|
+
box = YoloBox(
|
|
76
|
+
xyxy=[float(x1), float(y1), float(x2), float(y2)], conf=np.array(1.0), cls=0
|
|
77
|
+
)
|
|
78
|
+
boxes.append(box)
|
|
79
|
+
|
|
80
|
+
return YoloResult(names=names, boxes=boxes)
|
|
81
|
+
|
|
82
|
+
|
|
83
|
+
class RapidOCRModel:
|
|
84
|
+
def __init__(self):
|
|
85
|
+
self.use_cuda = False
|
|
86
|
+
self.use_dml = False
|
|
87
|
+
available_providers = onnxruntime.get_available_providers()
|
|
88
|
+
for provider in available_providers:
|
|
89
|
+
if re.match(r"dml", provider, re.IGNORECASE):
|
|
90
|
+
self.use_dml = True
|
|
91
|
+
elif re.match(r"cuda", provider, re.IGNORECASE):
|
|
92
|
+
self.use_cuda = True
|
|
93
|
+
|
|
94
|
+
self.model = RapidOCR(
|
|
95
|
+
det_model_path=get_table_detection_rapidocr_model_path(),
|
|
96
|
+
det_use_cuda=self.use_cuda,
|
|
97
|
+
det_use_dml=self.use_dml,
|
|
98
|
+
)
|
|
99
|
+
self.names = {0: "table_text"}
|
|
100
|
+
|
|
101
|
+
@property
|
|
102
|
+
def stride(self):
|
|
103
|
+
return 32
|
|
104
|
+
|
|
105
|
+
def resize_and_pad_image(self, image, new_shape):
|
|
106
|
+
"""
|
|
107
|
+
Resize and pad the image to the specified size, ensuring dimensions are multiples of stride.
|
|
108
|
+
|
|
109
|
+
Parameters:
|
|
110
|
+
- image: Input image
|
|
111
|
+
- new_shape: Target size (integer or (height, width) tuple)
|
|
112
|
+
- stride: Padding alignment stride, default 32
|
|
113
|
+
|
|
114
|
+
Returns:
|
|
115
|
+
- Processed image
|
|
116
|
+
"""
|
|
117
|
+
if isinstance(new_shape, int):
|
|
118
|
+
new_shape = (new_shape, new_shape)
|
|
119
|
+
|
|
120
|
+
h, w = image.shape[:2]
|
|
121
|
+
new_h, new_w = new_shape
|
|
122
|
+
|
|
123
|
+
# Calculate scaling ratio
|
|
124
|
+
r = min(new_h / h, new_w / w)
|
|
125
|
+
resized_h, resized_w = int(round(h * r)), int(round(w * r))
|
|
126
|
+
|
|
127
|
+
# Resize image
|
|
128
|
+
image = cv2.resize(
|
|
129
|
+
image,
|
|
130
|
+
(resized_w, resized_h),
|
|
131
|
+
interpolation=cv2.INTER_LINEAR,
|
|
132
|
+
)
|
|
133
|
+
|
|
134
|
+
# Calculate padding size and align to stride multiple
|
|
135
|
+
pad_w = (new_w - resized_w) % self.stride
|
|
136
|
+
pad_h = (new_h - resized_h) % self.stride
|
|
137
|
+
top, bottom = pad_h // 2, pad_h - pad_h // 2
|
|
138
|
+
left, right = pad_w // 2, pad_w - pad_w // 2
|
|
139
|
+
|
|
140
|
+
# Add padding
|
|
141
|
+
image = cv2.copyMakeBorder(
|
|
142
|
+
image,
|
|
143
|
+
top,
|
|
144
|
+
bottom,
|
|
145
|
+
left,
|
|
146
|
+
right,
|
|
147
|
+
cv2.BORDER_CONSTANT,
|
|
148
|
+
value=(114, 114, 114),
|
|
149
|
+
)
|
|
150
|
+
|
|
151
|
+
return image
|
|
152
|
+
|
|
153
|
+
def scale_boxes(self, img1_shape, boxes, img0_shape):
|
|
154
|
+
"""
|
|
155
|
+
Rescales bounding boxes (in the format of xyxy by default) from the shape of the image they were originally
|
|
156
|
+
specified in (img1_shape) to the shape of a different image (img0_shape).
|
|
157
|
+
|
|
158
|
+
Args:
|
|
159
|
+
img1_shape (tuple): The shape of the image that the bounding boxes are for,
|
|
160
|
+
in the format of (height, width).
|
|
161
|
+
boxes (torch.Tensor): the bounding boxes of the objects in the image, in the format of (x1, y1, x2, y2)
|
|
162
|
+
img0_shape (tuple): the shape of the target image, in the format of (height, width).
|
|
163
|
+
|
|
164
|
+
Returns:
|
|
165
|
+
boxes (torch.Tensor): The scaled bounding boxes, in the format of (x1, y1, x2, y2)
|
|
166
|
+
"""
|
|
167
|
+
|
|
168
|
+
# Calculate scaling ratio
|
|
169
|
+
gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1])
|
|
170
|
+
|
|
171
|
+
# Calculate padding size
|
|
172
|
+
pad_x = round((img1_shape[1] - img0_shape[1] * gain) / 2 - 0.1)
|
|
173
|
+
pad_y = round((img1_shape[0] - img0_shape[0] * gain) / 2 - 0.1)
|
|
174
|
+
|
|
175
|
+
# Remove padding and scale boxes
|
|
176
|
+
boxes[..., :4] = (boxes[..., :4] - [pad_x, pad_y, pad_x, pad_y]) / gain
|
|
177
|
+
return boxes
|
|
178
|
+
|
|
179
|
+
def predict(self, image, imgsz=800, batch_size=16, **kwargs):
|
|
180
|
+
"""
|
|
181
|
+
Predict the layout of document pages.
|
|
182
|
+
|
|
183
|
+
Args:
|
|
184
|
+
image: A single image or a list of images of document pages.
|
|
185
|
+
imgsz: Resize the image to this size. Must be a multiple of the stride.
|
|
186
|
+
batch_size: Number of images to process in one batch.
|
|
187
|
+
**kwargs: Additional arguments.
|
|
188
|
+
|
|
189
|
+
Returns:
|
|
190
|
+
A YoloResult object containing the detected boxes.
|
|
191
|
+
"""
|
|
192
|
+
# Handle single image input
|
|
193
|
+
assert isinstance(image, np.ndarray) and len(image.shape) == 3
|
|
194
|
+
|
|
195
|
+
# Calculate target size based on the maximum height in the batch
|
|
196
|
+
target_imgsz = 1024
|
|
197
|
+
|
|
198
|
+
orig_shape = (image.shape[0], image.shape[1])
|
|
199
|
+
|
|
200
|
+
pix = self.resize_and_pad_image(image, new_shape=target_imgsz)
|
|
201
|
+
# pix = np.transpose(pix, (2, 0, 1)) # CHW
|
|
202
|
+
# pix = pix.astype(np.float32) / 255.0 # Normalize to [0, 1]
|
|
203
|
+
input_ = pix
|
|
204
|
+
|
|
205
|
+
new_h, new_w = input_.shape[:2]
|
|
206
|
+
|
|
207
|
+
# Run inference
|
|
208
|
+
preds = self.model(input_, use_det=True, use_cls=False, use_rec=False)
|
|
209
|
+
|
|
210
|
+
# Process each prediction in the batch
|
|
211
|
+
if len(preds) > 0:
|
|
212
|
+
preds_np = np.array(preds[0])[:, [0, 2], :].reshape([-1, 4])
|
|
213
|
+
preds_np[..., :4] = self.scale_boxes(
|
|
214
|
+
(new_h, new_w),
|
|
215
|
+
preds_np[..., :4],
|
|
216
|
+
orig_shape,
|
|
217
|
+
)
|
|
218
|
+
|
|
219
|
+
# Convert predictions to YoloResult format
|
|
220
|
+
return create_yolo_result_from_nested_coords(preds_np, self.names)
|
|
221
|
+
else:
|
|
222
|
+
# Return empty YoloResult if no predictions
|
|
223
|
+
return YoloResult(names=self.names, boxes=[])
|
|
224
|
+
|
|
225
|
+
def handle_document(
|
|
226
|
+
self,
|
|
227
|
+
pages: list[babeldoc.document_il.il_version_1.Page],
|
|
228
|
+
mupdf_doc: pymupdf.Document,
|
|
229
|
+
translate_config,
|
|
230
|
+
save_debug_image,
|
|
231
|
+
) -> Generator[
|
|
232
|
+
tuple[babeldoc.document_il.il_version_1.Page, YoloResult], None, None
|
|
233
|
+
]:
|
|
234
|
+
for page in pages:
|
|
235
|
+
translate_config.raise_if_cancelled()
|
|
236
|
+
pix = mupdf_doc[page.page_number].get_pixmap(dpi=72)
|
|
237
|
+
image = np.fromstring(pix.samples, np.uint8).reshape(
|
|
238
|
+
pix.height,
|
|
239
|
+
pix.width,
|
|
240
|
+
3,
|
|
241
|
+
)[:, :, ::-1]
|
|
242
|
+
|
|
243
|
+
table_boxes = []
|
|
244
|
+
for layout in page.page_layout:
|
|
245
|
+
if layout.class_name == "table":
|
|
246
|
+
table_boxes.append(layout.box)
|
|
247
|
+
|
|
248
|
+
predict_result = self.predict(image)
|
|
249
|
+
|
|
250
|
+
ok_boxes = []
|
|
251
|
+
for box in predict_result.boxes:
|
|
252
|
+
# Convert the box coordinates to float for proper comparison
|
|
253
|
+
box_xyxy = [float(coord) for coord in box.xyxy]
|
|
254
|
+
|
|
255
|
+
# Check if this box is inside any of the table boxes
|
|
256
|
+
for table_box in table_boxes:
|
|
257
|
+
# Determine if box is inside or overlapping with table_box with image dimensions
|
|
258
|
+
if self._is_box_in_table(
|
|
259
|
+
box_xyxy, table_box, page, image.shape[1], image.shape[0]
|
|
260
|
+
):
|
|
261
|
+
ok_boxes.append(box)
|
|
262
|
+
break
|
|
263
|
+
|
|
264
|
+
yolo_result = YoloResult(names=self.names, boxes=ok_boxes)
|
|
265
|
+
save_debug_image(
|
|
266
|
+
image,
|
|
267
|
+
yolo_result,
|
|
268
|
+
page.page_number + 1,
|
|
269
|
+
)
|
|
270
|
+
yield page, yolo_result
|
|
271
|
+
|
|
272
|
+
def _is_box_in_table(self, box_xyxy, table_box, page, img_width, img_height):
|
|
273
|
+
"""
|
|
274
|
+
Check if a box from image coordinates is inside a table box from PDF coordinates.
|
|
275
|
+
|
|
276
|
+
Args:
|
|
277
|
+
box_xyxy (list): Box coordinates in image coordinate system [x1, y1, x2, y2]
|
|
278
|
+
table_box (Box): Table box in PDF coordinate system
|
|
279
|
+
page: The page object containing information for coordinate conversion
|
|
280
|
+
img_width: Width of the image
|
|
281
|
+
img_height: Height of the image
|
|
282
|
+
|
|
283
|
+
Returns:
|
|
284
|
+
bool: True if the box is inside or significantly overlapping with the table box
|
|
285
|
+
"""
|
|
286
|
+
|
|
287
|
+
# Get table box coordinates in PDF coordinate system
|
|
288
|
+
table_pdf_x1 = table_box.x
|
|
289
|
+
table_pdf_y1 = table_box.y
|
|
290
|
+
table_pdf_x2 = table_box.x2
|
|
291
|
+
table_pdf_y2 = table_box.y2
|
|
292
|
+
|
|
293
|
+
# Convert table box to image coordinates
|
|
294
|
+
table_img_x1 = table_pdf_x1
|
|
295
|
+
table_img_y1 = img_height - table_pdf_y2
|
|
296
|
+
table_img_x2 = table_pdf_x2
|
|
297
|
+
table_img_y2 = img_height - table_pdf_y1
|
|
298
|
+
|
|
299
|
+
# Now check for overlap between the boxes
|
|
300
|
+
# Calculate the area of overlap
|
|
301
|
+
x_overlap = max(
|
|
302
|
+
0, min(box_xyxy[2], table_img_x2) - max(box_xyxy[0], table_img_x1)
|
|
303
|
+
)
|
|
304
|
+
y_overlap = max(
|
|
305
|
+
0, min(box_xyxy[3], table_img_y2) - max(box_xyxy[1], table_img_y1)
|
|
306
|
+
)
|
|
307
|
+
overlap_area = x_overlap * y_overlap
|
|
308
|
+
|
|
309
|
+
# Calculate area of the detected box
|
|
310
|
+
box_area = (box_xyxy[2] - box_xyxy[0]) * (box_xyxy[3] - box_xyxy[1])
|
|
311
|
+
|
|
312
|
+
# If overlap area is significant relative to the box area, consider it inside
|
|
313
|
+
if box_area > 0 and overlap_area / box_area > 0.5:
|
|
314
|
+
return True
|
|
315
|
+
|
|
316
|
+
return False
|