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

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

@@ -1,13 +1,13 @@
 Metadata-Version: 2.4
 Name: doc-page-extractor
-Version: 0.1.1
+Version: 0.2.0
 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,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.1 → doc_page_extractor-0.2.0}/doc_page_extractor/__init__.py

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

{doc_page_extractor-0.1.1 → doc_page_extractor-0.2.0}/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,50 +16,53 @@ 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_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
-    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"),
+      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,
+      device=device,
     )
     self._layout_order: LayoutOrder = LayoutOrder(
-      model_path=os.path.join(model_dir_path, "layoutreader"),
+      get_model_dir=self._models_downloader.layoutreader,
     )
 
   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)
@@ -70,13 +72,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)
@@ -138,16 +140,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] = []
@@ -191,14 +199,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.2.0/doc_page_extractor/latex.py

@@ -0,0 +1,31 @@
+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 .types import GetModelDir
+
+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
+    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._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": self._device == "cpu",
+        "no_resize": False,
+      }))
+    return self._model

{doc_page_extractor-0.1.1 → doc_page_extractor-0.2.0}/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,18 +17,17 @@ 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
+    self._device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
   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,
+      ).to(device=self._device)
     return self._model
 
   def sort(self, layouts: list[Layout], size: tuple[int, int]) -> list[Layout]:

doc_page_extractor-0.2.0/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.1 → doc_page_extractor-0.2.0}/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]:
@@ -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
@@ -142,14 +141,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.1 → doc_page_extractor-0.2.0}/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.1 → doc_page_extractor-0.2.0}/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.1 → doc_page_extractor-0.2.0}/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.1 → doc_page_extractor-0.2.0}/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.1 → doc_page_extractor-0.2.0}/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.1 → doc_page_extractor-0.2.0}/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, List
 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):
@@ -59,9 +59,33 @@ class FormulaLayout(BaseLayout):
 
 Layout = PlainLayout | TableLayout | FormulaLayout
 
+
 @dataclass
 class ExtractedResult:
   rotation: float
-  layouts: list[Layout]
-  extracted_image: Image
-  adjusted_image: Image | None
+  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
+
+

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

@@ -1,13 +1,13 @@
 Metadata-Version: 2.4
 Name: doc-page-extractor
-Version: 0.1.1
+Version: 0.2.0
 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,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.1 → doc_page_extractor-0.2.0}/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.1 → doc_page_extractor-0.2.0}/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,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.1 → doc_page_extractor-0.2.0}/setup.py

@@ -5,7 +5,7 @@ if "doc_page_extractor.struct_eqtable" not in find_packages():
 
 setup(
   name="doc-page-extractor",
-  version="0.1.1",
+  version="0.2.0",
   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,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.1 → doc_page_extractor-0.2.0}/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.1/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