doc-page-extractor 0.2.0__tar.gz → 0.2.2__tar.gz

{doc_page_extractor-0.2.0 → doc_page_extractor-0.2.2}/PKG-INFO

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

{doc_page_extractor-0.2.0 → doc_page_extractor-0.2.2}/doc_page_extractor/__init__.py

@@ -2,6 +2,7 @@ from .extractor import DocExtractor
 from .clipper import clip, clip_from_image
 from .plot import plot
 from .rectangle import Point, Rectangle
+from .model import Model, HuggingfaceModel
 from .types import (
   ExtractedResult,
   OCRFragment,
@@ -12,5 +13,4 @@ from .types import (
   PlainLayout,
   FormulaLayout,
   TableLayout,
-  ModelsDownloader
 )

{doc_page_extractor-0.2.0 → doc_page_extractor-0.2.2}/doc_page_extractor/extractor.py

@@ -1,10 +1,11 @@
+import torch
 
+from os import PathLike
 from typing import Literal, Generator
 from PIL.Image import Image
 from doclayout_yolo import YOLOv10
-from logging import Logger, getLogger
 
-from .models import HuggingfaceModelsDownloader
+from .model import Model, HuggingfaceModel
 from .ocr import OCR
 from .ocr_corrector import correct_fragments
 from .raw_optimizer import RawOptimizer
@@ -16,7 +17,6 @@ from .overlap import merge_fragments_as_line, remove_overlap_layouts
 from .clipper import clip_from_image
 from .types import (
   ExtractedResult,
-  ModelsDownloader,
   OCRFragment,
   Layout,
   LayoutClass,
@@ -29,32 +29,35 @@ from .types import (
 
 class DocExtractor:
   def __init__(
-      self,
-      model_cache_dir: str | None = None,
-      device: Literal["cpu", "cuda"] = "cpu",
-      models_downloader: ModelsDownloader | None = None,
-      logger: Logger | None = None,
-    ):
-    self._logger = logger or getLogger(__name__)
-    self._models_downloader = models_downloader or HuggingfaceModelsDownloader(self._logger, model_cache_dir)
+        self,
+        model_cache_dir: PathLike | None = None,
+        device: Literal["cpu", "cuda"] = "cpu",
+        model: Model | None = None,
+      ) -> None:
+
+    if model is None:
+      if model_cache_dir is None:
+        raise ValueError("You must provide a model_cache_dir or a model instance.")
+      model = HuggingfaceModel(model_cache_dir)
+
+    if device == "cuda" and not torch.cuda.is_available():
+      device = "cpu"
+      print("CUDA is not available. Using CPU instead.")
 
     self._device: Literal["cpu", "cuda"] = device
+    self._model: Model = model
     self._yolo: YOLOv10 | None = None
-    self._ocr: OCR = OCR(
-      device=device,
-      get_model_dir=self._models_downloader.onnx_ocr,
-    )
-    self._table: Table = Table(
-      device=device,
-      get_model_dir=self._models_downloader.struct_eqtable,
-    )
-    self._latex: LaTeX = LaTeX(
-      get_model_dir=self._models_downloader.latex,
-      device=device,
-    )
-    self._layout_order: LayoutOrder = LayoutOrder(
-      get_model_dir=self._models_downloader.layoutreader,
-    )
+    self._ocr: OCR = OCR(device, model)
+    self._table: Table = Table(device, model)
+    self._latex: LaTeX = LaTeX(device, model)
+    self._layout_order: LayoutOrder = LayoutOrder(device, model)
+
+  def prepare_models(self):
+    self._model.get_onnx_ocr_path()
+    self._model.get_yolo_path()
+    self._model.get_layoutreader_path()
+    self._model.get_struct_eqtable_path()
+    self._model.get_latex_path()
 
   def extract(
       self,
@@ -199,7 +202,7 @@ class DocExtractor:
 
   def _get_yolo(self) -> YOLOv10:
     if self._yolo is None:
-      model_path = self._models_downloader.yolo()
+      model_path = self._model.get_yolo_path()
       self._yolo = YOLOv10(str(model_path))
     return self._yolo
 

{doc_page_extractor-0.2.0 → doc_page_extractor-0.2.2}/doc_page_extractor/latex.py

@@ -6,12 +6,13 @@ from pix2tex.cli import LatexOCR
 from PIL.Image import Image
 from typing import Literal
 from .utils import expand_image
-from .types import GetModelDir
+from .model import Model
+
 
 class LaTeX:
-  def __init__(self, device: Literal["cpu", "cuda"],get_model_dir: GetModelDir):
-    self._model_path: str = get_model_dir()
-    self._model: LatexOCR | None = None
+  def __init__(self, device: Literal["cpu", "cuda"], model: Model) -> None:
+    self._model: Model = model
+    self._latex_model: LatexOCR | None = None
     self._device: Literal["cpu", "cuda"] = device
 
   def extract(self, image: Image) -> str | None:
@@ -21,11 +22,12 @@ class LaTeX:
       return model(image)
 
   def _get_model(self) -> LatexOCR:
-    if self._model is None:
-      self._model = LatexOCR(Munch({
+    if self._latex_model is None:
+      model_path = self._model.get_latex_path()
+      self._latex_model = LatexOCR(Munch({
         "config": os.path.join("settings", "config.yaml"),
-        "checkpoint": os.path.join(self._model_path, "checkpoints", "weights.pth"),
+        "checkpoint": os.path.join(model_path, "checkpoints", "weights.pth"),
         "no_cuda": self._device == "cpu",
         "no_resize": False,
       }))
-    return self._model
+    return self._latex_model

{doc_page_extractor-0.2.0 → doc_page_extractor-0.2.2}/doc_page_extractor/layout_order.py

@@ -1,10 +1,11 @@
 import torch
 
-from typing import Generator
+from typing import Generator, Literal
 from dataclasses import dataclass
 from transformers import LayoutLMv3ForTokenClassification
 
-from .types import Layout, LayoutClass, GetModelDir
+from .types import Layout, LayoutClass
+from .model import Model
 from .layoutreader import prepare_inputs, boxes2inputs, parse_logits
 
 
@@ -17,18 +18,19 @@ class _BBox:
   value: tuple[float, float, float, float]
 
 class LayoutOrder:
-  def __init__(self, get_model_dir: GetModelDir):
-    self._model_path: str = get_model_dir()
-    self._model: LayoutLMv3ForTokenClassification | None = None
-    self._device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+  def __init__(self, device: Literal["cpu", "cuda"], model: Model):
+    self._model: Model = model
+    self._order_model: LayoutLMv3ForTokenClassification | None = None
+    self._device: Literal["cpu", "cuda"] = device
 
   def _get_model(self) -> LayoutLMv3ForTokenClassification:
-    if self._model is None:
-      self._model = LayoutLMv3ForTokenClassification.from_pretrained(
-        pretrained_model_name_or_path=self._model_path,
+    if self._order_model is None:
+      model_path = self._model.get_layoutreader_path()
+      self._order_model = LayoutLMv3ForTokenClassification.from_pretrained(
+        pretrained_model_name_or_path=model_path,
         local_files_only=True,
       ).to(device=self._device)
-    return self._model
+    return self._order_model
 
   def sort(self, layouts: list[Layout], size: tuple[int, int]) -> list[Layout]:
     width, height = size

doc_page_extractor-0.2.2/doc_page_extractor/model.py

@@ -0,0 +1,114 @@
+from os import PathLike
+from typing import runtime_checkable, Protocol
+from pathlib import Path
+from threading import Lock
+from huggingface_hub import hf_hub_download, snapshot_download, try_to_load_from_cache
+
+
+@runtime_checkable
+class Model(Protocol):
+  def get_onnx_ocr_path(self) -> Path:
+    pass
+
+  def get_yolo_path(self) -> Path:
+    pass
+
+  def get_layoutreader_path(self) -> Path:
+    pass
+
+  def get_struct_eqtable_path(self) -> Path:
+    pass
+
+  def get_latex_path(self) -> Path:
+    pass
+
+class HuggingfaceModel(Model):
+  def __init__(self, model_cache_dir: PathLike):
+    super().__init__()
+    self._lock: Lock = Lock()
+    self._model_cache_dir: Path = Path(model_cache_dir)
+
+  def get_onnx_ocr_path(self) -> Path:
+    return self._get_model_path(
+      repo_id="moskize/OnnxOCR",
+      filename="README.md",
+      repo_type=None,
+      is_snapshot=True,
+      wanna_dir_path=True,
+    )
+
+  def get_yolo_path(self) -> Path:
+    return self._get_model_path(
+      repo_id="opendatalab/PDF-Extract-Kit-1.0",
+      filename="models/Layout/YOLO/doclayout_yolo_ft.pt",
+      repo_type=None,
+      is_snapshot=False,
+      wanna_dir_path=False,
+    )
+
+  def get_layoutreader_path(self) -> Path:
+    return self._get_model_path(
+      repo_id="hantian/layoutreader",
+      filename="model.safetensors",
+      repo_type=None,
+      is_snapshot=True,
+      wanna_dir_path=True,
+    )
+
+  def get_struct_eqtable_path(self) -> Path:
+    return self._get_model_path(
+      repo_id="U4R/StructTable-InternVL2-1B",
+      filename="model.safetensors",
+      repo_type=None,
+      is_snapshot=True,
+      wanna_dir_path=True,
+    )
+
+  def get_latex_path(self) -> Path:
+    return self._get_model_path(
+      repo_id="lukbl/LaTeX-OCR",
+      filename="checkpoints/weights.pth",
+      repo_type="space",
+      is_snapshot=True,
+      wanna_dir_path=True,
+    )
+
+  def _get_model_path(
+        self,
+        repo_id: str,
+        filename: str,
+        repo_type: str | None,
+        is_snapshot: bool,
+        wanna_dir_path: bool,
+      ) -> Path:
+
+    with self._lock:
+      model_path = try_to_load_from_cache(
+        repo_id=repo_id,
+        filename=filename,
+        repo_type=repo_type,
+        cache_dir=self._model_cache_dir
+      )
+      if isinstance(model_path, str):
+        model_path = Path(model_path)
+        if wanna_dir_path:
+          for _ in Path(filename).parts:
+            model_path = model_path.parent
+
+      else:
+        if is_snapshot:
+          model_path = snapshot_download(
+            cache_dir=self._model_cache_dir,
+            repo_id=repo_id,
+            repo_type=repo_type,
+          )
+        else:
+          model_path = hf_hub_download(
+            cache_dir=self._model_cache_dir,
+            repo_id=repo_id,
+            repo_type=repo_type,
+            filename=filename,
+          )
+        model_path = Path(model_path)
+
+      return model_path

{doc_page_extractor-0.2.0 → doc_page_extractor-0.2.2}/doc_page_extractor/ocr.py

@@ -5,7 +5,8 @@ import os
 from typing import Literal, Generator
 from dataclasses import dataclass
 from .onnxocr import TextSystem
-from .types import GetModelDir, OCRFragment
+from .types import OCRFragment
+from .model import Model
 from .rectangle import Rectangle
 from .utils import is_space_text
 
@@ -46,17 +47,10 @@ class _OONXParams:
   det_model_dir: str
   rec_char_dict_path: str
 
-
-
-
 class OCR:
-  def __init__(
-      self,
-      device: Literal["cpu", "cuda"],
-      get_model_dir: GetModelDir,
-    ):
+  def __init__(self, device: Literal["cpu", "cuda"], model: Model):
     self._device: Literal["cpu", "cuda"] = device
-    self._get_model_dir: GetModelDir = get_model_dir
+    self._model: Model = model
     self._text_system: TextSystem | None = None
 
   def search_fragments(self, image: np.ndarray) -> Generator[OCRFragment, None, None]:
@@ -89,17 +83,9 @@ 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:
-      model_paths = self.make_model_paths()
+      model_paths = self._make_model_paths()
       self._text_system = TextSystem(_OONXParams(
         use_angle_cls=True,
         use_gpu=(self._device != "cpu"),
@@ -127,9 +113,16 @@ class OCR:
         det_model_dir=model_paths[2],
         rec_char_dict_path=model_paths[3],
       ))
-
     return self._text_system
 
+  def _make_model_paths(self) -> list[str]:
+    model_paths: list[str] = []
+    model_dir = self._model.get_onnx_ocr_path()
+    for model_path in _MODELS:
+      file_name = os.path.join(*model_path)
+      model_paths.append(str(model_dir / file_name))
+    return model_paths
+
   def _preprocess_image(self, image: np.ndarray) -> np.ndarray:
     image = self._alpha_to_color(image, (255, 255, 255))
     # image = cv2.bitwise_not(image) # inv

{doc_page_extractor-0.2.0 → doc_page_extractor-0.2.2}/doc_page_extractor/table.py

@@ -1,24 +1,20 @@
-import os
 import torch
 
 from typing import Literal, Any
 from PIL.Image import Image
-from .types import TableLayoutParsedFormat, GetModelDir
+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"],
-      get_model_dir: GetModelDir,
-    ):
-    self._model: Any | None = None
-    self._model_path: str = get_model_dir()
+  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" or not torch.cuda.is_available():
+    if device == "cpu":
       self._ban = True
 
   def predict(self, image: Image, format: TableLayoutParsedFormat) -> str | None:
@@ -47,24 +43,18 @@ class Table:
 
     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)
-
+  def _get_model(self) -> Any:
+    if self._table_model is None:
       from .struct_eqtable import build_model
-      model = build_model(
-        model_ckpt=self._model_path,
+      model_path = self._model.get_struct_eqtable_path()
+      table_model = build_model(
+        model_ckpt=model_path,
         max_new_tokens=1024,
         max_time=30,
         lmdeploy=False,
         flash_attn=True,
         batch_size=1,
-        local_files_only=local_files_only,
+        local_files_only=False,
       )
-      self._model = model.cuda()
-    return self._model
+      self._table_model = table_model.cuda()
+    return self._table_model

{doc_page_extractor-0.2.0 → doc_page_extractor-0.2.2}/doc_page_extractor/types.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from typing import Literal, Callable, Protocol, runtime_checkable, List
+from typing import Literal
 from enum import auto, Enum
 from PIL.Image import Image
 from .rectangle import Rectangle
@@ -32,7 +32,7 @@ class TableLayoutParsedFormat(Enum):
 @dataclass
 class BaseLayout:
   rect: Rectangle
-  fragments: List[OCRFragment]
+  fragments: list[OCRFragment]
 
 @dataclass
 class PlainLayout(BaseLayout):
@@ -63,29 +63,6 @@ Layout = PlainLayout | TableLayout | FormulaLayout
 @dataclass
 class ExtractedResult:
   rotation: float
-  layouts: List[Layout]
+  layouts: list[Layout]
   extracted_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
-
-
+  adjusted_image: Image | None

{doc_page_extractor-0.2.0 → doc_page_extractor-0.2.2}/doc_page_extractor.egg-info/PKG-INFO

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

{doc_page_extractor-0.2.0 → doc_page_extractor-0.2.2}/doc_page_extractor.egg-info/SOURCES.txt

@@ -8,7 +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/model.py
 doc_page_extractor/ocr.py
 doc_page_extractor/ocr_corrector.py
 doc_page_extractor/overlap.py

{doc_page_extractor-0.2.0 → doc_page_extractor-0.2.2}/doc_page_extractor.egg-info/requires.txt

@@ -7,4 +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
+huggingface_hub<1.0,>=0.30.2

{doc_page_extractor-0.2.0 → doc_page_extractor-0.2.2}/setup.py

@@ -5,7 +5,7 @@ if "doc_page_extractor.struct_eqtable" not in find_packages():
 
 setup(
   name="doc-page-extractor",
-  version="0.2.0",
+  version="0.2.2",
   author="Tao Zeyu",
   author_email="i@taozeyu.com",
   url="https://github.com/Moskize91/doc-page-extractor",
@@ -23,6 +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",
+    "huggingface_hub>=0.30.2,<1.0",
   ],
 )

doc_page_extractor-0.2.0/doc_page_extractor/models.py

@@ -1,92 +0,0 @@
-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",
-      )