doc-page-extractor 0.0.5__py3-none-any.whl → 0.0.7__py3-none-any.whl

doc_page_extractor/__init__.py

@@ -1,4 +1,4 @@
-from .extractor import PaddleLang, DocExtractor
+from .extractor import DocExtractor
 from .clipper import clip, clip_from_image
 from .plot import plot
 from .types import ExtractedResult, OCRFragment, LayoutClass, Layout

doc_page_extractor/downloader.py

@@ -5,9 +5,12 @@ from pathlib import Path
 
 def download(url: str, file_path: Path):
   response = requests.get(url, stream=True, timeout=60)
+  if response.status_code != 200:
+    raise FileNotFoundError(f"Failed to download file from {url}: {response.status_code}")
   try:
     with open(file_path, "wb") as file:
       file.write(response.content)
   except Exception as e:
-    os.remove(file_path)
+    if os.path.exists(file_path):
+      os.remove(file_path)
     raise e

doc_page_extractor/extractor.py

@@ -1,5 +1,4 @@
 import os
-import torch
 
 from typing import Literal, Iterable
 from pathlib import Path
@@ -8,7 +7,7 @@ from transformers import LayoutLMv3ForTokenClassification
 from doclayout_yolo import YOLOv10
 
 from .layoutreader import prepare_inputs, boxes2inputs, parse_logits
-from .ocr import OCR, PaddleLang
+from .ocr import OCR
 from .ocr_corrector import correct_fragments
 from .raw_optimizer import RawOptimizer
 from .rectangle import intersection_area, Rectangle
@@ -30,23 +29,18 @@ class DocExtractor:
     self._device: Literal["cpu", "cuda"] = device
     self._ocr_for_each_layouts: bool = ocr_for_each_layouts
     self._order_by_layoutreader: bool = order_by_layoutreader
-    self._ocr: OCR = OCR(device, os.path.join(model_dir_path, "paddle"))
+    self._ocr: OCR = OCR(device, model_dir_path)
     self._yolo: YOLOv10 | None = None
     self._layout: LayoutLMv3ForTokenClassification | None = None
 
-    if self._device.startswith("cuda") and not torch.cuda.is_available():
-      self._device = "cpu"
-      print("Warn: cuda is not available, use cpu instead")
-
   def extract(
       self,
       image: Image,
-      lang: PaddleLang,
       adjust_points: bool = False,
     ) -> ExtractedResult:
 
     raw_optimizer = RawOptimizer(image, adjust_points)
-    fragments = list(self._ocr.search_fragments(raw_optimizer.image_np, lang))
+    fragments = list(self._ocr.search_fragments(raw_optimizer.image_np))
     raw_optimizer.receive_raw_fragments(fragments)
 
     layouts = self._get_layouts(raw_optimizer.image)
@@ -54,7 +48,7 @@ class DocExtractor:
     layouts = remove_overlap_layouts(layouts)
 
     if self._ocr_for_each_layouts:
-      self._correct_fragments_by_ocr_layouts(raw_optimizer.image, layouts, lang)
+      self._correct_fragments_by_ocr_layouts(raw_optimizer.image, layouts)
 
     if self._order_by_layoutreader:
       width, height = raw_optimizer.image.size
@@ -84,7 +78,7 @@ class DocExtractor:
       source=source,
       imgsz=1024,
       conf=0.2,
-      device=self._device    # Device to use (e.g., "cuda:0" or "cpu")
+      device=self._device    # Device to use (e.g., "cuda" or "cpu")
     )
     boxes = det_res[0].__dict__["boxes"]
     layouts: list[Layout] = []
@@ -118,9 +112,9 @@ class DocExtractor:
           break
     return layouts
 
-  def _correct_fragments_by_ocr_layouts(self, source: Image, layouts: list[Layout], lang: PaddleLang):
+  def _correct_fragments_by_ocr_layouts(self, source: Image, layouts: list[Layout]):
     for layout in layouts:
-      correct_fragments(self._ocr, source, layout, lang)
+      correct_fragments(self._ocr, source, layout)
 
   def _split_layouts_by_group(self, layouts: list[Layout]):
     texts_layouts: list[Layout] = []

doc_page_extractor/ocr.py

@@ -1,18 +1,52 @@
-import os
 import numpy as np
 import cv2
+import os
 
-from typing import Any, Literal, Generator
-from paddleocr import PaddleOCR
+from typing import Literal, Generator
+from dataclasses import dataclass
+from .onnxocr import TextSystem
 from .types import OCRFragment
 from .rectangle import Rectangle
-from .utils import is_space_text, ensure_dir
+from .downloader import download
+from .utils import is_space_text
+
+
+_MODELS = (
+  ("ppocrv4", "rec", "rec.onnx"),
+  ("ppocrv4", "cls", "cls.onnx"),
+  ("ppocrv4", "det", "det.onnx"),
+  ("ch_ppocr_server_v2.0", "ppocr_keys_v1.txt"),
+)
 
+@dataclass
+class _OONXParams:
+  use_angle_cls: bool
+  use_gpu: bool
+  rec_image_shape: tuple[int, int, int]
+  cls_image_shape: tuple[int, int, int]
+  cls_batch_num: int
+  cls_thresh: float
+  label_list: list[str]
 
-# https://github.com/PaddlePaddle/PaddleOCR/blob/2c0c4beb0606819735a16083cdebf652939c781a/paddleocr.py#L108-L157
-PaddleLang = Literal["ch", "en", "korean", "japan", "chinese_cht", "ta", "te", "ka", "latin", "arabic", "cyrillic", "devanagari"]
+  det_algorithm: str
+  det_limit_side_len: int
+  det_limit_type: str
+  det_db_thresh: float
+  det_db_box_thresh: float
+  det_db_unclip_ratio: float
+  use_dilation: bool
+  det_db_score_mode: str
+  det_box_type: str
+  rec_batch_num: int
+  drop_score: float
+  save_crop_res: bool
+  rec_algorithm: str
+  use_space_char: bool
+  rec_model_dir: str
+  cls_model_dir: str
+  det_model_dir: str
+  rec_char_dict_path: str
 
-# https://paddlepaddle.github.io/PaddleOCR/latest/quick_start.html#_2
 class OCR:
   def __init__(
       self,
@@ -21,60 +55,78 @@ class OCR:
     ):
     self._device: Literal["cpu", "cuda"] = device
     self._model_dir_path: str = model_dir_path
-    self._ocr_and_lan: tuple[PaddleOCR, PaddleLang] | None = None
+    self._text_system: TextSystem | None = None
 
-  def search_fragments(self, image: np.ndarray, lang: PaddleLang) -> Generator[OCRFragment, None, None]:
+  def search_fragments(self, image: np.ndarray) -> Generator[OCRFragment, None, None]:
     index: int = 0
-    for item in self._handle(lang, image):
-      for line in item:
-        react: list[list[float]] = line[0]
-        text, rank = line[1]
-        if is_space_text(text):
-          continue
-        yield OCRFragment(
-          order=index,
-          text=text,
-          rank=rank,
-          rect=Rectangle(
-            lt=(react[0][0], react[0][1]),
-            rt=(react[1][0], react[1][1]),
-            rb=(react[2][0], react[2][1]),
-            lb=(react[3][0], react[3][1]),
-          ),
-        )
-        index += 1
-
-  def _handle(self, lang: PaddleLang, image: np.ndarray) -> list[Any]:
-    ocr = self._get_ocr(lang)
+    for box, res in self._ocr(image):
+      text, rank = res
+      if is_space_text(text):
+        continue
+      yield OCRFragment(
+        order=index,
+        text=text,
+        rank=rank,
+        rect=Rectangle(
+          lt=(box[0][0], box[0][1]),
+          rt=(box[1][0], box[1][1]),
+          rb=(box[2][0], box[2][1]),
+          lb=(box[3][0], box[3][1]),
+        ),
+      )
+      index += 1
+
+  def _ocr(self, image: np.ndarray) -> Generator[tuple[list[list[float]], tuple[str, float]], None, None]:
+    text_system = self._get_text_system()
     image = self._preprocess_image(image)
-    # about img parameter to see
-    # https://github.com/PaddlePaddle/PaddleOCR/blob/2c0c4beb0606819735a16083cdebf652939c781a/paddleocr.py#L582-L619
-    ocr_list = ocr.ocr(img=image, cls=True)
-    # there will be some None
-    return [e for e in ocr_list if e is not None]
-
-  def _get_ocr(self, lang: PaddleLang) -> PaddleOCR:
-    if self._ocr_and_lan is not None:
-      ocr, origin_lang = self._ocr_and_lan
-      if lang == origin_lang:
-        return ocr
-
-    ocr = PaddleOCR(
-      lang=lang,
-      use_angle_cls=True,
-      use_gpu=self._device.startswith("cuda"),
-      det_model_dir=ensure_dir(
-        os.path.join(self._model_dir_path, "det"),
-      ),
-      rec_model_dir=ensure_dir(
-        os.path.join(self._model_dir_path, "rec"),
-      ),
-      cls_model_dir=ensure_dir(
-        os.path.join(self._model_dir_path, "cls"),
-      ),
-    )
-    self._ocr_and_lan = (ocr, lang)
-    return ocr
+    dt_boxes, rec_res = text_system(image)
+
+    for box, res in zip(dt_boxes, rec_res):
+      yield box.tolist(), res
+
+  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)
+
+      self._text_system = TextSystem(_OONXParams(
+        use_angle_cls=True,
+        use_gpu=(self._device != "cpu"),
+        rec_image_shape=(3, 48, 320),
+        cls_image_shape=(3, 48, 192),
+        cls_batch_num=6,
+        cls_thresh=0.9,
+        label_list=["0", "180"],
+        det_algorithm="DB",
+        det_limit_side_len=960,
+        det_limit_type="max",
+        det_db_thresh=0.3,
+        det_db_box_thresh=0.6,
+        det_db_unclip_ratio=1.5,
+        use_dilation=False,
+        det_db_score_mode="fast",
+        det_box_type="quad",
+        rec_batch_num=6,
+        drop_score=0.5,
+        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]),
+      ))
+
+    return self._text_system
 
   def _preprocess_image(self, image: np.ndarray) -> np.ndarray:
     image = self._alpha_to_color(image, (255, 255, 255))

doc_page_extractor/ocr_corrector.py

@@ -4,14 +4,14 @@ from typing import Iterable
 from shapely.geometry import Polygon
 from PIL.Image import new, Image, Resampling
 from .types import Layout, OCRFragment
-from .ocr import OCR, PaddleLang
+from .ocr import OCR
 from .overlap import overlap_rate
 from .rectangle import Point, Rectangle
 
 
 _MIN_RATE = 0.5
 
-def correct_fragments(ocr: OCR, source: Image, layout: Layout, lang: PaddleLang):
+def correct_fragments(ocr: OCR, source: Image, layout: Layout):
   x1, y1, x2, y2 = layout.rect.wrapper
   image: Image = source.crop((
     round(x1), round(y1),
@@ -19,7 +19,7 @@ def correct_fragments(ocr: OCR, source: Image, layout: Layout, lang: PaddleLang)
   ))
   image, dx, dy, scale = _adjust_image(image)
   image_np = np.array(image)
-  ocr_fragments = list(ocr.search_fragments(image_np, lang))
+  ocr_fragments = list(ocr.search_fragments(image_np))
   corrected_fragments: list[OCRFragment] = []
 
   for fragment in ocr_fragments:

doc_page_extractor/onnxocr/__init__.py

@@ -0,0 +1 @@
+from .predict_system import TextSystem

doc_page_extractor/onnxocr/cls_postprocess.py

@@ -0,0 +1,26 @@
+class ClsPostProcess  (object):
+  """ Convert between text-label and text-index """
+
+  def __init__(self, label_list=None, key=None, **kwargs):
+    super(ClsPostProcess, self).__init__()
+    self.label_list = label_list
+    self.key = key
+
+  def __call__(self, preds, label=None, *args, **kwargs):
+    if self.key is not None:
+      preds = preds[self.key]
+
+    label_list = self.label_list
+    if label_list is None:
+      label_list = {idx: idx for idx in range(preds.shape[-1])}
+
+    # if isinstance(preds, paddle.Tensor):
+    #   preds = preds.numpy()
+
+    pred_idxs = preds.argmax(axis=1)
+    decode_out = [(label_list[idx], preds[i, idx])
+            for i, idx in enumerate(pred_idxs)]
+    if label is None:
+      return decode_out
+    label = [(label_list[idx], 1.0) for idx in label]
+    return decode_out, label

doc_page_extractor/onnxocr/db_postprocess.py

@@ -0,0 +1,246 @@
+# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+This code is refered from:
+https://github.com/WenmuZhou/DBNet.pytorch/blob/master/post_processing/seg_detector_representer.py
+"""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+import cv2
+# import paddle
+from shapely.geometry import Polygon
+import pyclipper
+
+
+class DBPostProcess(object):
+  """
+  The post process for Differentiable Binarization (DB).
+  """
+
+  def __init__(self,
+         thresh=0.3,
+         box_thresh=0.7,
+         max_candidates=1000,
+         unclip_ratio=2.0,
+         use_dilation=False,
+         score_mode="fast",
+         box_type='quad',
+         **kwargs):
+    self.thresh = thresh
+    self.box_thresh = box_thresh
+    self.max_candidates = max_candidates
+    self.unclip_ratio = unclip_ratio
+    self.min_size = 3
+    self.score_mode = score_mode
+    self.box_type = box_type
+    assert score_mode in [
+      "slow", "fast"
+    ], "Score mode must be in [slow, fast] but got: {}".format(score_mode)
+
+    self.dilation_kernel = None if not use_dilation else np.array(
+      [[1, 1], [1, 1]])
+
+  def polygons_from_bitmap(self, pred, _bitmap, dest_width, dest_height):
+    '''
+    _bitmap: single map with shape (1, H, W),
+      whose values are binarized as {0, 1}
+    '''
+
+    bitmap = _bitmap
+    height, width = bitmap.shape
+
+    boxes = []
+    scores = []
+
+    contours, _ = cv2.findContours((bitmap * 255).astype(np.uint8),
+                     cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
+
+    for contour in contours[:self.max_candidates]:
+      epsilon = 0.002 * cv2.arcLength(contour, True)
+      approx = cv2.approxPolyDP(contour, epsilon, True)
+      points = approx.reshape((-1, 2))
+      if points.shape[0] < 4:
+        continue
+
+      score = self.box_score_fast(pred, points.reshape(-1, 2))
+      if self.box_thresh > score:
+        continue
+
+      if points.shape[0] > 2:
+        box = self.unclip(points, self.unclip_ratio)
+        if len(box) > 1:
+          continue
+      else:
+        continue
+      box = box.reshape(-1, 2)
+
+      _, sside = self.get_mini_boxes(box.reshape((-1, 1, 2)))
+      if sside < self.min_size + 2:
+        continue
+
+      box = np.array(box)
+      box[:, 0] = np.clip(
+        np.round(box[:, 0] / width * dest_width), 0, dest_width)
+      box[:, 1] = np.clip(
+        np.round(box[:, 1] / height * dest_height), 0, dest_height)
+      boxes.append(box.tolist())
+      scores.append(score)
+    return boxes, scores
+
+  def boxes_from_bitmap(self, pred, _bitmap, dest_width, dest_height):
+    '''
+    _bitmap: single map with shape (1, H, W),
+        whose values are binarized as {0, 1}
+    '''
+
+    bitmap = _bitmap
+    height, width = bitmap.shape
+
+    outs = cv2.findContours((bitmap * 255).astype(np.uint8), cv2.RETR_LIST,
+                cv2.CHAIN_APPROX_SIMPLE)
+    if len(outs) == 3:
+      img, contours, _ = outs[0], outs[1], outs[2]
+    elif len(outs) == 2:
+      contours, _ = outs[0], outs[1]
+
+    num_contours = min(len(contours), self.max_candidates)
+
+    boxes = []
+    scores = []
+    for index in range(num_contours):
+      contour = contours[index]
+      points, sside = self.get_mini_boxes(contour)
+      if sside < self.min_size:
+        continue
+      points = np.array(points)
+      if self.score_mode == "fast":
+        score = self.box_score_fast(pred, points.reshape(-1, 2))
+      else:
+        score = self.box_score_slow(pred, contour)
+      if self.box_thresh > score:
+        continue
+
+      box = self.unclip(points, self.unclip_ratio).reshape(-1, 1, 2)
+      box, sside = self.get_mini_boxes(box)
+      if sside < self.min_size + 2:
+        continue
+      box = np.array(box)
+
+      box[:, 0] = np.clip(
+        np.round(box[:, 0] / width * dest_width), 0, dest_width)
+      box[:, 1] = np.clip(
+        np.round(box[:, 1] / height * dest_height), 0, dest_height)
+      boxes.append(box.astype("int32"))
+      scores.append(score)
+    return np.array(boxes, dtype="int32"), scores
+
+  def unclip(self, box, unclip_ratio):
+    poly = Polygon(box)
+    distance = poly.area * unclip_ratio / poly.length
+    offset = pyclipper.PyclipperOffset()
+    offset.AddPath(box, pyclipper.JT_ROUND, pyclipper.ET_CLOSEDPOLYGON)
+    expanded = np.array(offset.Execute(distance))
+    return expanded
+
+  def get_mini_boxes(self, contour):
+    bounding_box = cv2.minAreaRect(contour)
+    points = sorted(list(cv2.boxPoints(bounding_box)), key=lambda x: x[0])
+
+    index_1, index_2, index_3, index_4 = 0, 1, 2, 3
+    if points[1][1] > points[0][1]:
+      index_1 = 0
+      index_4 = 1
+    else:
+      index_1 = 1
+      index_4 = 0
+    if points[3][1] > points[2][1]:
+      index_2 = 2
+      index_3 = 3
+    else:
+      index_2 = 3
+      index_3 = 2
+
+    box = [
+      points[index_1], points[index_2], points[index_3], points[index_4]
+    ]
+    return box, min(bounding_box[1])
+
+  def box_score_fast(self, bitmap, _box):
+    '''
+    box_score_fast: use bbox mean score as the mean score
+    '''
+    h, w = bitmap.shape[:2]
+    box = _box.copy()
+    xmin = np.clip(np.floor(box[:, 0].min()).astype("int32"), 0, w - 1)
+    xmax = np.clip(np.ceil(box[:, 0].max()).astype("int32"), 0, w - 1)
+    ymin = np.clip(np.floor(box[:, 1].min()).astype("int32"), 0, h - 1)
+    ymax = np.clip(np.ceil(box[:, 1].max()).astype("int32"), 0, h - 1)
+
+    mask = np.zeros((ymax - ymin + 1, xmax - xmin + 1), dtype=np.uint8)
+    box[:, 0] = box[:, 0] - xmin
+    box[:, 1] = box[:, 1] - ymin
+    cv2.fillPoly(mask, box.reshape(1, -1, 2).astype("int32"), 1)
+    return cv2.mean(bitmap[ymin:ymax + 1, xmin:xmax + 1], mask)[0]
+
+  def box_score_slow(self, bitmap, contour):
+    '''
+    box_score_slow: use polyon mean score as the mean score
+    '''
+    h, w = bitmap.shape[:2]
+    contour = contour.copy()
+    contour = np.reshape(contour, (-1, 2))
+
+    xmin = np.clip(np.min(contour[:, 0]), 0, w - 1)
+    xmax = np.clip(np.max(contour[:, 0]), 0, w - 1)
+    ymin = np.clip(np.min(contour[:, 1]), 0, h - 1)
+    ymax = np.clip(np.max(contour[:, 1]), 0, h - 1)
+
+    mask = np.zeros((ymax - ymin + 1, xmax - xmin + 1), dtype=np.uint8)
+
+    contour[:, 0] = contour[:, 0] - xmin
+    contour[:, 1] = contour[:, 1] - ymin
+
+    cv2.fillPoly(mask, contour.reshape(1, -1, 2).astype("int32"), 1)
+    return cv2.mean(bitmap[ymin:ymax + 1, xmin:xmax + 1], mask)[0]
+
+  def __call__(self, outs_dict, shape_list):
+    pred = outs_dict['maps']
+    # if isinstance(pred, paddle.Tensor):
+    #   pred = pred.numpy()
+    pred = pred[:, 0, :, :]
+    segmentation = pred > self.thresh
+
+    boxes_batch = []
+    for batch_index in range(pred.shape[0]):
+      src_h, src_w, ratio_h, ratio_w = shape_list[batch_index]
+      if self.dilation_kernel is not None:
+        mask = cv2.dilate(
+          np.array(segmentation[batch_index]).astype(np.uint8),
+          self.dilation_kernel)
+      else:
+        mask = segmentation[batch_index]
+      if self.box_type == 'poly':
+        boxes, scores = self.polygons_from_bitmap(pred[batch_index],
+                              mask, src_w, src_h)
+      elif self.box_type == 'quad':
+        boxes, scores = self.boxes_from_bitmap(pred[batch_index], mask,
+                             src_w, src_h)
+      else:
+        raise ValueError("box_type can only be one of ['quad', 'poly']")
+
+      boxes_batch.append({'points': boxes})
+    return boxes_batch

doc_page_extractor/onnxocr/imaug.py

@@ -0,0 +1,32 @@
+from .operators import *
+
+
+def transform(data, ops=None):
+  """transform"""
+  if ops is None:
+    ops = []
+  for op in ops:
+    data = op(data)
+    if data is None:
+      return None
+  return data
+
+
+def create_operators(op_param_list, global_config=None):
+  """
+  create operators based on the config
+
+  Args:
+    params(list): a dict list, used to create some operators
+  """
+  assert isinstance(op_param_list, list), "operator config should be a list"
+  ops = []
+  for operator in op_param_list:
+    assert isinstance(operator, dict) and len(operator) == 1, "yaml format error"
+    op_name = list(operator)[0]
+    param = {} if operator[op_name] is None else operator[op_name]
+    if global_config is not None:
+      param.update(global_config)
+    op = eval(op_name)(**param)
+    ops.append(op)
+  return ops