doc-page-extractor 0.2.0__py3-none-any.whl → 1.0.2__py3-none-any.whl

doc_page_extractor/layoutreader.py

@@ -1,126 +0,0 @@
-# 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[int]]) -> 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/models.py

@@ -1,92 +0,0 @@
-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/ocr.py

@@ -1,200 +0,0 @@
-import numpy as np
-import cv2
-import os
-
-from typing import Literal, Generator
-from dataclasses import dataclass
-from .onnxocr import TextSystem
-from .types import GetModelDir, OCRFragment
-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"],
-      get_model_dir: GetModelDir,
-    ):
-    self._device: Literal["cpu", "cuda"] = device
-    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]:
-    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(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:
-      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 _preprocess_image(self, image: np.ndarray) -> np.ndarray:
-    image = self._alpha_to_color(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(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

@@ -1,126 +0,0 @@
-import numpy as np
-
-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
-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[OCRFragment] = 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)
-    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

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

doc_page_extractor/onnxocr/cls_postprocess.py

@@ -1,26 +0,0 @@
-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