doc-page-extractor 0.2.4__cp310-cp310-macosx_15_0_arm64.whl

doc_page_extractor/layoutreader.py

@@ -0,0 +1,126 @@
+# Copy from https://github.com/ppaanngggg/layoutreader/blob/main/v3/helpers.py
+from collections import defaultdict
+from typing import List, Dict
+
+import torch
+from transformers import LayoutLMv3ForTokenClassification
+
+MAX_LEN = 510
+CLS_TOKEN_ID = 0
+UNK_TOKEN_ID = 3
+EOS_TOKEN_ID = 2
+
+
+class DataCollator:
+  def __call__(self, features: List[dict]) -> Dict[str, torch.Tensor]:
+    bbox = []
+    labels = []
+    input_ids = []
+    attention_mask = []
+
+    # clip bbox and labels to max length, build input_ids and attention_mask
+    for feature in features:
+      _bbox = feature["source_boxes"]
+      if len(_bbox) > MAX_LEN:
+        _bbox = _bbox[:MAX_LEN]
+      _labels = feature["target_index"]
+      if len(_labels) > MAX_LEN:
+        _labels = _labels[:MAX_LEN]
+      _input_ids = [UNK_TOKEN_ID] * len(_bbox)
+      _attention_mask = [1] * len(_bbox)
+      assert len(_bbox) == len(_labels) == len(_input_ids) == len(_attention_mask)
+      bbox.append(_bbox)
+      labels.append(_labels)
+      input_ids.append(_input_ids)
+      attention_mask.append(_attention_mask)
+
+    # add CLS and EOS tokens
+    for i in range(len(bbox)):
+      bbox[i] = [[0, 0, 0, 0]] + bbox[i] + [[0, 0, 0, 0]]
+      labels[i] = [-100] + labels[i] + [-100]
+      input_ids[i] = [CLS_TOKEN_ID] + input_ids[i] + [EOS_TOKEN_ID]
+      attention_mask[i] = [1] + attention_mask[i] + [1]
+
+    # padding to max length
+    max_len = max(len(x) for x in bbox)
+    for i in range(len(bbox)):
+      bbox[i] = bbox[i] + [[0, 0, 0, 0]] * (max_len - len(bbox[i]))
+      labels[i] = labels[i] + [-100] * (max_len - len(labels[i]))
+      input_ids[i] = input_ids[i] + [EOS_TOKEN_ID] * (max_len - len(input_ids[i]))
+      attention_mask[i] = attention_mask[i] + [0] * (
+        max_len - len(attention_mask[i])
+      )
+
+    ret = {
+      "bbox": torch.tensor(bbox),
+      "attention_mask": torch.tensor(attention_mask),
+      "labels": torch.tensor(labels),
+      "input_ids": torch.tensor(input_ids),
+    }
+    # set label > MAX_LEN to -100, because original labels may be > MAX_LEN
+    ret["labels"][ret["labels"] > MAX_LEN] = -100
+    # set label > 0 to label-1, because original labels are 1-indexed
+    ret["labels"][ret["labels"] > 0] -= 1
+    return ret
+
+
+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]
+  return {
+    "bbox": torch.tensor([bbox]),
+    "attention_mask": torch.tensor([attention_mask]),
+    "input_ids": torch.tensor([input_ids]),
+  }
+
+
+def prepare_inputs(
+  inputs: Dict[str, torch.Tensor], model: LayoutLMv3ForTokenClassification
+) -> Dict[str, torch.Tensor]:
+  ret = {}
+  for k, v in inputs.items():
+    v = v.to(model.device)
+    if torch.is_floating_point(v):
+      v = v.to(model.dtype)
+    ret[k] = v
+  return ret
+
+
+def parse_logits(logits: torch.Tensor, length: int) -> List[int]:
+  """
+  parse logits to orders
+
+  :param logits: logits from model
+  :param length: input length
+  :return: orders
+  """
+  logits = logits[1 : length + 1, :length]
+  orders = logits.argsort(descending=False).tolist()
+  ret = [o.pop() for o in orders]
+  while True:
+    order_to_idxes = defaultdict(list)
+    for idx, order in enumerate(ret):
+      order_to_idxes[order].append(idx)
+    # filter idxes len > 1
+    order_to_idxes = {k: v for k, v in order_to_idxes.items() if len(v) > 1}
+    if not order_to_idxes:
+      break
+    # filter
+    for order, idxes in order_to_idxes.items():
+      # find original logits of idxes
+      idxes_to_logit = {}
+      for idx in idxes:
+        idxes_to_logit[idx] = logits[idx, order]
+      idxes_to_logit = sorted(
+        idxes_to_logit.items(), key=lambda x: x[1], reverse=True
+      )
+      # keep the highest logit as order, set others to next candidate
+      for idx, _ in idxes_to_logit[1:]:
+        ret[idx] = orders[idx].pop()
+
+  return ret
+
+
+def check_duplicate(a: List[int]) -> bool:
+  return len(a) != len(set(a))

doc_page_extractor/model.py

@@ -0,0 +1,133 @@
+from os import PathLike
+from time import sleep
+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
+
+
+_RETRY_TIMES = 6
+_RETRY_SLEEP = 3.5
+
+@runtime_checkable
+class Model(Protocol):
+  def get_onnx_ocr_path(self) -> Path:
+    raise NotImplementedError()
+
+  def get_yolo_path(self) -> Path:
+    raise NotImplementedError()
+
+  def get_layoutreader_path(self) -> Path:
+    raise NotImplementedError()
+
+  def get_struct_eqtable_path(self) -> Path:
+    raise NotImplementedError()
+
+  def get_latex_path(self) -> Path:
+    raise NotImplementedError()
+
+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="README.md",
+      repo_type=None,
+      is_snapshot=True,
+      wanna_dir_path=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,
+      wanna_dir_path=False,
+    )
+
+  def get_layoutreader_path(self) -> Path:
+    return self._get_model_path(
+      repo_id="hantian/layoutreader",
+      filename="model.safetensors",
+      repo_type=None,
+      is_snapshot=True,
+      wanna_dir_path=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,
+      wanna_dir_path=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,
+      wanna_dir_path=True,
+    )
+
+  def _get_model_path(
+        self,
+        repo_id: str,
+        filename: str,
+        repo_type: str | None,
+        is_snapshot: bool,
+        wanna_dir_path: bool,
+      ) -> Path:
+
+    with self._lock:
+      model_path = try_to_load_from_cache(
+        repo_id=repo_id,
+        filename=filename,
+        repo_type=repo_type,
+        cache_dir=self._model_cache_dir
+      )
+      if isinstance(model_path, str):
+        model_path = Path(model_path)
+        if wanna_dir_path:
+          for _ in Path(filename).parts:
+            model_path = model_path.parent
+
+      else:
+        # https://github.com/huggingface/huggingface_hub/issues/1542#issuecomment-1630465844
+        latest_error: ConnectionError | None = None
+        for i in range(_RETRY_TIMES + 1):
+          if latest_error is not None:
+            print(f"Retrying to download {repo_id} model, attempt {i + 1}/{_RETRY_TIMES}...")
+            sleep(_RETRY_SLEEP)
+          try:
+            if is_snapshot:
+              model_path = snapshot_download(
+                cache_dir=self._model_cache_dir,
+                repo_id=repo_id,
+                repo_type=repo_type,
+                resume_download=True,
+              )
+            else:
+              model_path = hf_hub_download(
+                cache_dir=self._model_cache_dir,
+                repo_id=repo_id,
+                repo_type=repo_type,
+                filename=filename,
+                resume_download=True,
+              )
+            latest_error = None
+          except ConnectionError as err:
+            latest_error = err
+
+        if latest_error is not None:
+          raise latest_error
+        model_path = Path(cast(PathLike, model_path))
+
+      return model_path

doc_page_extractor/ocr.py

@@ -0,0 +1,196 @@
+import numpy as np
+import cv2
+import os
+
+from typing import cast, Any, Iterable, Literal, Generator
+from dataclasses import dataclass
+from .onnxocr import TextSystem
+from .types import OCRFragment
+from .model import Model
+from .rectangle import Rectangle
+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]
+
+  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
+
+class OCR:
+  def __init__(self, device: Literal["cpu", "cuda"], model: Model):
+    self._device: Literal["cpu", "cuda"] = device
+    self._model: Model = model
+    self._text_system: TextSystem | None = None
+
+  def search_fragments(self, image: np.ndarray) -> Generator[OCRFragment, None, None]:
+    for box, res in self._ocr(image):
+      text, rank = res
+      if is_space_text(text):
+        continue
+
+      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]),
+      )
+      if not rect.is_valid or rect.area == 0.0:
+        continue
+
+      yield OCRFragment(
+        order=0,
+        text=text,
+        rank=rank,
+        rect=rect,
+      )
+
+  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)
+    dt_boxes, rec_res = text_system(image)
+
+    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:
+    if self._text_system is None:
+      model_paths = self._make_model_paths()
+      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=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
+
+  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, 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(
+      src=image,
+      dst=np.zeros((image.shape[0], image.shape[1])),
+      alpha=0,
+      beta=255,
+      norm_type=cv2.NORM_MINMAX,
+    )
+    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(cast(Any, 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
+
+  def _alpha_to_color(self, image: np.ndarray, alpha_color: tuple[float, float, float]) -> np.ndarray:
+    if len(image.shape) == 3 and image.shape[2] == 4:
+      B, G, R, A = cv2.split(image)
+      alpha = A / 255
+
+      R = (alpha_color[0] * (1 - alpha) + R * alpha).astype(np.uint8)
+      G = (alpha_color[1] * (1 - alpha) + G * alpha).astype(np.uint8)
+      B = (alpha_color[2] * (1 - alpha) + B * alpha).astype(np.uint8)
+
+      image = cv2.merge((B, G, R))
+
+    return image
+
+  def _binarize_img(self, image: np.ndarray):
+    if len(image.shape) == 3 and image.shape[2] == 3:
+      gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)  # conversion to grayscale image
+      # use cv2 threshold binarization
+      _, gray = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
+      image = cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR)
+    return image

doc_page_extractor/ocr_corrector.py

@@ -0,0 +1,126 @@
+import numpy as np
+
+from typing import cast, Iterable
+from shapely.geometry import Polygon
+from PIL.Image import new, Image, Resampling
+from .types import Layout, OCRFragment
+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):
+  x1, y1, x2, y2 = layout.rect.wrapper
+  image: Image = source.crop((
+    round(x1), round(y1),
+    round(x2), round(y2),
+  ))
+  image, dx, dy, scale = _adjust_image(image)
+  image_np = np.array(image)
+  ocr_fragments = list(ocr.search_fragments(image_np))
+  corrected_fragments: list[OCRFragment] = []
+
+  for fragment in ocr_fragments:
+    _apply_fragment(fragment.rect, layout, dx, dy, scale)
+
+  matched_fragments, not_matched_fragments = _match_fragments(
+    zone_rect=layout.rect,
+    fragments1=layout.fragments,
+    fragments2=ocr_fragments,
+  )
+  for fragment1, fragment2 in matched_fragments:
+    if fragment1.rank > fragment2.rank:
+      corrected_fragments.append(fragment1)
+    else:
+      corrected_fragments.append(fragment2)
+
+  corrected_fragments.extend(not_matched_fragments)
+  layout.fragments = corrected_fragments
+
+def _adjust_image(image: Image) -> tuple[Image, int, int, float]:
+  # after testing, adding white borders to images can reduce
+  # the possibility of some text not being recognized
+  border_size: int = 50
+  adjusted_size: int = 1024 - 2 * border_size
+  width, height = image.size
+  core_width = float(max(adjusted_size, width))
+  core_height = float(max(adjusted_size, height))
+
+  scale_x = core_width / width
+  scale_y = core_height / height
+  scale = min(scale_x, scale_y)
+  adjusted_width = width * scale
+  adjusted_height = height * scale
+
+  dx = (core_width - adjusted_width) / 2.0
+  dy = (core_height - adjusted_height) / 2.0
+  dx = round(dx) + border_size
+  dy = round(dy) + border_size
+
+  if scale != 1.0:
+    width = round(width * scale)
+    height = round(height * scale)
+    image = image.resize((width, height), Resampling.BICUBIC)
+
+  width = round(core_width) + 2 * border_size
+  height = round(core_height) + 2 * border_size
+  new_image = new("RGB", (width, height), (255, 255, 255))
+  new_image.paste(image, (dx, dy))
+
+  return new_image, dx, dy, scale
+
+def _apply_fragment(rect: Rectangle, layout: Layout, dx: int, dy: int, scale: float):
+  rect.lt = _apply_point(rect.lt, layout, dx, dy, scale)
+  rect.lb = _apply_point(rect.lb, layout, dx, dy, scale)
+  rect.rb = _apply_point(rect.rb, layout, dx, dy, scale)
+  rect.rt = _apply_point(rect.rt, layout, dx, dy, scale)
+
+def _apply_point(point: Point, layout: Layout, dx: int, dy: int, scale: float) -> Point:
+  x, y = point
+  x = (x - dx) / scale + layout.rect.lt[0]
+  y = (y - dy) / scale + layout.rect.lt[1]
+  return x, y
+
+def _match_fragments(
+    zone_rect: Rectangle,
+    fragments1: Iterable[OCRFragment],
+    fragments2: Iterable[OCRFragment],
+  ) -> tuple[list[tuple[OCRFragment, OCRFragment]], list[OCRFragment]]:
+
+  zone_polygon = Polygon(zone_rect)
+  fragments2 = list(fragments2)
+  matched_fragments: list[tuple[OCRFragment, OCRFragment]] = []
+  not_matched_fragments: list[OCRFragment] = []
+
+  for fragment1 in fragments1:
+    polygon1 = Polygon(fragment1.rect)
+    polygon1 = cast(Polygon, zone_polygon.intersection(polygon1))
+    if polygon1.is_empty:
+      continue
+
+    beast_j = -1
+    beast_rate = 0.0
+
+    for j, fragment2 in enumerate(fragments2):
+      polygon2 = Polygon(fragment2.rect)
+      rate = overlap_rate(polygon1, polygon2)
+      if rate < _MIN_RATE:
+        continue
+
+      if rate > beast_rate:
+        beast_j = j
+        beast_rate = rate
+
+    if beast_j != -1:
+      matched_fragments.append((
+        fragment1,
+        fragments2[beast_j],
+      ))
+      del fragments2[beast_j]
+    else:
+      not_matched_fragments.append(fragment1)
+
+  not_matched_fragments.extend(fragments2)
+  return matched_fragments, not_matched_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