doc-page-extractor 0.0.9__py3-none-any.whl → 0.1.0__py3-none-any.whl

doc_page_extractor/table.py

@@ -0,0 +1,71 @@
+import os
+import torch
+
+from typing import Literal, Any
+from PIL.Image import Image
+from .types import TableLayoutParsedFormat
+from .utils import expand_image
+
+
+OutputFormat = Literal["latex", "markdown", "html"]
+
+class Table:
+  def __init__(
+      self,
+      device: Literal["cpu", "cuda"],
+      model_path: str,
+    ):
+    self._model: Any | None = None
+    self._model_path: str = model_path
+    self._ban: bool = False
+    if device == "cpu" or not torch.cuda.is_available():
+      self._ban = True
+
+  def predict(self, image: Image, format: TableLayoutParsedFormat) -> str | None:
+    if self._ban:
+      print("CUDA is not available. You cannot parse table from image.")
+      return None
+
+    output_format: str
+    if format == TableLayoutParsedFormat.LATEX:
+      output_format = "latex"
+    elif format == TableLayoutParsedFormat.MARKDOWN:
+      output_format = "markdown"
+    elif format == TableLayoutParsedFormat.HTML:
+      output_format = "html"
+    else:
+      raise ValueError(f"Table format {format} is not supported.")
+
+    image = expand_image(image, 0.1)
+    model = self._get_model()
+
+    with torch.no_grad():
+      results = model([image], output_format=output_format)
+
+    if len(results) == 0:
+      return None
+
+    return results[0]
+
+  def _get_model(self):
+    if self._model is None:
+      local_files_only: bool
+      if os.path.exists(self._model_path):
+        local_files_only = True
+      else:
+        local_files_only = False
+        os.makedirs(self._model_path)
+
+      from .struct_eqtable import build_model
+      model = build_model(
+        model_ckpt="U4R/StructTable-InternVL2-1B",
+        max_new_tokens=1024,
+        max_time=30,
+        lmdeploy=False,
+        flash_attn=True,
+        batch_size=1,
+        cache_dir=self._model_path,
+        local_files_only=local_files_only,
+      )
+      self._model = model.cuda()
+    return self._model

doc_page_extractor/types.py

@@ -1,8 +1,10 @@
 from dataclasses import dataclass
-from enum import Enum
+from typing import Literal
+from enum import auto, Enum
 from PIL.Image import Image
 from .rectangle import Rectangle
 
+
 @dataclass
 class OCRFragment:
   order: int
@@ -22,12 +24,41 @@ class LayoutClass(Enum):
   ISOLATE_FORMULA = 8
   FORMULA_CAPTION = 9
 
+class TableLayoutParsedFormat(Enum):
+  LATEX = auto()
+  MARKDOWN = auto()
+  HTML = auto()
+
 @dataclass
-class Layout:
-  cls: LayoutClass
+class BaseLayout:
   rect: Rectangle
   fragments: list[OCRFragment]
 
+@dataclass
+class PlainLayout(BaseLayout):
+  cls: Literal[
+    LayoutClass.TITLE,
+    LayoutClass.PLAIN_TEXT,
+    LayoutClass.ABANDON,
+    LayoutClass.FIGURE,
+    LayoutClass.FIGURE_CAPTION,
+    LayoutClass.TABLE_CAPTION,
+    LayoutClass.TABLE_FOOTNOTE,
+    LayoutClass.FORMULA_CAPTION,
+  ]
+
+@dataclass
+class TableLayout(BaseLayout):
+  parsed: tuple[str, TableLayoutParsedFormat] | None
+  cls: LayoutClass.TABLE
+
+@dataclass
+class FormulaLayout(BaseLayout):
+  latex: str | None
+  cls: LayoutClass.ISOLATE_FORMULA
+
+Layout = PlainLayout | TableLayout | FormulaLayout
+
 @dataclass
 class ExtractedResult:
   rotation: float

doc_page_extractor/utils.py

@@ -1,10 +1,32 @@
 import os
 import re
 
+from math import ceil
+from PIL.Image import Image
+from PIL.ImageOps import expand
+
+
 def ensure_dir(path: str) -> str:
   path = os.path.abspath(path)
   os.makedirs(path, exist_ok=True)
   return path
 
 def is_space_text(text: str) -> bool:
-  return re.match(r"^\s*$", text)
+  return re.match(r"^\s*$", text)
+
+def expand_image(image: Image, percent: float):
+  width, height = image.size
+  border_width = ceil(width * percent)
+  border_height = ceil(height * percent)
+  fill_color: tuple[int, ...]
+
+  if image.mode == "RGBA":
+    fill_color = (255, 255, 255, 255)
+  else:
+    fill_color = (255, 255, 255)
+
+  return expand(
+    image=image,
+    border=(border_width, border_height),
+    fill=fill_color,
+  )

{doc_page_extractor-0.0.9.dist-info → doc_page_extractor-0.1.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: doc-page-extractor
-Version: 0.0.9
+Version: 0.1.0
 Summary: doc page extractor can identify text and format in images and return structured data.
 Home-page: https://github.com/Moskize91/doc-page-extractor
 Author: Tao Zeyu
@@ -12,8 +12,10 @@ Requires-Dist: pillow<11.0,>=10.3
 Requires-Dist: pyclipper<2.0,>=1.2.0
 Requires-Dist: numpy<2.0,>=1.24.0
 Requires-Dist: shapely<3.0,>=2.0.0
-Requires-Dist: transformers<5.0,>=4.48.0
+Requires-Dist: transformers<=4.47,>=4.42.4
 Requires-Dist: doclayout_yolo>=0.0.3
+Requires-Dist: pix2tex<=0.2.0,>=0.1.4
+Requires-Dist: accelerate<2.0,>=1.6.0
 Dynamic: author
 Dynamic: author-email
 Dynamic: description
@@ -78,3 +80,5 @@ The code of `doc_page_extractor/onnxocr` in this repo comes from [OnnxOCR](https
 - [DocLayout-YOLO](https://github.com/opendatalab/DocLayout-YOLO)
 - [OnnxOCR](https://github.com/jingsongliujing/OnnxOCR)
 - [layoutreader](https://github.com/ppaanngggg/layoutreader)
+- [StructEqTable](https://github.com/Alpha-Innovator/StructEqTable-Deploy)
+- [LaTeX-OCR](https://github.com/lukas-blecher/LaTeX-OCR)

{doc_page_extractor-0.0.9.dist-info → doc_page_extractor-0.1.0.dist-info}/RECORD

@@ -1,17 +1,20 @@
-doc_page_extractor/__init__.py,sha256=jCf5lo3A9JVDquflYMlvH8nJIs3EjBt8AG5y8mwfS68,210
+doc_page_extractor/__init__.py,sha256=9rWKSMTgzP7Xv15zA4upsyPaR8S8JeNpMyhWElRCW0M,311
 doc_page_extractor/clipper.py,sha256=5S1TI0aqMebwlPv_Ih4Nxpp6MchEjOih-CiZfMWUAhI,3201
 doc_page_extractor/downloader.py,sha256=NbGN9ARnER8-gd4T1uc3W98WMEClVxMrqnShq8HibTw,455
-doc_page_extractor/extractor.py,sha256=RPoqc_UtbHe4BoQvtaJxWuKyHwoDkK_fuYvuyA31p0Y,10200
+doc_page_extractor/extractor.py,sha256=njLl8VdOgm-noWPDYTfjIAUU1giNc-wLvCSR1pHkfS8,7267
+doc_page_extractor/latex.py,sha256=W_zAcksNRuru-WjCq4CSn07s_SWrDhikadJSy_Cg3Do,1954
+doc_page_extractor/layout_order.py,sha256=NwMzTPr4xsriz4slCwqwhw2-vrMu-qfwtcFsDu8d1yM,7426
 doc_page_extractor/layoutreader.py,sha256=BdC4oPbtpXoLmYhjuSFrKn6SNoT2zWw_gi95sGAUwrk,4031
-doc_page_extractor/ocr.py,sha256=tEalIlEz_80bJN8BXnHWi3UCMoPdJ5ylgkmaBV1dzVs,5201
+doc_page_extractor/ocr.py,sha256=KJ5PqtBa4_n8LAfMLGApUVNPUS1DBEwVKcC-zck283I,5161
 doc_page_extractor/ocr_corrector.py,sha256=RfRA1jESEuqC8_a2kUEvHblT_B4xBjE0OApLMl1JiRg,3917
-doc_page_extractor/overlap.py,sha256=Nx7Z6Ih_7SBE-JAjqlMx51PYiciXomPxIOAS37x6fco,5327
+doc_page_extractor/overlap.py,sha256=z1DF4_2OPvauDHwmz1SC1WosULkE84HKaRfNEgexPzc,5337
 doc_page_extractor/plot.py,sha256=4uibjS_x1SyEyjaJJd0YsBbzkgldDOCct4Ry2cOhdXU,2556
 doc_page_extractor/raw_optimizer.py,sha256=1KghECq_rJwuZZITTLQnGTKYivFKg_qDvMLN9g17sks,2844
 doc_page_extractor/rectangle.py,sha256=Tp__NPiY6JlYwYxejST7BUXhv_bl8tkmDXi4JgHCK6E,1539
 doc_page_extractor/rotation.py,sha256=QCZ-HqfDxIhnQw8KRHki2myj6-UusvNY7Mpjsu-wI-4,4334
-doc_page_extractor/types.py,sha256=UWghDwajMtEKEYUcOOjr5dM-MWiJ-P-8nWYagouf9ds,631
-doc_page_extractor/utils.py,sha256=3rtIxiTJ7W5yOuY0UHedUJ3G34tPOw0jdHnUdOQ1tWI,207
+doc_page_extractor/table.py,sha256=AWymTRbOet55uImW8QJqb90Qs_v2V2U1mZv0U6rSz3c,1891
+doc_page_extractor/types.py,sha256=7blT8YNKrOsc4qQdAhM7J7MEQjFcBwE0QV8-lipZBeQ,1305
+doc_page_extractor/utils.py,sha256=ZlQVOLPUg_v5J8u6SoD8XtMG_JkF-ERgjubc4LO5_Lg,688
 doc_page_extractor/onnxocr/__init__.py,sha256=BK4YpX4pU0nRxbcI5f5cbIVfdBEsx4W980QYmpNQaH0,38
 doc_page_extractor/onnxocr/cls_postprocess.py,sha256=o879Ned0RMUERYLviuToZ0xTvhn2UsYAb-yPC5gj8h4,822
 doc_page_extractor/onnxocr/db_postprocess.py,sha256=R3yXXfReiQgLaYIvvfnrFfshI202LjHMvcZwcLpjmTY,7913
@@ -24,10 +27,18 @@ doc_page_extractor/onnxocr/predict_rec.py,sha256=qQrCs5jzCf5PYp-iEKJ53pcx_xRoJdJ
 doc_page_extractor/onnxocr/predict_system.py,sha256=yoqXunAsoboPsWe7qQjvQf2_SMW1T1QMriEoiGdX3BM,2721
 doc_page_extractor/onnxocr/rec_postprocess.py,sha256=qZt5Ripal7z9hniKq5e7azOkD9e6NR1ylWpRpznhweg,29556
 doc_page_extractor/onnxocr/utils.py,sha256=AQoHgQyv-jpPo4BsVzq3r7_ze698EZ-a7LJobm2fwUI,1864
-doc_page_extractor-0.0.9.dist-info/licenses/LICENSE,sha256=TfPDBt3ar0uv_f9cqCDMZ5rIzW3CY8anRRd4PkL6ejs,34522
+doc_page_extractor/struct_eqtable/__init__.py,sha256=QoTsNuJfpNSrMIMd6Cot1jJqWk88_lDqFP_C2rcVJO4,1329
+doc_page_extractor/struct_eqtable/internvl/__init__.py,sha256=2aOsU-aHkFv_gjdP8LeUXjj_9-0d4x79iyxh4cCzaEw,79
+doc_page_extractor/struct_eqtable/internvl/conversation.py,sha256=s7DceRlM6JtHmxgyuE6vqu5XVT1fHzhzCL_I6r8MI1c,15129
+doc_page_extractor/struct_eqtable/internvl/internvl.py,sha256=ovVZG-PuBrsj_9lEoNPOygJ-2en3v6gPzRfWjDpWNOM,7678
+doc_page_extractor/struct_eqtable/internvl/internvl_lmdeploy.py,sha256=ACHxFntxS38G43PzE955Nv4fjKk_-Oz4y_o9JEjQwlg,2608
+doc_page_extractor/struct_eqtable/pix2s/__init__.py,sha256=cXRo4eg6u1-TXktZ8rQf0HIzLmmScIwYQhbxMKl-MyA,76
+doc_page_extractor/struct_eqtable/pix2s/pix2s.py,sha256=fCNve8PNeJ3-AWJIhSeGtp-mYKoMXfW0CIpszkQnAaA,2535
+doc_page_extractor/struct_eqtable/pix2s/pix2s_trt.py,sha256=zSGw45JhWdZ3iuJel5Chsy-NzsOHx9QyPQIUAzzCjFE,43880
+doc_page_extractor-0.1.0.dist-info/licenses/LICENSE,sha256=TfPDBt3ar0uv_f9cqCDMZ5rIzW3CY8anRRd4PkL6ejs,34522
 tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tests/test_history_bus.py,sha256=WaCUW3U75SESMcLq0f5FKnpVUVRDvmfxLFE7Zo83e48,2517
-doc_page_extractor-0.0.9.dist-info/METADATA,sha256=vwcUU9DE9rUhpbz1UigpDN3YhNnNSRbHGAyJroHPV7Q,2225
-doc_page_extractor-0.0.9.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-doc_page_extractor-0.0.9.dist-info/top_level.txt,sha256=ErNybD_lBzAmw8mVBAK4htsAH_hp14jioZVex-tUqvM,25
-doc_page_extractor-0.0.9.dist-info/RECORD,,
+doc_page_extractor-0.1.0.dist-info/METADATA,sha256=8AM05x2gY75j70rD2HNi5zyaKE1okNJG7Pw5iLuIlnQ,2436
+doc_page_extractor-0.1.0.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+doc_page_extractor-0.1.0.dist-info/top_level.txt,sha256=ErNybD_lBzAmw8mVBAK4htsAH_hp14jioZVex-tUqvM,25
+doc_page_extractor-0.1.0.dist-info/RECORD,,