doc-page-extractor 0.0.10__tar.gz → 0.1.1__tar.gz

{doc_page_extractor-0.0.10 → doc_page_extractor-0.1.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: doc-page-extractor
-Version: 0.0.10
+Version: 0.1.1
 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.10 → doc_page_extractor-0.1.1}/README.md

@@ -52,4 +52,6 @@ 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)
+- [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.1.1/doc_page_extractor/__init__.py

@@ -0,0 +1,15 @@
+from .extractor import DocExtractor
+from .clipper import clip, clip_from_image
+from .plot import plot
+from .rectangle import Point, Rectangle
+from .types import (
+  ExtractedResult,
+  OCRFragment,
+  LayoutClass,
+  TableLayoutParsedFormat,
+  Layout,
+  BaseLayout,
+  PlainLayout,
+  FormulaLayout,
+  TableLayout,
+)

{doc_page_extractor-0.0.10 → doc_page_extractor-0.1.1}/doc_page_extractor/extractor.py

@@ -1,6 +1,6 @@
 import os
 
-from typing import Literal
+from typing import Literal, Generator
 from pathlib import Path
 from PIL.Image import Image
 from doclayout_yolo import YOLOv10
@@ -9,10 +9,22 @@ from .ocr import OCR
 from .ocr_corrector import correct_fragments
 from .raw_optimizer import RawOptimizer
 from .rectangle import intersection_area, Rectangle
-from .types import ExtractedResult, OCRFragment, LayoutClass, Layout
 from .downloader import download
+from .table import Table
+from .latex import LaTeX
 from .layout_order import LayoutOrder
 from .overlap import merge_fragments_as_line, remove_overlap_layouts
+from .clipper import clip_from_image
+from .types import (
+  ExtractedResult,
+  OCRFragment,
+  TableLayoutParsedFormat,
+  Layout,
+  LayoutClass,
+  PlainLayout,
+  TableLayout,
+  FormulaLayout,
+)
 
 
 class DocExtractor:
@@ -21,12 +33,26 @@ class DocExtractor:
       model_dir_path: str,
       device: Literal["cpu", "cuda"] = "cpu",
       ocr_for_each_layouts: bool = True,
+      extract_formula: bool = True,
+      extract_table_format: TableLayoutParsedFormat | None = None,
     ):
     self._model_dir_path: str = model_dir_path
     self._device: Literal["cpu", "cuda"] = device
     self._ocr_for_each_layouts: bool = ocr_for_each_layouts
-    self._ocr: OCR = OCR(device, model_dir_path)
+    self._extract_formula: bool = extract_formula
+    self._extract_table_format: TableLayoutParsedFormat | None = extract_table_format
     self._yolo: YOLOv10 | None = None
+    self._ocr: OCR = OCR(
+      device=device,
+      model_dir_path=os.path.join(model_dir_path, "onnx_ocr"),
+    )
+    self._table: Table = Table(
+      device=device,
+      model_path=os.path.join(model_dir_path, "struct_eqtable"),
+    )
+    self._latex: LaTeX = LaTeX(
+      model_path=os.path.join(model_dir_path, "latex"),
+    )
     self._layout_order: LayoutOrder = LayoutOrder(
       model_path=os.path.join(model_dir_path, "layoutreader"),
     )
@@ -40,7 +66,7 @@ class DocExtractor:
     raw_optimizer = RawOptimizer(image, adjust_points)
     fragments = list(self._ocr.search_fragments(raw_optimizer.image_np))
     raw_optimizer.receive_raw_fragments(fragments)
-    layouts = self._get_layouts(raw_optimizer.image)
+    layouts = list(self._yolo_extract_layouts(raw_optimizer.image))
     layouts = self._layouts_matched_by_fragments(fragments, layouts)
     layouts = remove_overlap_layouts(layouts)
 
@@ -50,6 +76,8 @@ class DocExtractor:
     layouts = self._layout_order.sort(layouts, raw_optimizer.image.size)
     layouts = [layout for layout in layouts if self._should_keep_layout(layout)]
 
+    self._parse_table_and_formula_layouts(layouts, raw_optimizer)
+
     for layout in layouts:
       layout.fragments = merge_fragments_as_line(layout.fragments)
 
@@ -62,7 +90,7 @@ class DocExtractor:
       adjusted_image=raw_optimizer.adjusted_image,
     )
 
-  def _get_layouts(self, source: Image) -> list[Layout]:
+  def _yolo_extract_layouts(self, source: Image) -> Generator[Layout, None, None]:
     # about source parameter to see:
     # https://github.com/opendatalab/DocLayout-YOLO/blob/7c4be36bc61f11b67cf4a44ee47f3c41e9800a91/doclayout_yolo/data/build.py#L157-L175
     det_res = self._get_yolo().predict(
@@ -72,7 +100,6 @@ class DocExtractor:
       device=self._device    # Device to use (e.g., "cuda" or "cpu")
     )
     boxes = det_res[0].__dict__["boxes"]
-    layouts: list[Layout] = []
 
     for cls_id, rect in zip(boxes.cls, boxes.xyxy):
       cls_id = cls_id.item()
@@ -89,9 +116,13 @@ class DocExtractor:
         lb=(x1, y2),
         rb=(x2, y2),
       )
-      layouts.append(Layout(cls, rect, []))
-
-    return layouts
+      if rect.is_valid:
+        if cls == LayoutClass.TABLE:
+          yield TableLayout(cls=cls, rect=rect, fragments=[], parsed=None)
+        elif cls == LayoutClass.ISOLATE_FORMULA:
+          yield FormulaLayout(cls=cls, rect=rect, fragments=[], latex=None)
+        else:
+          yield PlainLayout(cls=cls, rect=rect, fragments=[])
 
   def _layouts_matched_by_fragments(self, fragments: list[OCRFragment], layouts: list[Layout]):
     layouts_group = self._split_layouts_by_group(layouts)
@@ -107,6 +138,17 @@ class DocExtractor:
     for layout in layouts:
       correct_fragments(self._ocr, source, layout)
 
+  def _parse_table_and_formula_layouts(self, layouts: list[Layout], raw_optimizer: RawOptimizer):
+    for layout in layouts:
+      if isinstance(layout, FormulaLayout) and self._extract_formula:
+        image = clip_from_image(raw_optimizer.image, layout.rect)
+        layout.latex = self._latex.extract(image)
+      elif isinstance(layout, TableLayout) and self._extract_table_format is not None:
+        image = clip_from_image(raw_optimizer.image, layout.rect)
+        parsed = self._table.predict(image, self._extract_table_format)
+        if parsed is not None:
+          layout.parsed = (parsed, self._extract_table_format)
+
   def _split_layouts_by_group(self, layouts: list[Layout]):
     texts_layouts: list[Layout] = []
     abandon_layouts: list[Layout] = []
@@ -149,9 +191,11 @@ class DocExtractor:
 
   def _get_yolo(self) -> YOLOv10:
     if self._yolo is None:
+      base_path = os.path.join(self._model_dir_path, "yolo")
+      os.makedirs(base_path, exist_ok=True)
       yolo_model_url = "https://huggingface.co/opendatalab/PDF-Extract-Kit-1.0/resolve/main/models/Layout/YOLO/doclayout_yolo_ft.pt"
       yolo_model_name = "doclayout_yolo_ft.pt"
-      yolo_model_path = Path(os.path.join(self._model_dir_path, yolo_model_name))
+      yolo_model_path = Path(os.path.join(base_path, yolo_model_name))
       if not yolo_model_path.exists():
         download(yolo_model_url, yolo_model_path)
       self._yolo = YOLOv10(str(yolo_model_path))

doc_page_extractor-0.1.1/doc_page_extractor/latex.py

@@ -0,0 +1,57 @@
+import os
+import torch
+import requests
+
+from munch import Munch
+from pix2tex.cli import LatexOCR
+from PIL.Image import Image
+from .utils import expand_image
+
+
+class LaTeX:
+  def __init__(self, model_path: str):
+    self._model_path: str = model_path
+    self._model: LatexOCR | None = None
+
+  def extract(self, image: Image) -> str | None:
+    image = expand_image(image, 0.1) # 添加边缘提高识别准确率
+    model = self._get_model()
+    with torch.no_grad():
+      return model(image)
+
+  def _get_model(self) -> LatexOCR:
+    if self._model is None:
+      if not os.path.exists(self._model_path):
+        self._download_model()
+
+      self._model = LatexOCR(Munch({
+        "config": os.path.join("settings", "config.yaml"),
+        "checkpoint": os.path.join(self._model_path, "weights.pth"),
+        "no_cuda": True,
+        "no_resize": False,
+      }))
+    return self._model
+
+  # from https://github.com/lukas-blecher/LaTeX-OCR/blob/5c1ac929bd19a7ecf86d5fb8d94771c8969fcb80/pix2tex/model/checkpoints/get_latest_checkpoint.py#L37-L45
+  def _download_model(self):
+    os.makedirs(self._model_path, exist_ok=True)
+    tag = "v0.0.1"
+    files: list[tuple[str, str]] = (
+      ("weights.pth", f"https://github.com/lukas-blecher/LaTeX-OCR/releases/download/{tag}/weights.pth"),
+      ("image_resizer.pth", f"https://github.com/lukas-blecher/LaTeX-OCR/releases/download/{tag}/image_resizer.pth")
+    )
+    for file_name, url in files:
+      file_path = os.path.join(self._model_path, file_name)
+      try:
+        with open(file_path, "wb") as file:
+          response = requests.get(url, stream=True, timeout=15)
+          response.raise_for_status()
+          for chunk in response.iter_content(chunk_size=8192):
+            if chunk:  # 过滤掉保持连接的新块
+              file.write(chunk)
+          file.flush()
+
+      except BaseException as e:
+        if os.path.exists(file_path):
+          os.remove(file_path)
+        raise e

{doc_page_extractor-0.0.10 → doc_page_extractor-0.1.1}/doc_page_extractor/ocr.py

@@ -69,7 +69,7 @@ class OCR:
         rb=(box[2][0], box[2][1]),
         lb=(box[3][0], box[3][1]),
       )
-      if rect.area == 0.0:
+      if not rect.is_valid or rect.area == 0.0:
         continue
 
       yield OCRFragment(

{doc_page_extractor-0.0.10 → doc_page_extractor-0.1.1}/doc_page_extractor/rectangle.py

@@ -19,6 +19,10 @@ class Rectangle:
     yield self.rb
     yield self.rt
 
+  @property
+  def is_valid(self) -> bool:
+    return Polygon(self).is_valid
+
   @property
   def segments(self) -> Generator[tuple[Point, Point], None, None]:
     yield (self.lt, self.lb)
@@ -60,6 +64,8 @@ class Rectangle:
 def intersection_area(rect1: Rectangle, rect2: Rectangle) -> float:
   poly1 = Polygon(rect1)
   poly2 = Polygon(rect2)
+  if not poly1.is_valid or not poly2.is_valid:
+    return 0.0
   intersection = poly1.intersection(poly2)
   if intersection.is_empty:
     return 0.0

doc_page_extractor-0.1.1/doc_page_extractor/struct_eqtable/__init__.py

@@ -0,0 +1,49 @@
+from .pix2s import Pix2Struct, Pix2StructTensorRT
+from .internvl import InternVL, InternVL_LMDeploy
+
+from transformers import AutoConfig
+
+
+__ALL_MODELS__ = {
+    'Pix2Struct': Pix2Struct,
+    'Pix2StructTensorRT': Pix2StructTensorRT,
+    'InternVL': InternVL,
+    'InternVL_LMDeploy': InternVL_LMDeploy,
+}
+
+
+def get_model_name(model_path):
+    model_config = AutoConfig.from_pretrained(
+        model_path,
+        trust_remote_code=True,
+    )
+
+    if 'Pix2Struct' in model_config.architectures[0]:
+        model_name = 'Pix2Struct'
+    elif 'InternVL' in model_config.architectures[0]:
+        model_name = 'InternVL'
+    else:
+        raise ValueError(f"Unsupported model type: {model_config.architectures[0]}")
+
+    return model_name
+
+
+def build_model(
+        model_ckpt='U4R/StructTable-InternVL2-1B',
+        cache_dir=None,
+        local_files_only=None,
+        **kwargs,
+    ):
+    model_name = get_model_name(model_ckpt)
+    if model_name == 'InternVL' and kwargs.get('lmdeploy', False):
+        model_name = 'InternVL_LMDeploy'
+    elif model_name == 'Pix2Struct' and kwargs.get('tensorrt_path', None):
+        model_name = 'Pix2StructTensorRT'
+
+    model = __ALL_MODELS__[model_name](
+        model_ckpt,
+        cache_dir=cache_dir,
+        local_files_only=local_files_only,
+        **kwargs
+    )
+    return model

doc_page_extractor-0.1.1/doc_page_extractor/struct_eqtable/internvl/__init__.py

@@ -0,0 +1,2 @@
+from .internvl import InternVL
+from .internvl_lmdeploy import InternVL_LMDeploy