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.
Files changed (95) hide show
  1. {babeldoc-0.2.13 → babeldoc-0.2.15}/PKG-INFO +3 -1
  2. {babeldoc-0.2.13 → babeldoc-0.2.15}/README.md +1 -0
  3. babeldoc-0.2.15/babeldoc/__init__.py +1 -0
  4. {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/assets/assets.py +35 -1
  5. {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/assets/embedding_assets_metadata.py +10 -0
  6. {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/const.py +1 -1
  7. {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/backend/pdf_creater.py +6 -1
  8. {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/midend/paragraph_finder.py +2 -0
  9. babeldoc-0.2.15/babeldoc/document_il/midend/table_parser.py +164 -0
  10. babeldoc-0.2.15/babeldoc/docvision/table_detection/rapidocr.py +316 -0
  11. {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/high_level.py +37 -23
  12. {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/main.py +14 -1
  13. {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/pdfinterp.py +7 -7
  14. {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/translation_config.py +3 -0
  15. {babeldoc-0.2.13 → babeldoc-0.2.15}/pyproject.toml +3 -2
  16. {babeldoc-0.2.13 → babeldoc-0.2.15}/uv.lock +100 -7
  17. babeldoc-0.2.13/babeldoc/__init__.py +0 -1
  18. {babeldoc-0.2.13 → babeldoc-0.2.15}/.cursorignore +0 -0
  19. {babeldoc-0.2.13 → babeldoc-0.2.15}/.github/ISSUE_TEMPLATE/bug_report.yaml +0 -0
  20. {babeldoc-0.2.13 → babeldoc-0.2.15}/.github/ISSUE_TEMPLATE/feature_request.yaml +0 -0
  21. {babeldoc-0.2.13 → babeldoc-0.2.15}/.github/dependabot.yml +0 -0
  22. {babeldoc-0.2.13 → babeldoc-0.2.15}/.github/labels.yml +0 -0
  23. {babeldoc-0.2.13 → babeldoc-0.2.15}/.github/release-drafter.yml +0 -0
  24. {babeldoc-0.2.13 → babeldoc-0.2.15}/.github/workflows/codeql.yml +0 -0
  25. {babeldoc-0.2.13 → babeldoc-0.2.15}/.github/workflows/docs.yml +0 -0
  26. {babeldoc-0.2.13 → babeldoc-0.2.15}/.github/workflows/labeler.yml +0 -0
  27. {babeldoc-0.2.13 → babeldoc-0.2.15}/.github/workflows/publish-to-pypi.yml +0 -0
  28. {babeldoc-0.2.13 → babeldoc-0.2.15}/.github/workflows/test.yml +0 -0
  29. {babeldoc-0.2.13 → babeldoc-0.2.15}/.gitignore +0 -0
  30. {babeldoc-0.2.13 → babeldoc-0.2.15}/.pre-commit-config.yaml +0 -0
  31. {babeldoc-0.2.13 → babeldoc-0.2.15}/LICENSE +0 -0
  32. {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/asynchronize/__init__.py +0 -0
  33. {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/converter.py +0 -0
  34. {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/__init__.py +0 -0
  35. {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/babeldoc_exception/BabelDOCException.py +0 -0
  36. {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/backend/__init__.py +0 -0
  37. {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/frontend/__init__.py +0 -0
  38. {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/frontend/il_creater.py +0 -0
  39. {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/il_version_1.py +0 -0
  40. {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/il_version_1.rnc +0 -0
  41. {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/il_version_1.rng +0 -0
  42. {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/il_version_1.xsd +0 -0
  43. {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/midend/__init__.py +0 -0
  44. {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/midend/add_debug_information.py +0 -0
  45. {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/midend/detect_scanned_file.py +0 -0
  46. {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/midend/il_translator.py +0 -0
  47. {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/midend/il_translator_llm_only.py +0 -0
  48. {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/midend/layout_parser.py +0 -0
  49. {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/midend/remove_descent.py +0 -0
  50. {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/midend/styles_and_formulas.py +0 -0
  51. {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/midend/typesetting.py +0 -0
  52. {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/translator/__init__.py +0 -0
  53. {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/translator/cache.py +0 -0
  54. {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/translator/translator.py +0 -0
  55. {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/utils/fontmap.py +0 -0
  56. {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/utils/layout_helper.py +0 -0
  57. {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/utils/style_helper.py +0 -0
  58. {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/document_il/xml_converter.py +0 -0
  59. {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/docvision/README.md +0 -0
  60. {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/docvision/__init__.py +0 -0
  61. {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/docvision/doclayout.py +0 -0
  62. {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/docvision/rpc_doclayout.py +0 -0
  63. {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/progress_monitor.py +0 -0
  64. {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/result_merger.py +0 -0
  65. {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/split_manager.py +0 -0
  66. {babeldoc-0.2.13 → babeldoc-0.2.15}/babeldoc/tools/generate_font_metadata.py +0 -0
  67. {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/CODE_OF_CONDUCT.md +0 -0
  68. {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/CONTRIBUTING.md +0 -0
  69. {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/CONTRIBUTOR_REWARD.md +0 -0
  70. {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/ImplementationDetails/AsyncTranslate/AsyncTranslate.md +0 -0
  71. {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/ImplementationDetails/ILTranslator/ILTranslator.md +0 -0
  72. {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/ImplementationDetails/PDFCreation/PDFCreation.md +0 -0
  73. {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/ImplementationDetails/PDFParsing/PDFParsing.md +0 -0
  74. {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/ImplementationDetails/ParagraphFinding/ParagraphFinding.md +0 -0
  75. {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/ImplementationDetails/README.md +0 -0
  76. {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/ImplementationDetails/StylesAndFormulas/StylesAndFormulas.md +0 -0
  77. {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/ImplementationDetails/Typesetting/Typesetting.md +0 -0
  78. {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/README.md +0 -0
  79. {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/deploy.sh +0 -0
  80. {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/images/babeldoc-banner.png +0 -0
  81. {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/images/babeldoc-big-logo-darkmode-with-transparent-background.png +0 -0
  82. {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/images/babeldoc-big-logo-darkmode-with-transparent-background.svg +0 -0
  83. {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/images/babeldoc-big-logo-with-transparent-background.png +0 -0
  84. {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/images/babeldoc-big-logo-with-transparent-background.svg +0 -0
  85. {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/images/babeldoc-big-logo.png +0 -0
  86. {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/images/babeldoc-contributor_reward_example.png +0 -0
  87. {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/images/babeldoc-preview.gif +0 -0
  88. {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/images/babeldoc-small-logo-with-transparent-background.png +0 -0
  89. {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/images/babeldoc-small-logo-with-transparent-background.svg +0 -0
  90. {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/images/babeldoc-small-logo.png +0 -0
  91. {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/index.md +0 -0
  92. {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/intro-to-pdf-object.md +0 -0
  93. {babeldoc-0.2.13 → babeldoc-0.2.15}/docs/requirements.txt +0 -0
  94. {babeldoc-0.2.13 → babeldoc-0.2.15}/mkdocs.yml +0 -0
  95. {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.13
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": {
@@ -3,7 +3,7 @@ import shutil
3
3
  import subprocess
4
4
  from pathlib import Path
5
5
 
6
- __version__ = "0.2.13"
6
+ __version__ = "0.2.15"
7
7
 
8
8
  CACHE_FOLDER = Path.home() / ".cache" / "babeldoc"
9
9
 
@@ -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
- pdf.update_stream(xobj.xref_id, draw_op.tobytes())
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