doc-page-extractor 0.1.0__tar.gz → 0.1.2__tar.gz

{doc_page_extractor-0.1.0 → doc_page_extractor-0.1.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: doc-page-extractor
-Version: 0.1.0
+Version: 0.1.2
 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
@@ -16,6 +16,7 @@ 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
+Requires-Dist: huggingface_hub>=0.30.2
 Dynamic: author
 Dynamic: author-email
 Dynamic: description

{doc_page_extractor-0.1.0 → doc_page_extractor-0.1.2}/doc_page_extractor/__init__.py

@@ -12,4 +12,5 @@ from .types import (
   PlainLayout,
   FormulaLayout,
   TableLayout,
-)
+  ModelsDownloader
+)

{doc_page_extractor-0.1.0 → doc_page_extractor-0.1.2}/doc_page_extractor/extractor.py

@@ -1,15 +1,14 @@
-import os
 
 from typing import Literal, Generator
-from pathlib import Path
 from PIL.Image import Image
 from doclayout_yolo import YOLOv10
+from logging import Logger, getLogger
 
+from .models import HuggingfaceModelsDownloader
 from .ocr import OCR
 from .ocr_corrector import correct_fragments
 from .raw_optimizer import RawOptimizer
 from .rectangle import intersection_area, Rectangle
-from .downloader import download
 from .table import Table
 from .latex import LaTeX
 from .layout_order import LayoutOrder
@@ -17,6 +16,7 @@ from .overlap import merge_fragments_as_line, remove_overlap_layouts
 from .clipper import clip_from_image
 from .types import (
   ExtractedResult,
+  ModelsDownloader,
   OCRFragment,
   TableLayoutParsedFormat,
   Layout,
@@ -30,13 +30,17 @@ from .types import (
 class DocExtractor:
   def __init__(
       self,
-      model_dir_path: str,
+      model_cache_dir: str | None = None,
       device: Literal["cpu", "cuda"] = "cpu",
       ocr_for_each_layouts: bool = True,
       extract_formula: bool = True,
       extract_table_format: TableLayoutParsedFormat | None = None,
+      models_downloader: ModelsDownloader | None = None,
+      logger: Logger | None = None,
     ):
-    self._model_dir_path: str = model_dir_path
+    self._logger = logger or getLogger(__name__)
+    self._models_downloader = models_downloader or HuggingfaceModelsDownloader(self._logger, model_cache_dir)
+
     self._device: Literal["cpu", "cuda"] = device
     self._ocr_for_each_layouts: bool = ocr_for_each_layouts
     self._extract_formula: bool = extract_formula
@@ -44,17 +48,17 @@ class DocExtractor:
     self._yolo: YOLOv10 | None = None
     self._ocr: OCR = OCR(
       device=device,
-      model_dir_path=os.path.join(model_dir_path, "onnx_ocr"),
+      get_model_dir=self._models_downloader.onnx_ocr,
     )
     self._table: Table = Table(
       device=device,
-      model_path=os.path.join(model_dir_path, "struct_eqtable"),
+      get_model_dir=self._models_downloader.struct_eqtable,
     )
     self._latex: LaTeX = LaTeX(
-      model_path=os.path.join(model_dir_path, "latex"),
+      get_model_dir=self._models_downloader.latex,
     )
     self._layout_order: LayoutOrder = LayoutOrder(
-      model_path=os.path.join(model_dir_path, "layoutreader"),
+      get_model_dir=self._models_downloader.layoutreader,
     )
 
   def extract(
@@ -116,12 +120,13 @@ class DocExtractor:
         lb=(x1, y2),
         rb=(x2, y2),
       )
-      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=[])
+      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)
@@ -190,14 +195,8 @@ 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(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))
+      model_path = self._models_downloader.yolo()
+      self._yolo = YOLOv10(str(model_path))
     return self._yolo
 
   def _should_keep_layout(self, layout: Layout) -> bool:

doc_page_extractor-0.1.2/doc_page_extractor/latex.py

@@ -0,0 +1,29 @@
+import os
+import torch
+
+from munch import Munch
+from pix2tex.cli import LatexOCR
+from PIL.Image import Image
+from .utils import expand_image
+from .types import GetModelDir
+
+class LaTeX:
+  def __init__(self, get_model_dir: GetModelDir):
+    self._model_path: str = get_model_dir()
+    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:
+      self._model = LatexOCR(Munch({
+        "config": os.path.join("settings", "config.yaml"),
+        "checkpoint": os.path.join(self._model_path, "checkpoints", "weights.pth"),
+        "no_cuda": True,
+        "no_resize": False,
+      }))
+    return self._model

{doc_page_extractor-0.1.0 → doc_page_extractor-0.1.2}/doc_page_extractor/layout_order.py

@@ -1,13 +1,11 @@
-import os
 import torch
 
 from typing import Generator
 from dataclasses import dataclass
 from transformers import LayoutLMv3ForTokenClassification
 
-from .types import Layout, LayoutClass
+from .types import Layout, LayoutClass, GetModelDir
 from .layoutreader import prepare_inputs, boxes2inputs, parse_logits
-from .utils import ensure_dir
 
 
 @dataclass
@@ -19,17 +17,15 @@ class _BBox:
   value: tuple[float, float, float, float]
 
 class LayoutOrder:
-  def __init__(self, model_path: str):
-    self._model_path: str = model_path
+  def __init__(self, get_model_dir: GetModelDir):
+    self._model_path: str = get_model_dir()
     self._model: LayoutLMv3ForTokenClassification | None = None
 
   def _get_model(self) -> LayoutLMv3ForTokenClassification:
     if self._model is None:
-      model_path = ensure_dir(self._model_path)
       self._model = LayoutLMv3ForTokenClassification.from_pretrained(
-        pretrained_model_name_or_path="hantian/layoutreader",
-        cache_dir=model_path,
-        local_files_only=os.path.exists(os.path.join(model_path, "models--hantian--layoutreader")),
+        pretrained_model_name_or_path=self._model_path,
+        local_files_only=True,
       )
     return self._model
 

doc_page_extractor-0.1.2/doc_page_extractor/models.py

@@ -0,0 +1,92 @@
+import os
+
+from logging import Logger
+from huggingface_hub import hf_hub_download, snapshot_download, try_to_load_from_cache
+from .types import ModelsDownloader
+
+class HuggingfaceModelsDownloader(ModelsDownloader):
+  def __init__(
+      self,
+      logger: Logger,
+      model_dir_path: str | None
+    ):
+    self._logger = logger
+    self._model_dir_path: str | None = model_dir_path
+
+  def onnx_ocr(self) -> str:
+    repo_path = try_to_load_from_cache(
+      repo_id="moskize/OnnxOCR",
+      filename="README.md",
+      cache_dir=self._model_dir_path
+    )
+    if isinstance(repo_path, str):
+      return os.path.dirname(repo_path)
+    else:
+      self._logger.info("Downloading OCR model...")
+      return snapshot_download(
+        cache_dir=self._model_dir_path,
+        repo_id="moskize/OnnxOCR",
+      )
+
+  def yolo(self) -> str:
+    yolo_file_path = try_to_load_from_cache(
+      repo_id="opendatalab/PDF-Extract-Kit-1.0",
+      filename="models/Layout/YOLO/doclayout_yolo_ft.pt",
+      cache_dir=self._model_dir_path
+    )
+    if isinstance(yolo_file_path, str):
+      return yolo_file_path
+    else:
+      self._logger.info("Downloading YOLO model...")
+      return hf_hub_download(
+        cache_dir=self._model_dir_path,
+        repo_id="opendatalab/PDF-Extract-Kit-1.0",
+        filename="models/Layout/YOLO/doclayout_yolo_ft.pt",
+      )
+
+  def layoutreader(self) -> str:
+    repo_path = try_to_load_from_cache(
+      repo_id="hantian/layoutreader",
+      filename="model.safetensors",
+      cache_dir=self._model_dir_path
+    )
+    if isinstance(repo_path, str):
+      return os.path.dirname(repo_path)
+    else:
+      self._logger.info("Downloading LayoutReader model...")
+      return snapshot_download(
+        cache_dir=self._model_dir_path,
+        repo_id="hantian/layoutreader",
+      )
+
+  def struct_eqtable(self) -> str:
+    repo_path = try_to_load_from_cache(
+      repo_id="U4R/StructTable-InternVL2-1B",
+      filename="model.safetensors",
+      cache_dir=self._model_dir_path
+    )
+    if isinstance(repo_path, str):
+      return os.path.dirname(repo_path)
+    else:
+      self._logger.info("Downloading StructEqTable model...")
+      return snapshot_download(
+        cache_dir=self._model_dir_path,
+        repo_id="U4R/StructTable-InternVL2-1B",
+      )
+
+  def latex(self):
+    repo_path = try_to_load_from_cache(
+      repo_id="lukbl/LaTeX-OCR",
+      filename="checkpoints/weights.pth",
+      repo_type="space",
+      cache_dir=self._model_dir_path
+    )
+    if isinstance(repo_path, str):
+      return os.path.dirname(os.path.dirname(repo_path))
+    else:
+      self._logger.info("Downloading LaTeX model...")
+      return snapshot_download(
+        cache_dir=self._model_dir_path,
+        repo_type="space",
+        repo_id="lukbl/LaTeX-OCR",
+      )

{doc_page_extractor-0.1.0 → doc_page_extractor-0.1.2}/doc_page_extractor/ocr.py

@@ -5,9 +5,8 @@ import os
 from typing import Literal, Generator
 from dataclasses import dataclass
 from .onnxocr import TextSystem
-from .types import OCRFragment
+from .types import GetModelDir, OCRFragment
 from .rectangle import Rectangle
-from .downloader import download
 from .utils import is_space_text
 
 
@@ -47,14 +46,17 @@ class _OONXParams:
   det_model_dir: str
   rec_char_dict_path: str
 
+
+
+
 class OCR:
   def __init__(
       self,
       device: Literal["cpu", "cuda"],
-      model_dir_path: str,
+      get_model_dir: GetModelDir,
     ):
     self._device: Literal["cpu", "cuda"] = device
-    self._model_dir_path: str = model_dir_path
+    self._get_model_dir: GetModelDir = get_model_dir
     self._text_system: TextSystem | None = None
 
   def search_fragments(self, image: np.ndarray) -> Generator[OCRFragment, None, None]:
@@ -69,7 +71,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(
@@ -87,20 +89,17 @@ class OCR:
     for box, res in zip(dt_boxes, rec_res):
       yield box.tolist(), res
 
+  def make_model_paths(self) -> list[str]:
+    model_paths = []
+    model_dir = self._get_model_dir()
+    for model_path in _MODELS:
+      file_name = os.path.join(*model_path)
+      model_paths.append(os.path.join(model_dir, file_name))
+    return model_paths
+
   def _get_text_system(self) -> TextSystem:
     if self._text_system is None:
-      for model_path in _MODELS:
-        file_path = os.path.join(self._model_dir_path, *model_path)
-        if os.path.exists(file_path):
-          continue
-
-        file_dir_path = os.path.dirname(file_path)
-        os.makedirs(file_dir_path, exist_ok=True)
-
-        url_path = "/".join(model_path)
-        url = f"https://huggingface.co/moskize/OnnxOCR/resolve/main/{url_path}"
-        download(url, file_path)
-
+      model_paths = self.make_model_paths()
       self._text_system = TextSystem(_OONXParams(
         use_angle_cls=True,
         use_gpu=(self._device != "cpu"),
@@ -123,10 +122,10 @@ class OCR:
         save_crop_res=False,
         rec_algorithm="SVTR_LCNet",
         use_space_char=True,
-        rec_model_dir=os.path.join(self._model_dir_path, *_MODELS[0]),
-        cls_model_dir=os.path.join(self._model_dir_path, *_MODELS[1]),
-        det_model_dir=os.path.join(self._model_dir_path, *_MODELS[2]),
-        rec_char_dict_path=os.path.join(self._model_dir_path, *_MODELS[3]),
+        rec_model_dir=model_paths[0],
+        cls_model_dir=model_paths[1],
+        det_model_dir=model_paths[2],
+        rec_char_dict_path=model_paths[3],
       ))
 
     return self._text_system

{doc_page_extractor-0.1.0 → doc_page_extractor-0.1.2}/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.0 → doc_page_extractor-0.1.2}/doc_page_extractor/table.py

@@ -3,7 +3,7 @@ import torch
 
 from typing import Literal, Any
 from PIL.Image import Image
-from .types import TableLayoutParsedFormat
+from .types import TableLayoutParsedFormat, GetModelDir
 from .utils import expand_image
 
 
@@ -13,10 +13,10 @@ class Table:
   def __init__(
       self,
       device: Literal["cpu", "cuda"],
-      model_path: str,
+      get_model_dir: GetModelDir,
     ):
     self._model: Any | None = None
-    self._model_path: str = model_path
+    self._model_path: str = get_model_dir()
     self._ban: bool = False
     if device == "cpu" or not torch.cuda.is_available():
       self._ban = True
@@ -58,13 +58,12 @@ class Table:
 
       from .struct_eqtable import build_model
       model = build_model(
-        model_ckpt="U4R/StructTable-InternVL2-1B",
+        model_ckpt=self._model_path,
         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()

{doc_page_extractor-0.1.0 → doc_page_extractor-0.1.2}/doc_page_extractor/types.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from typing import Literal
+from typing import Literal, Callable, Protocol, runtime_checkable
 from enum import auto, Enum
 from PIL.Image import Image
 from .rectangle import Rectangle
@@ -64,4 +64,25 @@ class ExtractedResult:
   rotation: float
   layouts: list[Layout]
   extracted_image: Image
-  adjusted_image: Image | None
+  adjusted_image: Image | None
+
+GetModelDir = Callable[[], str]
+
+
+@runtime_checkable
+class ModelsDownloader(Protocol):
+
+  def onnx_ocr(self) -> str:
+    pass
+
+  def yolo(self) -> str:
+    pass
+
+  def layoutreader(self) -> str:
+    pass
+
+  def struct_eqtable(self) -> str:
+    pass
+
+  def latex(self) -> str:
+    pass

{doc_page_extractor-0.1.0 → doc_page_extractor-0.1.2}/doc_page_extractor.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: doc-page-extractor
-Version: 0.1.0
+Version: 0.1.2
 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
@@ -16,6 +16,7 @@ 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
+Requires-Dist: huggingface_hub>=0.30.2
 Dynamic: author
 Dynamic: author-email
 Dynamic: description

{doc_page_extractor-0.1.0 → doc_page_extractor-0.1.2}/doc_page_extractor.egg-info/SOURCES.txt

@@ -8,6 +8,7 @@ doc_page_extractor/extractor.py
 doc_page_extractor/latex.py
 doc_page_extractor/layout_order.py
 doc_page_extractor/layoutreader.py
+doc_page_extractor/models.py
 doc_page_extractor/ocr.py
 doc_page_extractor/ocr_corrector.py
 doc_page_extractor/overlap.py

{doc_page_extractor-0.1.0 → doc_page_extractor-0.1.2}/doc_page_extractor.egg-info/requires.txt

@@ -7,3 +7,4 @@ transformers<=4.47,>=4.42.4
 doclayout_yolo>=0.0.3
 pix2tex<=0.2.0,>=0.1.4
 accelerate<2.0,>=1.6.0
+huggingface_hub>=0.30.2

{doc_page_extractor-0.1.0 → doc_page_extractor-0.1.2}/setup.py

@@ -5,7 +5,7 @@ if "doc_page_extractor.struct_eqtable" not in find_packages():
 
 setup(
   name="doc-page-extractor",
-  version="0.1.0",
+  version="0.1.2",
   author="Tao Zeyu",
   author_email="i@taozeyu.com",
   url="https://github.com/Moskize91/doc-page-extractor",
@@ -23,5 +23,6 @@ setup(
     "doclayout_yolo>=0.0.3",
     "pix2tex>=0.1.4,<=0.2.0",
     "accelerate>=1.6.0,<2.0",
+    "huggingface_hub>=0.30.2",
   ],
 )

doc_page_extractor-0.1.0/doc_page_extractor/latex.py

@@ -1,57 +0,0 @@
-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