doc-page-extractor 0.2.4__cp310-cp310-macosx_15_0_arm64.whl

doc_page_extractor/table.py

@@ -0,0 +1,60 @@
+import torch
+
+from typing import Literal, Any
+from PIL.Image import Image
+from .types import TableLayoutParsedFormat
+from .model import Model
+from .utils import expand_image
+
+
+OutputFormat = Literal["latex", "markdown", "html"]
+
+class Table:
+  def __init__(self, device: Literal["cpu", "cuda"], model: Model) -> None:
+    self._model: Model = model
+    self._table_model: Any | None = None
+    self._ban: bool = False
+    if device == "cpu":
+      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) -> Any:
+    if self._table_model is None:
+      from .struct_eqtable import build_model
+      model_path = self._model.get_struct_eqtable_path()
+      table_model = build_model(
+        model_ckpt=str(model_path),
+        max_new_tokens=1024,
+        max_time=30,
+        lmdeploy=False,
+        flash_attn=True,
+        batch_size=1,
+        local_files_only=False,
+      )
+      self._table_model = table_model.cuda()
+    return self._table_model

doc_page_extractor/types.py

@@ -0,0 +1,68 @@
+from dataclasses import dataclass
+from typing import Literal
+from enum import auto, Enum
+from PIL.Image import Image
+from .rectangle import Rectangle
+
+
+@dataclass
+class OCRFragment:
+  order: int
+  text: str
+  rank: float
+  rect: Rectangle
+
+class LayoutClass(Enum):
+  TITLE = 0
+  PLAIN_TEXT = 1
+  ABANDON = 2
+  FIGURE = 3
+  FIGURE_CAPTION = 4
+  TABLE = 5
+  TABLE_CAPTION = 6
+  TABLE_FOOTNOTE = 7
+  ISOLATE_FORMULA = 8
+  FORMULA_CAPTION = 9
+
+class TableLayoutParsedFormat(Enum):
+  LATEX = auto()
+  MARKDOWN = auto()
+  HTML = auto()
+
+@dataclass
+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: Literal[LayoutClass.TABLE]
+
+@dataclass
+class FormulaLayout(BaseLayout):
+  latex: str | None
+  cls: Literal[LayoutClass.ISOLATE_FORMULA]
+
+Layout = PlainLayout | TableLayout | FormulaLayout
+
+
+@dataclass
+class ExtractedResult:
+  rotation: float
+  layouts: list[Layout]
+  extracted_image: Image | None
+  adjusted_image: Image | None

doc_page_extractor/utils.py

@@ -0,0 +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 bool(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,
+  )