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

doc_page_extractor/onnxocr/db_postprocess.py

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

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

doc_page_extractor/onnxocr/operators.py

@@ -1,187 +0,0 @@
-import numpy as np
-import cv2
-import sys
-import math
-
-
-class NormalizeImage(object):
-  """ normalize image such as substract mean, divide std
-  """
-
-  def __init__(self, scale=None, mean=None, std=None, order='chw', **kwargs):
-    if isinstance(scale, str):
-      scale = eval(scale)
-    self.scale = np.float32(scale if scale is not None else 1.0 / 255.0)
-    mean = mean if mean is not None else [0.485, 0.456, 0.406]
-    std = std if std is not None else [0.229, 0.224, 0.225]
-
-    shape = (3, 1, 1) if order == 'chw' else (1, 1, 3)
-    self.mean = np.array(mean).reshape(shape).astype('float32')
-    self.std = np.array(std).reshape(shape).astype('float32')
-
-  def __call__(self, data):
-    img = data['image']
-    from PIL import Image
-    if isinstance(img, Image.Image):
-      img = np.array(img)
-    assert isinstance(img,
-              np.ndarray), "invalid input 'img' in NormalizeImage"
-    data['image'] = (
-      img.astype('float32') * self.scale - self.mean) / self.std
-    return data
-
-
-class DetResizeForTest(object):
-  def __init__(self, **kwargs):
-    super(DetResizeForTest, self).__init__()
-    self.resize_type = 0
-    self.keep_ratio = False
-    if 'image_shape' in kwargs:
-      self.image_shape = kwargs['image_shape']
-      self.resize_type = 1
-      if 'keep_ratio' in kwargs:
-        self.keep_ratio = kwargs['keep_ratio']
-    elif 'limit_side_len' in kwargs:
-      self.limit_side_len = kwargs['limit_side_len']
-      self.limit_type = kwargs.get('limit_type', 'min')
-    elif 'resize_long' in kwargs:
-      self.resize_type = 2
-      self.resize_long = kwargs.get('resize_long', 960)
-    else:
-      self.limit_side_len = 736
-      self.limit_type = 'min'
-
-  def __call__(self, data):
-    img = data['image']
-    src_h, src_w, _ = img.shape
-    if sum([src_h, src_w]) < 64:
-      img = self.image_padding(img)
-
-    if self.resize_type == 0:
-      # img, shape = self.resize_image_type0(img)
-      img, [ratio_h, ratio_w] = self.resize_image_type0(img)
-    elif self.resize_type == 2:
-      img, [ratio_h, ratio_w] = self.resize_image_type2(img)
-    else:
-      # img, shape = self.resize_image_type1(img)
-      img, [ratio_h, ratio_w] = self.resize_image_type1(img)
-    data['image'] = img
-    data['shape'] = np.array([src_h, src_w, ratio_h, ratio_w])
-    return data
-
-  def image_padding(self, im, value=0):
-    h, w, c = im.shape
-    im_pad = np.zeros((max(32, h), max(32, w), c), np.uint8) + value
-    im_pad[:h, :w, :] = im
-    return im_pad
-
-  def resize_image_type1(self, img):
-    resize_h, resize_w = self.image_shape
-    ori_h, ori_w = img.shape[:2]  # (h, w, c)
-    if self.keep_ratio is True:
-      resize_w = ori_w * resize_h / ori_h
-      N = math.ceil(resize_w / 32)
-      resize_w = N * 32
-    ratio_h = float(resize_h) / ori_h
-    ratio_w = float(resize_w) / ori_w
-    img = cv2.resize(img, (int(resize_w), int(resize_h)))
-    # return img, np.array([ori_h, ori_w])
-    return img, [ratio_h, ratio_w]
-
-  def resize_image_type0(self, img):
-    """
-    resize image to a size multiple of 32 which is required by the network
-    args:
-      img(array): array with shape [h, w, c]
-    return(tuple):
-      img, (ratio_h, ratio_w)
-    """
-    limit_side_len = self.limit_side_len
-    h, w, c = img.shape
-
-    # limit the max side
-    if self.limit_type == 'max':
-      if max(h, w) > limit_side_len:
-        if h > w:
-          ratio = float(limit_side_len) / h
-        else:
-          ratio = float(limit_side_len) / w
-      else:
-        ratio = 1.
-    elif self.limit_type == 'min':
-      if min(h, w) < limit_side_len:
-        if h < w:
-          ratio = float(limit_side_len) / h
-        else:
-          ratio = float(limit_side_len) / w
-      else:
-        ratio = 1.
-    elif self.limit_type == 'resize_long':
-      ratio = float(limit_side_len) / max(h, w)
-    else:
-      raise Exception('not support limit type, image ')
-    resize_h = int(h * ratio)
-    resize_w = int(w * ratio)
-
-    resize_h = max(int(round(resize_h / 32) * 32), 32)
-    resize_w = max(int(round(resize_w / 32) * 32), 32)
-
-    try:
-      if int(resize_w) <= 0 or int(resize_h) <= 0:
-        return None, (None, None)
-      img = cv2.resize(img, (int(resize_w), int(resize_h)))
-    except:
-      print(img.shape, resize_w, resize_h)
-      sys.exit(0)
-    ratio_h = resize_h / float(h)
-    ratio_w = resize_w / float(w)
-    return img, [ratio_h, ratio_w]
-
-  def resize_image_type2(self, img):
-    h, w, _ = img.shape
-
-    resize_w = w
-    resize_h = h
-
-    if resize_h > resize_w:
-      ratio = float(self.resize_long) / resize_h
-    else:
-      ratio = float(self.resize_long) / resize_w
-
-    resize_h = int(resize_h * ratio)
-    resize_w = int(resize_w * ratio)
-
-    max_stride = 128
-    resize_h = (resize_h + max_stride - 1) // max_stride * max_stride
-    resize_w = (resize_w + max_stride - 1) // max_stride * max_stride
-    img = cv2.resize(img, (int(resize_w), int(resize_h)))
-    ratio_h = resize_h / float(h)
-    ratio_w = resize_w / float(w)
-
-    return img, [ratio_h, ratio_w]
-
-class ToCHWImage(object):
-  """ convert hwc image to chw image
-  """
-
-  def __init__(self, **kwargs):
-    pass
-
-  def __call__(self, data):
-    img = data['image']
-    from PIL import Image
-    if isinstance(img, Image.Image):
-      img = np.array(img)
-    data['image'] = img.transpose((2, 0, 1))
-    return data
-
-
-class KeepKeys(object):
-  def __init__(self, keep_keys, **kwargs):
-    self.keep_keys = keep_keys
-
-  def __call__(self, data):
-    data_list = []
-    for key in self.keep_keys:
-      data_list.append(data[key])
-    return data_list

doc_page_extractor/onnxocr/predict_base.py

@@ -1,57 +0,0 @@
-class PredictBase(object):
-  def __init__(self):
-    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
-    if use_gpu:
-      providers = providers=['CUDAExecutionProvider']
-    else:
-      providers = providers = ['CPUExecutionProvider']
-
-    onnx_session = self.onnxruntime.InferenceSession(model_dir, None, providers=providers)
-
-    # print("providers:", onnxruntime.get_device())
-    return onnx_session
-
-
-  def get_output_name(self, onnx_session):
-    """
-    output_name = onnx_session.get_outputs()[0].name
-    :param onnx_session:
-    :return:
-    """
-    output_name = []
-    for node in onnx_session.get_outputs():
-      output_name.append(node.name)
-    return output_name
-
-  def get_input_name(self, onnx_session):
-    """
-    input_name = onnx_session.get_inputs()[0].name
-    :param onnx_session:
-    :return:
-    """
-    input_name = []
-    for node in onnx_session.get_inputs():
-      input_name.append(node.name)
-    return input_name
-
-  def get_input_feed(self, input_name, image_numpy):
-    """
-    input_feed={self.input_name: image_numpy}
-    :param input_name:
-    :param image_numpy:
-    :return:
-    """
-    input_feed = {}
-    for name in input_name:
-      input_feed[name] = image_numpy
-    return input_feed

doc_page_extractor/onnxocr/predict_cls.py

@@ -1,109 +0,0 @@
-import cv2
-import copy
-import numpy as np
-import math
-
-from .cls_postprocess import ClsPostProcess
-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 = 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
-    h = img.shape[0]
-    w = img.shape[1]
-    ratio = w / float(h)
-    if math.ceil(imgH * ratio) > imgW:
-      resized_w = imgW
-    else:
-      resized_w = int(math.ceil(imgH * ratio))
-    resized_image = cv2.resize(img, (resized_w, imgH))
-    resized_image = resized_image.astype("float32")
-    if self.cls_image_shape[0] == 1:
-      resized_image = resized_image / 255
-      resized_image = resized_image[np.newaxis, :]
-    else:
-      resized_image = resized_image.transpose((2, 0, 1)) / 255
-    resized_image -= 0.5
-    resized_image /= 0.5
-    padding_im = np.zeros((imgC, imgH, imgW), dtype=np.float32)
-    padding_im[:, :, 0:resized_w] = resized_image
-    return padding_im
-
-  def __call__(self, img_list):
-    img_list = copy.deepcopy(img_list)
-    img_num = len(img_list)
-    # Calculate the aspect ratio of all text bars
-    width_list = []
-    for img in img_list:
-      width_list.append(img.shape[1] / float(img.shape[0]))
-    # Sorting can speed up the cls process
-    indices = np.argsort(np.array(width_list))
-
-    cls_res = [["", 0.0]] * img_num
-    batch_num = self.cls_batch_num
-
-    for beg_img_no in range(0, img_num, batch_num):
-
-      end_img_no = min(img_num, beg_img_no + batch_num)
-      norm_img_batch = []
-      max_wh_ratio = 0
-
-      for ino in range(beg_img_no, end_img_no):
-        h, w = img_list[indices[ino]].shape[0:2]
-        wh_ratio = w * 1.0 / h
-        max_wh_ratio = max(max_wh_ratio, wh_ratio)
-      for ino in range(beg_img_no, end_img_no):
-        norm_img = self.resize_norm_img(img_list[indices[ino]])
-        norm_img = norm_img[np.newaxis, :]
-        norm_img_batch.append(norm_img)
-      norm_img_batch = np.concatenate(norm_img_batch)
-      norm_img_batch = norm_img_batch.copy()
-
-      input_feed = self.get_input_feed(self.cls_input_name, norm_img_batch)
-      outputs = self.cls_onnx_session.run(
-        self.cls_output_name, input_feed=input_feed
-      )
-
-      prob_out = outputs[0]
-
-      cls_result = self.postprocess_op(prob_out)
-      for rno in range(len(cls_result)):
-        label, score = cls_result[rno]
-        cls_res[indices[beg_img_no + rno]] = [label, score]
-        if "180" in label and score > self.cls_thresh:
-          img_list[indices[beg_img_no + rno]] = cv2.rotate(
-            img_list[indices[beg_img_no + rno]], 1
-          )
-    return img_list, cls_res