doc-page-extractor 0.2.3__tar.gz → 0.2.4__tar.gz

{doc_page_extractor-0.2.3 → doc_page_extractor-0.2.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: doc-page-extractor
-Version: 0.2.3
+Version: 0.2.4
 Summary: 
 License: AGPL-3.0
 Author: Tao Zeyu
@@ -16,10 +16,14 @@ Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
+Provides-Extra: cpu
+Provides-Extra: cuda
 Requires-Dist: accelerate (>=1.6.0,<2.0)
 Requires-Dist: doclayout_yolo (>=0.0.3)
 Requires-Dist: huggingface_hub (>=0.33.0,<1.0)
 Requires-Dist: numpy (>=1.24.0,<2.0)
+Requires-Dist: onnxruntime (==1.21.0) ; extra == "cpu"
+Requires-Dist: onnxruntime-gpu (==1.21.0) ; extra == "cuda"
 Requires-Dist: opencv-python (>=4.10.0,<5.0)
 Requires-Dist: pillow (>=10.3,<11.0)
 Requires-Dist: pix2tex (>=0.1.4,<=0.2.0)
@@ -40,21 +44,17 @@ doc page extractor can identify text and format in images and return structured
 ## Installation
 
 ```shell
-pip install doc-page-extractor
-```
-
-```shell
-pip install onnxruntime==1.21.0
+pip install doc-page-extractor[cpu]
 ```
 
 ## Using CUDA
 
 Please refer to the introduction of [PyTorch](https://pytorch.org/get-started/locally/) and select the appropriate command to install according to your operating system.
 
-In addition, replace the command to install `onnxruntime` in the previous article with the following:
+The installation mentioned above uses the following command.
 
 ```shell
-pip install onnxruntime-gpu==1.21.0
+pip install doc-page-extractor[cuda]
 ```
 
 ## Example

{doc_page_extractor-0.2.3 → doc_page_extractor-0.2.4}/README.md

@@ -9,21 +9,17 @@ doc page extractor can identify text and format in images and return structured
 ## Installation
 
 ```shell
-pip install doc-page-extractor
-```
-
-```shell
-pip install onnxruntime==1.21.0
+pip install doc-page-extractor[cpu]
 ```
 
 ## Using CUDA
 
 Please refer to the introduction of [PyTorch](https://pytorch.org/get-started/locally/) and select the appropriate command to install according to your operating system.
 
-In addition, replace the command to install `onnxruntime` in the previous article with the following:
+The installation mentioned above uses the following command.
 
 ```shell
-pip install onnxruntime-gpu==1.21.0
+pip install doc-page-extractor[cuda]
 ```
 
 ## Example

{doc_page_extractor-0.2.3 → doc_page_extractor-0.2.4}/doc_page_extractor/clipper.py

@@ -13,12 +13,12 @@ def clip(
     wrapped_width: float = 0.0,
     wrapped_height: float = 0.0,
   ) -> Image:
-  image: Image
+  image: Image | None
   if extracted_result.adjusted_image is None:
     image = extracted_result.extracted_image
   else:
     image = extracted_result.adjusted_image
-
+  assert image is not None, "Image must not be None"
   return clip_from_image(
     image, layout.rect,
     wrapped_width, wrapped_height,
@@ -91,7 +91,7 @@ def _size_and_wrapper(rect: Rectangle):
 
   return width, height, max_width, max_height
 
-def _to_pillow_matrix(matrix: np.array):
+def _to_pillow_matrix(matrix):
   return (
     matrix[0][0], matrix[0][1], matrix[0][2],
     matrix[1][0], matrix[1][1], matrix[1][2],

{doc_page_extractor-0.2.3 → doc_page_extractor-0.2.4}/doc_page_extractor/extractor.py

@@ -1,7 +1,7 @@
 import torch
 
 from os import PathLike
-from typing import Literal, Generator
+from typing import cast, Any, Literal, Generator
 from PIL.Image import Image
 from doclayout_yolo import YOLOv10
 
@@ -99,7 +99,7 @@ class DocExtractor:
     # about source parameter to see:
     # https://github.com/opendatalab/DocLayout-YOLO/blob/7c4be36bc61f11b67cf4a44ee47f3c41e9800a91/doclayout_yolo/data/build.py#L157-L175
     det_res = self._get_yolo().predict(
-      source=source,
+      source=cast(Any, source),
       imgsz=1024,
       conf=0.2,
       device=self._device    # Device to use (e.g., "cuda" or "cpu")
@@ -180,7 +180,7 @@ class DocExtractor:
 
   def _find_matched_layout(self, fragment: OCRFragment, layouts: list[Layout]) -> Layout | None:
     fragment_area = fragment.rect.area
-    primary_layouts: list[(Layout, float)] = []
+    primary_layouts: list[tuple[Layout, float]] = []
 
     if fragment_area == 0.0:
       return None

{doc_page_extractor-0.2.3 → doc_page_extractor-0.2.4}/doc_page_extractor/layout_order.py

@@ -88,7 +88,7 @@ class LayoutOrder:
     y_scale = layoutreader_size / float(height)
 
     for bbox in bbox_list:
-      x0, y0, x1, y1 = self._squeeze(bbox.value, width, height)
+      x0, y0, x1, y1 = self._squeeze(bbox, width, height)
       x0 = round(x0 * x_scale)
       y0 = round(y0 * y_scale)
       x1 = round(x1 * x_scale)
@@ -223,8 +223,8 @@ class LayoutOrder:
       mid2 = sorted_numbers[n // 2]
       return float((mid1 + mid2) / 2)
 
-  def _squeeze(self, bbox: _BBox, width: int, height: int) -> _BBox:
-    x0, y0, x1, y1 = bbox
+  def _squeeze(self, bbox: _BBox, width: int, height: int) -> tuple[float, float, float, float]:
+    x0, y0, x1, y1 = bbox.value
     x0 = self._squeeze_value(x0, width)
     x1 = self._squeeze_value(x1, width)
     y0 = self._squeeze_value(y0, height)

{doc_page_extractor-0.2.3 → doc_page_extractor-0.2.4}/doc_page_extractor/layoutreader.py

@@ -64,7 +64,7 @@ class DataCollator:
     return ret
 
 
-def boxes2inputs(boxes: List[List[int]]) -> Dict[str, torch.Tensor]:
+def boxes2inputs(boxes: List[List[float]]) -> Dict[str, torch.Tensor]:
   bbox = [[0, 0, 0, 0]] + boxes + [[0, 0, 0, 0]]
   input_ids = [CLS_TOKEN_ID] + [UNK_TOKEN_ID] * len(boxes) + [EOS_TOKEN_ID]
   attention_mask = [1] + [1] * len(boxes) + [1]

{doc_page_extractor-0.2.3 → doc_page_extractor-0.2.4}/doc_page_extractor/model.py

@@ -1,6 +1,6 @@
 from os import PathLike
 from time import sleep
-from typing import runtime_checkable, Protocol
+from typing import cast, 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
@@ -12,19 +12,19 @@ _RETRY_SLEEP = 3.5
 @runtime_checkable
 class Model(Protocol):
   def get_onnx_ocr_path(self) -> Path:
-    pass
+    raise NotImplementedError()
 
   def get_yolo_path(self) -> Path:
-    pass
+    raise NotImplementedError()
 
   def get_layoutreader_path(self) -> Path:
-    pass
+    raise NotImplementedError()
 
   def get_struct_eqtable_path(self) -> Path:
-    pass
+    raise NotImplementedError()
 
   def get_latex_path(self) -> Path:
-    pass
+    raise NotImplementedError()
 
 class HuggingfaceModel(Model):
   def __init__(self, model_cache_dir: PathLike):
@@ -128,6 +128,6 @@ class HuggingfaceModel(Model):
 
         if latest_error is not None:
           raise latest_error
-        model_path = Path(model_path)
+        model_path = Path(cast(PathLike, model_path))
 
       return model_path

{doc_page_extractor-0.2.3 → doc_page_extractor-0.2.4}/doc_page_extractor/ocr.py

@@ -2,7 +2,7 @@ import numpy as np
 import cv2
 import os
 
-from typing import Literal, Generator
+from typing import cast, Any, Iterable, Literal, Generator
 from dataclasses import dataclass
 from .onnxocr import TextSystem
 from .types import OCRFragment
@@ -80,7 +80,10 @@ class OCR:
     image = self._preprocess_image(image)
     dt_boxes, rec_res = text_system(image)
 
-    for box, res in zip(dt_boxes, rec_res):
+    for box, res in zip(
+      cast(Iterable[Any], dt_boxes),
+      cast(Iterable[Any], rec_res),
+    ):
       yield box.tolist(), res
 
   def _get_text_system(self) -> TextSystem:
@@ -123,8 +126,8 @@ class OCR:
       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))
+  def _preprocess_image(self, np_image: np.ndarray) -> np.ndarray:
+    image = self._alpha_to_color(np_image, (255, 255, 255))
     # image = cv2.bitwise_not(image) # inv
     # image = self._binarize_img(image) # bin
     image = cv2.normalize(
@@ -148,7 +151,7 @@ class OCR:
       image = gpu_frame.download()
     elif cv2.ocl.haveOpenCL():
       cv2.ocl.setUseOpenCL(True)
-      gpu_frame = cv2.UMat(image)
+      gpu_frame = cv2.UMat(cast(Any, image))
       image = cv2.fastNlMeansDenoisingColored(
         src=gpu_frame,
         dst=None,

{doc_page_extractor-0.2.3 → doc_page_extractor-0.2.4}/doc_page_extractor/ocr_corrector.py

@@ -1,6 +1,6 @@
 import numpy as np
 
-from typing import Iterable
+from typing import cast, Iterable
 from shapely.geometry import Polygon
 from PIL.Image import new, Image, Resampling
 from .types import Layout, OCRFragment
@@ -90,13 +90,13 @@ def _match_fragments(
   ) -> tuple[list[tuple[OCRFragment, OCRFragment]], list[OCRFragment]]:
 
   zone_polygon = Polygon(zone_rect)
-  fragments2: list[OCRFragment] = list(fragments2)
+  fragments2 = list(fragments2)
   matched_fragments: list[tuple[OCRFragment, OCRFragment]] = []
   not_matched_fragments: list[OCRFragment] = []
 
   for fragment1 in fragments1:
     polygon1 = Polygon(fragment1.rect)
-    polygon1 = zone_polygon.intersection(polygon1)
+    polygon1 = cast(Polygon, zone_polygon.intersection(polygon1))
     if polygon1.is_empty:
       continue
 

{doc_page_extractor-0.2.3 → doc_page_extractor-0.2.4}/doc_page_extractor/overlap.py

@@ -1,4 +1,4 @@
-from typing import Generator
+from typing import cast, Generator
 from shapely.geometry import Polygon
 from .types import Layout, OCRFragment
 from .rectangle import Rectangle
@@ -92,7 +92,7 @@ def merge_fragments_as_line(origin_fragments: list[OCRFragment]) -> list[OCRFrag
       continue
 
     fragments.append(OCRFragment(
-      order=min_order,
+      order=round(min_order),
       text=" ".join(texts),
       rank=text_rate_weights / proto_texts_len,
       rect=Rectangle(
@@ -141,7 +141,7 @@ def _split_fragments_into_groups(fragments: list[OCRFragment]) -> Generator[list
 # they are very sensitive to changes in height because they are very thin and long.
 # In order to make it equally sensitive to length and width, the ratio of area is not used.
 def overlap_rate(polygon1: Polygon, polygon2: Polygon) -> float:
-  intersection: Polygon = polygon1.intersection(polygon2)
+  intersection = cast(Polygon, polygon1.intersection(polygon2))
   if intersection.is_empty:
     return 0.0
   else:

{doc_page_extractor-0.2.3 → doc_page_extractor-0.2.4}/doc_page_extractor/plot.py

@@ -1,4 +1,4 @@
-from typing import Iterable
+from typing import cast, Iterable
 from PIL import ImageDraw
 from PIL.ImageFont import load_default, FreeTypeFont
 from PIL.Image import Image
@@ -9,8 +9,8 @@ _FRAGMENT_COLOR = (0x49, 0xCF, 0xCB) # Light Green
 _Color = tuple[int, int, int]
 
 def plot(image: Image, layouts: Iterable[Layout]) -> None:
-  layout_font = load_default(size=35)
-  fragment_font = load_default(size=25)
+  layout_font = cast(FreeTypeFont, load_default(size=35))
+  fragment_font = cast(FreeTypeFont, load_default(size=25))
   draw = ImageDraw.Draw(image, mode="RGBA")
 
   def _draw_number(position: Point, number: int, font: FreeTypeFont, bold: bool, color: _Color) -> None:
@@ -88,4 +88,6 @@ def _layout_color(layout: Layout) -> _Color:
   elif cls == LayoutClass.ISOLATE_FORMULA:
     return (0xFA, 0x38, 0x27) # Red
   elif cls == LayoutClass.FORMULA_CAPTION:
-    return (0xFF, 0x9D, 0x24) # Orange
+    return (0xFF, 0x9D, 0x24) # Orange
+  else:
+    return (0x00, 0x00, 0x00)

{doc_page_extractor-0.2.3 → doc_page_extractor-0.2.4}/doc_page_extractor/raw_optimizer.py

@@ -3,7 +3,7 @@ import numpy as np
 from dataclasses import dataclass
 from PIL.Image import Image
 from math import pi
-from .types import OCRFragment, Layout
+from .types import Layout, OCRFragment
 from .rotation import calculate_rotation, RotationAdjuster
 from .rectangle import Rectangle
 

doc_page_extractor-0.2.4/doc_page_extractor/struct_eqtable/__init__.py

@@ -0,0 +1,49 @@
+from .pix2s import Pix2Struct, Pix2StructTensorRT
+from .internvl import InternVL, InternVL_LMDeploy
+
+from transformers import AutoConfig
+
+
+__ALL_MODELS__ = {
+    'Pix2Struct': Pix2Struct,
+    'Pix2StructTensorRT': Pix2StructTensorRT,
+    'InternVL': InternVL,
+    'InternVL_LMDeploy': InternVL_LMDeploy,
+}
+
+
+def get_model_name(model_path):
+    model_config = AutoConfig.from_pretrained(
+        model_path,
+        trust_remote_code=True,
+    )
+
+    if 'Pix2Struct' in model_config.architectures[0]:
+        model_name = 'Pix2Struct'
+    elif 'InternVL' in model_config.architectures[0]:
+        model_name = 'InternVL'
+    else:
+        raise ValueError(f"Unsupported model type: {model_config.architectures[0]}")
+
+    return model_name
+
+
+def build_model(
+        model_ckpt='U4R/StructTable-InternVL2-1B',
+        cache_dir=None,
+        local_files_only=None,
+        **kwargs,
+    ):
+    model_name = get_model_name(model_ckpt)
+    if model_name == 'InternVL' and kwargs.get('lmdeploy', False):
+        model_name = 'InternVL_LMDeploy'
+    elif model_name == 'Pix2Struct' and kwargs.get('tensorrt_path', None):
+        model_name = 'Pix2StructTensorRT'
+
+    model = __ALL_MODELS__[model_name](
+        model_ckpt,
+        cache_dir=cache_dir,
+        local_files_only=local_files_only,
+        **kwargs
+    )
+    return model

doc_page_extractor-0.2.4/doc_page_extractor/struct_eqtable/internvl/__init__.py

@@ -0,0 +1,2 @@
+from .internvl import InternVL
+from .internvl_lmdeploy import InternVL_LMDeploy