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

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

@@ -1,13 +1,13 @@
 Metadata-Version: 2.4
 Name: doc-page-extractor
-Version: 0.1.2
+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

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

@@ -18,12 +18,12 @@ from .types import (
   ExtractedResult,
   ModelsDownloader,
   OCRFragment,
-  TableLayoutParsedFormat,
   Layout,
   LayoutClass,
   PlainLayout,
   TableLayout,
   FormulaLayout,
+  TableLayoutParsedFormat
 )
 
 
@@ -32,9 +32,6 @@ class DocExtractor:
       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,
     ):
@@ -42,9 +39,6 @@ class DocExtractor:
     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,
@@ -56,6 +50,7 @@ class DocExtractor:
     )
     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,
@@ -64,7 +59,10 @@ class DocExtractor:
   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 +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)
@@ -142,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] = []

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

@@ -4,13 +4,15 @@ 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, get_model_dir: GetModelDir):
+  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) # 添加边缘提高识别准确率
@@ -23,7 +25,7 @@ class LaTeX:
       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_cuda": self._device == "cpu",
         "no_resize": False,
       }))
     return self._model

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

@@ -20,13 +20,14 @@ 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 _get_model(self) -> LayoutLMv3ForTokenClassification:
     if self._model is None:
       self._model = LayoutLMv3ForTokenClassification.from_pretrained(
         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.1.2 → doc_page_extractor-0.2.0}/doc_page_extractor/ocr.py

@@ -141,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.2 → 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.2 → 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.2 → 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.2 → 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.2 → doc_page_extractor-0.2.0}/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, 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,11 +59,12 @@ class FormulaLayout(BaseLayout):
 
 Layout = PlainLayout | TableLayout | FormulaLayout
 
+
 @dataclass
 class ExtractedResult:
   rotation: float
-  layouts: list[Layout]
-  extracted_image: Image
+  layouts: List[Layout]
+  extracted_image: Image | None
   adjusted_image: Image | None
 
 GetModelDir = Callable[[], str]
@@ -86,3 +87,5 @@ class ModelsDownloader(Protocol):
 
   def latex(self) -> str:
     pass
+
+

{doc_page_extractor-0.1.2 → 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.2
+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

{doc_page_extractor-0.1.2 → 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

{doc_page_extractor-0.1.2 → 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.2",
+  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",

{doc_page_extractor-0.1.2 → 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, [