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

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

@@ -1,13 +1,13 @@
 Metadata-Version: 2.4
 Name: doc-page-extractor
-Version: 0.1.2
+Version: 0.2.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
 Author-email: i@taozeyu.com
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: opencv-python<5.0,>=4.11.0
+Requires-Dist: opencv-python<5.0,>=4.10.0
 Requires-Dist: pillow<11.0,>=10.3
 Requires-Dist: pyclipper<2.0,>=1.2.0
 Requires-Dist: numpy<2.0,>=1.24.0
@@ -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.1.2 → doc_page_extractor-0.2.1}/doc_page_extractor/__init__.py

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

{doc_page_extractor-0.1.2 → doc_page_extractor-0.2.1}/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,55 +17,48 @@ from .overlap import merge_fragments_as_line, remove_overlap_layouts
 from .clipper import clip_from_image
 from .types import (
   ExtractedResult,
-  ModelsDownloader,
   OCRFragment,
-  TableLayoutParsedFormat,
   Layout,
   LayoutClass,
   PlainLayout,
   TableLayout,
   FormulaLayout,
+  TableLayoutParsedFormat
 )
 
 
 class DocExtractor:
   def __init__(
-      self,
-      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._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._ocr_for_each_layouts: bool = ocr_for_each_layouts
-    self._extract_formula: bool = extract_formula
-    self._extract_table_format: TableLayoutParsedFormat | None = extract_table_format
+    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,
-    )
-    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 extract(
       self,
       image: Image,
-      adjust_points: bool = False,
+      extract_formula: bool,
+      extract_table_format: TableLayoutParsedFormat | None = None,
+      ocr_for_each_layouts: bool = False,
+      adjust_points: bool = False
     ) -> ExtractedResult:
 
     raw_optimizer = RawOptimizer(image, adjust_points)
@@ -74,13 +68,13 @@ class DocExtractor:
     layouts = self._layouts_matched_by_fragments(fragments, layouts)
     layouts = remove_overlap_layouts(layouts)
 
-    if self._ocr_for_each_layouts:
+    if ocr_for_each_layouts:
       self._correct_fragments_by_ocr_layouts(raw_optimizer.image, layouts)
 
     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)
+    self._parse_table_and_formula_layouts(layouts, raw_optimizer, extract_formula=extract_formula, extract_table_format=extract_table_format)
 
     for layout in layouts:
       layout.fragments = merge_fragments_as_line(layout.fragments)
@@ -142,16 +136,22 @@ 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):
+  def _parse_table_and_formula_layouts(
+      self,
+      layouts: list[Layout],
+      raw_optimizer: RawOptimizer,
+      extract_formula: bool,
+      extract_table_format: TableLayoutParsedFormat | None,
+  ):
     for layout in layouts:
-      if isinstance(layout, FormulaLayout) and self._extract_formula:
+      if isinstance(layout, FormulaLayout) and 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:
+      elif isinstance(layout, TableLayout) and extract_table_format is not None:
         image = clip_from_image(raw_optimizer.image, layout.rect)
-        parsed = self._table.predict(image, self._extract_table_format)
+        parsed = self._table.predict(image, extract_table_format)
         if parsed is not None:
-          layout.parsed = (parsed, self._extract_table_format)
+          layout.parsed = (parsed, extract_table_format)
 
   def _split_layouts_by_group(self, layouts: list[Layout]):
     texts_layouts: list[Layout] = []
@@ -195,7 +195,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.1/doc_page_extractor/latex.py

@@ -0,0 +1,33 @@
+import os
+import torch
+
+from munch import Munch
+from pix2tex.cli import LatexOCR
+from PIL.Image import Image
+from typing import Literal
+from .utils import expand_image
+from .model import Model
+
+
+class LaTeX:
+  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:
+    image = expand_image(image, 0.1) # 添加边缘提高识别准确率
+    model = self._get_model()
+    with torch.no_grad():
+      return model(image)
+
+  def _get_model(self) -> LatexOCR:
+    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(model_path, "checkpoints", "weights.pth"),
+        "no_cuda": self._device == "cpu",
+        "no_resize": False,
+      }))
+    return self._latex_model

{doc_page_extractor-0.1.2 → doc_page_extractor-0.2.1}/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,17 +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
+  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,
-      )
-    return self._model
+      ).to(device=self._device)
+    return self._order_model
 
   def sort(self, layouts: list[Layout], size: tuple[int, int]) -> list[Layout]:
     width, height = size

doc_page_extractor-0.2.1/doc_page_extractor/model.py

@@ -0,0 +1,103 @@
+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=None,
+      repo_type=None,
+      is_snapshot=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,
+    )
+
+  def get_layoutreader_path(self) -> Path:
+    return self._get_model_path(
+      repo_id="hantian/layoutreader",
+      filename=None,
+      repo_type=None,
+      is_snapshot=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,
+    )
+
+  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,
+    )
+
+  def _get_model_path(
+        self,
+        repo_id: str,
+        filename: str | None,
+        repo_type: str | None,
+        is_snapshot: bool,
+      ) -> Path:
+    with self._lock:
+      cache_filename = "README.md"
+      if filename is not None:
+        cache_filename = filename
+      model_path = try_to_load_from_cache(
+        repo_id=repo_id,
+        filename=cache_filename,
+        repo_type=repo_type,
+        cache_dir=self._model_cache_dir
+      )
+      if isinstance(model_path, str):
+        if filename is None:
+          model_path = Path(model_path).parent
+
+      elif is_snapshot:
+        model_path = snapshot_download(
+          cache_dir=self._model_cache_dir,
+          repo_id=repo_id,
+        )
+      else:
+        model_path = hf_hub_download(
+          cache_dir=self._model_cache_dir,
+          repo_id=repo_id,
+          filename=filename,
+        )
+      return Path(model_path)

{doc_page_extractor-0.1.2 → doc_page_extractor-0.2.1}/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
@@ -141,14 +134,40 @@ class OCR:
       beta=255,
       norm_type=cv2.NORM_MINMAX,
     )
-    image = cv2.fastNlMeansDenoisingColored(
-      src=image,
-      dst=None,
-      h=10,
-      hColor=10,
-      templateWindowSize=7,
-      searchWindowSize=15,
-    )
+    if cv2.cuda.getCudaEnabledDeviceCount() > 0:
+      gpu_frame = cv2.cuda.GpuMat()
+      gpu_frame.upload(image)
+      image = cv2.cuda.fastNlMeansDenoisingColored(
+        src=gpu_frame,
+        dst=None,
+        h_luminance=10,
+        photo_render=10,
+        search_window=15,
+        block_size=7,
+      )
+      image = gpu_frame.download()
+    elif cv2.ocl.haveOpenCL():
+      cv2.ocl.setUseOpenCL(True)
+      gpu_frame = cv2.UMat(image)
+      image = cv2.fastNlMeansDenoisingColored(
+        src=gpu_frame,
+        dst=None,
+        h=10,
+        hColor=10,
+        templateWindowSize=7,
+        searchWindowSize=15,
+      )
+      image = image.get()
+    else:
+      image = cv2.fastNlMeansDenoisingColored(
+        src=image,
+        dst=None,
+        h=10,
+        hColor=10,
+        templateWindowSize=7,
+        searchWindowSize=15,
+      )
+
     # image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) # image to gray
     return image
 

{doc_page_extractor-0.1.2 → doc_page_extractor-0.2.1}/doc_page_extractor/onnxocr/predict_base.py

@@ -1,8 +1,13 @@
-import onnxruntime
-
 class PredictBase(object):
   def __init__(self):
-    pass
+    self._onnxruntime = None
+
+  @property
+  def onnxruntime(self):
+    if self._onnxruntime is None:
+      import onnxruntime
+      self._onnxruntime = onnxruntime
+    return self._onnxruntime
 
   def get_onnx_session(self, model_dir, use_gpu):
     # 使用gpu
@@ -11,7 +16,7 @@ class PredictBase(object):
     else:
       providers = providers = ['CPUExecutionProvider']
 
-    onnx_session = onnxruntime.InferenceSession(model_dir, None,providers=providers)
+    onnx_session = self.onnxruntime.InferenceSession(model_dir, None, providers=providers)
 
     # print("providers:", onnxruntime.get_device())
     return onnx_session

{doc_page_extractor-0.1.2 → doc_page_extractor-0.2.1}/doc_page_extractor/onnxocr/predict_cls.py

@@ -9,15 +9,35 @@ from .predict_base import PredictBase
 
 class TextClassifier(PredictBase):
   def __init__(self, args):
+    super().__init__()
     self.cls_image_shape = args.cls_image_shape
     self.cls_batch_num = args.cls_batch_num
     self.cls_thresh = args.cls_thresh
     self.postprocess_op = ClsPostProcess(label_list=args.label_list)
+    self._args = args
 
     # 初始化模型
-    self.cls_onnx_session = self.get_onnx_session(args.cls_model_dir, args.use_gpu)
-    self.cls_input_name = self.get_input_name(self.cls_onnx_session)
-    self.cls_output_name = self.get_output_name(self.cls_onnx_session)
+    self._cls_onnx_session = None
+    self._cls_input_name = None
+    self._cls_output_name = None
+
+  @property
+  def cls_onnx_session(self):
+    if self._cls_onnx_session is None:
+      self._cls_onnx_session = self.get_onnx_session(self._args.cls_model_dir, self._args.use_gpu)
+    return self._cls_onnx_session
+
+  @property
+  def cls_input_name(self):
+    if self._cls_input_name is None:
+      self._cls_input_name = self.get_input_name(self.cls_onnx_session)
+    return self._cls_input_name
+
+  @property
+  def cls_output_name(self):
+    if self._cls_output_name is None:
+      self._cls_output_name = self.get_output_name(self.cls_onnx_session)
+    return self._cls_output_name
 
   def resize_norm_img(self, img):
     imgC, imgH, imgW = self.cls_image_shape

{doc_page_extractor-0.1.2 → doc_page_extractor-0.2.1}/doc_page_extractor/onnxocr/predict_det.py

@@ -6,7 +6,8 @@ from .predict_base import PredictBase
 
 class TextDetector(PredictBase):
   def __init__(self, args):
-    self.args = args
+    super().__init__()
+    self._args = args
     self.det_algorithm = args.det_algorithm
     pre_process_list = [
       {
@@ -43,9 +44,27 @@ class TextDetector(PredictBase):
     self.postprocess_op = DBPostProcess(**postprocess_params)
 
     # 初始化模型
-    self.det_onnx_session = self.get_onnx_session(args.det_model_dir, args.use_gpu)
-    self.det_input_name = self.get_input_name(self.det_onnx_session)
-    self.det_output_name = self.get_output_name(self.det_onnx_session)
+    self._det_onnx_session = None
+    self._det_input_name = None
+    self._det_output_name = None
+
+  @property
+  def det_onnx_session(self):
+    if self._det_onnx_session is None:
+      self._det_onnx_session = self.get_onnx_session(self._args.det_model_dir, self._args.use_gpu)
+    return self._det_onnx_session
+
+  @property
+  def det_input_name(self):
+    if self._det_input_name is None:
+      self._det_input_name = self.get_input_name(self.det_onnx_session)
+    return self._det_input_name
+
+  @property
+  def det_output_name(self):
+    if self._det_output_name is None:
+      self._det_output_name = self.get_output_name(self.det_onnx_session)
+    return self._det_output_name
 
   def order_points_clockwise(self, pts):
     rect = np.zeros((4, 2), dtype="float32")
@@ -112,7 +131,7 @@ class TextDetector(PredictBase):
     post_result = self.postprocess_op(preds, shape_list)
     dt_boxes = post_result[0]["points"]
 
-    if self.args.det_box_type == "poly":
+    if self._args.det_box_type == "poly":
       dt_boxes = self.filter_tag_det_res_only_clip(dt_boxes, ori_im.shape)
     else:
       dt_boxes = self.filter_tag_det_res(dt_boxes, ori_im.shape)

{doc_page_extractor-0.1.2 → doc_page_extractor-0.2.1}/doc_page_extractor/onnxocr/predict_rec.py

@@ -10,6 +10,8 @@ from .predict_base import PredictBase
 
 class TextRecognizer(PredictBase):
   def __init__(self, args):
+    super().__init__()
+    self._args = args
     self.rec_image_shape = args.rec_image_shape
     self.rec_batch_num = args.rec_batch_num
     self.rec_algorithm = args.rec_algorithm
@@ -19,9 +21,29 @@ class TextRecognizer(PredictBase):
     )
 
     # 初始化模型
-    self.rec_onnx_session = self.get_onnx_session(args.rec_model_dir, args.use_gpu)
-    self.rec_input_name = self.get_input_name(self.rec_onnx_session)
-    self.rec_output_name = self.get_output_name(self.rec_onnx_session)
+    self._rec_onnx_session = None
+    self._rec_input_name = None
+    self._rec_output_name = None
+
+  @property
+  def rec_onnx_session(self):
+    if self._rec_onnx_session is None:
+      self._rec_onnx_session = self.get_onnx_session(
+        self._args.rec_model_dir, self._args.use_gpu
+      )
+    return self._rec_onnx_session
+
+  @property
+  def rec_input_name(self):
+    if self._rec_input_name is None:
+      self._rec_input_name = self.get_input_name(self.rec_onnx_session)
+    return self._rec_input_name
+
+  @property
+  def rec_output_name(self):
+    if self._rec_output_name is None:
+      self._rec_output_name = self.get_output_name(self.rec_onnx_session)
+    return self._rec_output_name
 
   def resize_norm_img(self, img, max_wh_ratio):
     imgC, imgH, imgW = self.rec_image_shape
@@ -30,9 +52,9 @@ class TextRecognizer(PredictBase):
       # return padding_im
       image_pil = Image.fromarray(np.uint8(img))
       if self.rec_algorithm == "ViTSTR":
-        img = image_pil.resize([imgW, imgH], Image.BICUBIC)
+        img = image_pil.resize([imgW, imgH], Image.Resampling.BICUBIC)
       else:
-        img = image_pil.resize([imgW, imgH], Image.ANTIALIAS)
+        img = image_pil.resize([imgW, imgH], Image.Resampling.LANCZOS)
       img = np.array(img)
       norm_img = np.expand_dims(img, -1)
       norm_img = norm_img.transpose((2, 0, 1))
@@ -250,8 +272,9 @@ class TextRecognizer(PredictBase):
   def norm_img_can(self, img, image_shape):
     img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)  # CAN only predict gray scale image
 
-    if self.inverse:
-      img = 255 - img
+    # FIXME
+    # if self.inverse:
+    #   img = 255 - img
 
     if self.rec_image_shape[0] == 1:
       h, w = img.shape

{doc_page_extractor-0.1.2 → doc_page_extractor-0.2.1}/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.1.2 → doc_page_extractor-0.2.1}/doc_page_extractor/types.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from typing import Literal, Callable, Protocol, runtime_checkable
+from typing import Literal
 from enum import auto, Enum
 from PIL.Image import Image
 from .rectangle import Rectangle
@@ -59,30 +59,10 @@ class FormulaLayout(BaseLayout):
 
 Layout = PlainLayout | TableLayout | FormulaLayout
 
+
 @dataclass
 class ExtractedResult:
   rotation: float
   layouts: list[Layout]
-  extracted_image: Image
-  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
+  extracted_image: Image | None
+  adjusted_image: Image | None

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

@@ -1,13 +1,13 @@
 Metadata-Version: 2.4
 Name: doc-page-extractor
-Version: 0.1.2
+Version: 0.2.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
 Author-email: i@taozeyu.com
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: opencv-python<5.0,>=4.11.0
+Requires-Dist: opencv-python<5.0,>=4.10.0
 Requires-Dist: pillow<11.0,>=10.3
 Requires-Dist: pyclipper<2.0,>=1.2.0
 Requires-Dist: numpy<2.0,>=1.24.0
@@ -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.1.2 → doc_page_extractor-0.2.1}/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.1.2 → doc_page_extractor-0.2.1}/doc_page_extractor.egg-info/requires.txt

@@ -1,4 +1,4 @@
-opencv-python<5.0,>=4.11.0
+opencv-python<5.0,>=4.10.0
 pillow<11.0,>=10.3
 pyclipper<2.0,>=1.2.0
 numpy<2.0,>=1.24.0
@@ -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.1.2 → doc_page_extractor-0.2.1}/setup.py

@@ -5,7 +5,7 @@ if "doc_page_extractor.struct_eqtable" not in find_packages():
 
 setup(
   name="doc-page-extractor",
-  version="0.1.2",
+  version="0.2.1",
   author="Tao Zeyu",
   author_email="i@taozeyu.com",
   url="https://github.com/Moskize91/doc-page-extractor",
@@ -14,7 +14,7 @@ setup(
   long_description=open("./README.md", encoding="utf8").read(),
   long_description_content_type="text/markdown",
   install_requires=[
-    "opencv-python>=4.11.0,<5.0",
+    "opencv-python>=4.10.0,<5.0",
     "pillow>=10.3,<11.0",
     "pyclipper>=1.2.0,<2.0",
     "numpy>=1.24.0,<2.0",
@@ -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.1.2 → doc_page_extractor-0.2.1}/tests/test_history_bus.py

@@ -15,7 +15,7 @@ class TestGroup(unittest.TestCase):
     layouts: list[tuple[LayoutClass, list[str]]]
 
     with Image.open(image_path) as image:
-      result = extractor.extract(image, "ch")
+      result = extractor.extract(image, extract_formula=False)
       layouts = [self._format_Layout(layout) for layout in result.layouts]
 
     self.assertEqual(layouts, [

doc_page_extractor-0.1.2/doc_page_extractor/latex.py

@@ -1,29 +0,0 @@
-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.2/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",
-      )