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

doc_page_extractor/onnxocr/predict_det.py

@@ -1,120 +0,0 @@
-import numpy as np
-from .imaug import transform, create_operators
-from .db_postprocess import DBPostProcess
-from .predict_base import PredictBase
-
-
-class TextDetector(PredictBase):
-  def __init__(self, args):
-    self.args = args
-    self.det_algorithm = args.det_algorithm
-    pre_process_list = [
-      {
-        "DetResizeForTest": {
-          "limit_side_len": args.det_limit_side_len,
-          "limit_type": args.det_limit_type,
-        }
-      },
-      {
-        "NormalizeImage": {
-          "std": [0.229, 0.224, 0.225],
-          "mean": [0.485, 0.456, 0.406],
-          "scale": "1./255.",
-          "order": "hwc",
-        }
-      },
-      {"ToCHWImage": None},
-      {"KeepKeys": {"keep_keys": ["image", "shape"]}},
-    ]
-    postprocess_params = {}
-    postprocess_params["name"] = "DBPostProcess"
-    postprocess_params["thresh"] = args.det_db_thresh
-    postprocess_params["box_thresh"] = args.det_db_box_thresh
-    postprocess_params["max_candidates"] = 1000
-    postprocess_params["unclip_ratio"] = args.det_db_unclip_ratio
-    postprocess_params["use_dilation"] = args.use_dilation
-    postprocess_params["score_mode"] = args.det_db_score_mode
-    postprocess_params["box_type"] = args.det_box_type
-
-    # 实例化预处理操作类
-    self.preprocess_op = create_operators(pre_process_list)
-    # self.postprocess_op = build_post_process(postprocess_params)
-    # 实例化后处理操作类
-    self.postprocess_op = DBPostProcess(**postprocess_params)
-
-    # 初始化模型
-    self.det_onnx_session = self.get_onnx_session(args.det_model_dir, args.use_gpu)
-    self.det_input_name = self.get_input_name(self.det_onnx_session)
-    self.det_output_name = self.get_output_name(self.det_onnx_session)
-
-  def order_points_clockwise(self, pts):
-    rect = np.zeros((4, 2), dtype="float32")
-    s = pts.sum(axis=1)
-    rect[0] = pts[np.argmin(s)]
-    rect[2] = pts[np.argmax(s)]
-    tmp = np.delete(pts, (np.argmin(s), np.argmax(s)), axis=0)
-    diff = np.diff(np.array(tmp), axis=1)
-    rect[1] = tmp[np.argmin(diff)]
-    rect[3] = tmp[np.argmax(diff)]
-    return rect
-
-  def clip_det_res(self, points, img_height, img_width):
-    for pno in range(points.shape[0]):
-      points[pno, 0] = int(min(max(points[pno, 0], 0), img_width - 1))
-      points[pno, 1] = int(min(max(points[pno, 1], 0), img_height - 1))
-    return points
-
-  def filter_tag_det_res(self, dt_boxes, image_shape):
-    img_height, img_width = image_shape[0:2]
-    dt_boxes_new = []
-    for box in dt_boxes:
-      if type(box) is list:
-        box = np.array(box)
-      box = self.order_points_clockwise(box)
-      box = self.clip_det_res(box, img_height, img_width)
-      rect_width = int(np.linalg.norm(box[0] - box[1]))
-      rect_height = int(np.linalg.norm(box[0] - box[3]))
-      if rect_width <= 3 or rect_height <= 3:
-        continue
-      dt_boxes_new.append(box)
-    dt_boxes = np.array(dt_boxes_new)
-    return dt_boxes
-
-  def filter_tag_det_res_only_clip(self, dt_boxes, image_shape):
-    img_height, img_width = image_shape[0:2]
-    dt_boxes_new = []
-    for box in dt_boxes:
-      if type(box) is list:
-        box = np.array(box)
-      box = self.clip_det_res(box, img_height, img_width)
-      dt_boxes_new.append(box)
-    dt_boxes = np.array(dt_boxes_new)
-    return dt_boxes
-
-  def __call__(self, img):
-    ori_im = img.copy()
-    data = {"image": img}
-
-    data = transform(data, self.preprocess_op)
-    img, shape_list = data
-    if img is None:
-      return None, 0
-    img = np.expand_dims(img, axis=0)
-    shape_list = np.expand_dims(shape_list, axis=0)
-    img = img.copy()
-
-    input_feed = self.get_input_feed(self.det_input_name, img)
-    outputs = self.det_onnx_session.run(self.det_output_name, input_feed=input_feed)
-
-    preds = {}
-    preds["maps"] = outputs[0]
-
-    post_result = self.postprocess_op(preds, shape_list)
-    dt_boxes = post_result[0]["points"]
-
-    if self.args.det_box_type == "poly":
-      dt_boxes = self.filter_tag_det_res_only_clip(dt_boxes, ori_im.shape)
-    else:
-      dt_boxes = self.filter_tag_det_res(dt_boxes, ori_im.shape)
-
-    return dt_boxes

doc_page_extractor/onnxocr/predict_rec.py

@@ -1,321 +0,0 @@
-import cv2
-import numpy as np
-import math
-from PIL import Image
-
-
-from .rec_postprocess import CTCLabelDecode
-from .predict_base import PredictBase
-
-
-class TextRecognizer(PredictBase):
-  def __init__(self, args):
-    self.rec_image_shape = args.rec_image_shape
-    self.rec_batch_num = args.rec_batch_num
-    self.rec_algorithm = args.rec_algorithm
-    self.postprocess_op = CTCLabelDecode(
-      character_dict_path=args.rec_char_dict_path,
-      use_space_char=args.use_space_char,
-    )
-
-    # 初始化模型
-    self.rec_onnx_session = self.get_onnx_session(args.rec_model_dir, args.use_gpu)
-    self.rec_input_name = self.get_input_name(self.rec_onnx_session)
-    self.rec_output_name = self.get_output_name(self.rec_onnx_session)
-
-  def resize_norm_img(self, img, max_wh_ratio):
-    imgC, imgH, imgW = self.rec_image_shape
-    if self.rec_algorithm == "NRTR" or self.rec_algorithm == "ViTSTR":
-      img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
-      # return padding_im
-      image_pil = Image.fromarray(np.uint8(img))
-      if self.rec_algorithm == "ViTSTR":
-        img = image_pil.resize([imgW, imgH], Image.BICUBIC)
-      else:
-        img = image_pil.resize([imgW, imgH], Image.ANTIALIAS)
-      img = np.array(img)
-      norm_img = np.expand_dims(img, -1)
-      norm_img = norm_img.transpose((2, 0, 1))
-      if self.rec_algorithm == "ViTSTR":
-        norm_img = norm_img.astype(np.float32) / 255.0
-      else:
-        norm_img = norm_img.astype(np.float32) / 128.0 - 1.0
-      return norm_img
-    elif self.rec_algorithm == "RFL":
-      img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
-      resized_image = cv2.resize(img, (imgW, imgH), interpolation=cv2.INTER_CUBIC)
-      resized_image = resized_image.astype("float32")
-      resized_image = resized_image / 255
-      resized_image = resized_image[np.newaxis, :]
-      resized_image -= 0.5
-      resized_image /= 0.5
-      return resized_image
-
-    assert imgC == img.shape[2]
-    imgW = int((imgH * max_wh_ratio))
-
-    # w = self.rec_onnx_session.get_inputs()[0].shape[3:][0]
-    # w = self.rec_onnx_session.get_inputs()[0].shape[3:][0]
-    # print(w)
-    # if w is not None and w > 0:
-    #   imgW = w
-
-    h, w = img.shape[:2]
-    ratio = w / float(h)
-    if math.ceil(imgH * ratio) > imgW:
-      resized_w = imgW
-    else:
-      resized_w = int(math.ceil(imgH * ratio))
-    if self.rec_algorithm == "RARE":
-      if resized_w > self.rec_image_shape[2]:
-        resized_w = self.rec_image_shape[2]
-      imgW = self.rec_image_shape[2]
-    resized_image = cv2.resize(img, (resized_w, imgH))
-    resized_image = resized_image.astype("float32")
-    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 resize_norm_img_vl(self, img, image_shape):
-    imgC, imgH, imgW = image_shape
-    img = img[:, :, ::-1]  # bgr2rgb
-    resized_image = cv2.resize(img, (imgW, imgH), interpolation=cv2.INTER_LINEAR)
-    resized_image = resized_image.astype("float32")
-    resized_image = resized_image.transpose((2, 0, 1)) / 255
-    return resized_image
-
-  def resize_norm_img_srn(self, img, image_shape):
-    imgC, imgH, imgW = image_shape
-
-    img_black = np.zeros((imgH, imgW))
-    im_hei = img.shape[0]
-    im_wid = img.shape[1]
-
-    if im_wid <= im_hei * 1:
-      img_new = cv2.resize(img, (imgH * 1, imgH))
-    elif im_wid <= im_hei * 2:
-      img_new = cv2.resize(img, (imgH * 2, imgH))
-    elif im_wid <= im_hei * 3:
-      img_new = cv2.resize(img, (imgH * 3, imgH))
-    else:
-      img_new = cv2.resize(img, (imgW, imgH))
-
-    img_np = np.asarray(img_new)
-    img_np = cv2.cvtColor(img_np, cv2.COLOR_BGR2GRAY)
-    img_black[:, 0 : img_np.shape[1]] = img_np
-    img_black = img_black[:, :, np.newaxis]
-
-    row, col, c = img_black.shape
-    c = 1
-
-    return np.reshape(img_black, (c, row, col)).astype(np.float32)
-
-  def srn_other_inputs(self, image_shape, num_heads, max_text_length):
-    imgC, imgH, imgW = image_shape
-    feature_dim = int((imgH / 8) * (imgW / 8))
-
-    encoder_word_pos = (
-      np.array(range(0, feature_dim)).reshape((feature_dim, 1)).astype("int64")
-    )
-    gsrm_word_pos = (
-      np.array(range(0, max_text_length))
-      .reshape((max_text_length, 1))
-      .astype("int64")
-    )
-
-    gsrm_attn_bias_data = np.ones((1, max_text_length, max_text_length))
-    gsrm_slf_attn_bias1 = np.triu(gsrm_attn_bias_data, 1).reshape(
-      [-1, 1, max_text_length, max_text_length]
-    )
-    gsrm_slf_attn_bias1 = np.tile(gsrm_slf_attn_bias1, [1, num_heads, 1, 1]).astype(
-      "float32"
-    ) * [-1e9]
-
-    gsrm_slf_attn_bias2 = np.tril(gsrm_attn_bias_data, -1).reshape(
-      [-1, 1, max_text_length, max_text_length]
-    )
-    gsrm_slf_attn_bias2 = np.tile(gsrm_slf_attn_bias2, [1, num_heads, 1, 1]).astype(
-      "float32"
-    ) * [-1e9]
-
-    encoder_word_pos = encoder_word_pos[np.newaxis, :]
-    gsrm_word_pos = gsrm_word_pos[np.newaxis, :]
-
-    return [
-      encoder_word_pos,
-      gsrm_word_pos,
-      gsrm_slf_attn_bias1,
-      gsrm_slf_attn_bias2,
-    ]
-
-  def process_image_srn(self, img, image_shape, num_heads, max_text_length):
-    norm_img = self.resize_norm_img_srn(img, image_shape)
-    norm_img = norm_img[np.newaxis, :]
-
-    [encoder_word_pos, gsrm_word_pos, gsrm_slf_attn_bias1, gsrm_slf_attn_bias2] = (
-      self.srn_other_inputs(image_shape, num_heads, max_text_length)
-    )
-
-    gsrm_slf_attn_bias1 = gsrm_slf_attn_bias1.astype(np.float32)
-    gsrm_slf_attn_bias2 = gsrm_slf_attn_bias2.astype(np.float32)
-    encoder_word_pos = encoder_word_pos.astype(np.int64)
-    gsrm_word_pos = gsrm_word_pos.astype(np.int64)
-
-    return (
-      norm_img,
-      encoder_word_pos,
-      gsrm_word_pos,
-      gsrm_slf_attn_bias1,
-      gsrm_slf_attn_bias2,
-    )
-
-  def resize_norm_img_sar(self, img, image_shape, width_downsample_ratio=0.25):
-    imgC, imgH, imgW_min, imgW_max = image_shape
-    h = img.shape[0]
-    w = img.shape[1]
-    valid_ratio = 1.0
-    # make sure new_width is an integral multiple of width_divisor.
-    width_divisor = int(1 / width_downsample_ratio)
-    # resize
-    ratio = w / float(h)
-    resize_w = math.ceil(imgH * ratio)
-    if resize_w % width_divisor != 0:
-      resize_w = round(resize_w / width_divisor) * width_divisor
-    if imgW_min is not None:
-      resize_w = max(imgW_min, resize_w)
-    if imgW_max is not None:
-      valid_ratio = min(1.0, 1.0 * resize_w / imgW_max)
-      resize_w = min(imgW_max, resize_w)
-    resized_image = cv2.resize(img, (resize_w, imgH))
-    resized_image = resized_image.astype("float32")
-    # norm
-    if 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
-    resize_shape = resized_image.shape
-    padding_im = -1.0 * np.ones((imgC, imgH, imgW_max), dtype=np.float32)
-    padding_im[:, :, 0:resize_w] = resized_image
-    pad_shape = padding_im.shape
-
-    return padding_im, resize_shape, pad_shape, valid_ratio
-
-  def resize_norm_img_spin(self, img):
-    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
-    # return padding_im
-    img = cv2.resize(img, tuple([100, 32]), cv2.INTER_CUBIC)
-    img = np.array(img, np.float32)
-    img = np.expand_dims(img, -1)
-    img = img.transpose((2, 0, 1))
-    mean = [127.5]
-    std = [127.5]
-    mean = np.array(mean, dtype=np.float32)
-    std = np.array(std, dtype=np.float32)
-    mean = np.float32(mean.reshape(1, -1))
-    stdinv = 1 / np.float32(std.reshape(1, -1))
-    img -= mean
-    img *= stdinv
-    return img
-
-  def resize_norm_img_svtr(self, img, image_shape):
-    imgC, imgH, imgW = image_shape
-    resized_image = cv2.resize(img, (imgW, imgH), interpolation=cv2.INTER_LINEAR)
-    resized_image = resized_image.astype("float32")
-    resized_image = resized_image.transpose((2, 0, 1)) / 255
-    resized_image -= 0.5
-    resized_image /= 0.5
-    return resized_image
-
-  def resize_norm_img_abinet(self, img, image_shape):
-    imgC, imgH, imgW = image_shape
-
-    resized_image = cv2.resize(img, (imgW, imgH), interpolation=cv2.INTER_LINEAR)
-    resized_image = resized_image.astype("float32")
-    resized_image = resized_image / 255.0
-
-    mean = np.array([0.485, 0.456, 0.406])
-    std = np.array([0.229, 0.224, 0.225])
-    resized_image = (resized_image - mean[None, None, ...]) / std[None, None, ...]
-    resized_image = resized_image.transpose((2, 0, 1))
-    resized_image = resized_image.astype("float32")
-
-    return resized_image
-
-  def norm_img_can(self, img, image_shape):
-    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)  # CAN only predict gray scale image
-
-    if self.inverse:
-      img = 255 - img
-
-    if self.rec_image_shape[0] == 1:
-      h, w = img.shape
-      _, imgH, imgW = self.rec_image_shape
-      if h < imgH or w < imgW:
-        padding_h = max(imgH - h, 0)
-        padding_w = max(imgW - w, 0)
-        img_padded = np.pad(
-          img,
-          ((0, padding_h), (0, padding_w)),
-          "constant",
-          constant_values=(255),
-        )
-        img = img_padded
-
-    img = np.expand_dims(img, 0) / 255.0  # h,w,c -> c,h,w
-    img = img.astype("float32")
-
-    return img
-
-  def __call__(self, 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 recognition process
-    indices = np.argsort(np.array(width_list))
-    rec_res = [["", 0.0]] * img_num
-    batch_num = self.rec_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 = []
-      imgC, imgH, imgW = self.rec_image_shape[:3]
-      max_wh_ratio = imgW / imgH
-      # 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]], max_wh_ratio)
-        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()
-
-      # img = img[:, :, ::-1].transpose(2, 0, 1)
-      # img = img[:, :, ::-1]
-      # img = img.transpose(2, 0, 1)
-      # img = img.astype(np.float32)
-      # img = np.expand_dims(img, axis=0)
-      # print(img.shape)
-      input_feed = self.get_input_feed(self.rec_input_name, norm_img_batch)
-      outputs = self.rec_onnx_session.run(
-        self.rec_output_name, input_feed=input_feed
-      )
-
-      preds = outputs[0]
-
-      rec_result = self.postprocess_op(preds)
-      for rno in range(len(rec_result)):
-        rec_res[indices[beg_img_no + rno]] = rec_result[rno]
-
-    return rec_res

doc_page_extractor/onnxocr/predict_system.py

@@ -1,97 +0,0 @@
-import os
-import cv2
-import copy
-
-from . import predict_det
-from . import predict_cls
-from . import predict_rec
-from .utils import get_rotate_crop_image, get_minarea_rect_crop
-
-class TextSystem:
-  def __init__(self, args):
-    self.text_detector = predict_det.TextDetector(args)
-    self.text_recognizer = predict_rec.TextRecognizer(args)
-    self.use_angle_cls = True
-    self.drop_score = args.drop_score
-    if self.use_angle_cls:
-      self.text_classifier = predict_cls.TextClassifier(args)
-
-    self.args = args
-    self.crop_image_res_index = 0
-
-  def draw_crop_rec_res(self, output_dir, img_crop_list, rec_res):
-    os.makedirs(output_dir, exist_ok=True)
-    bbox_num = len(img_crop_list)
-    for bno in range(bbox_num):
-      cv2.imwrite(
-        os.path.join(
-          output_dir, f"mg_crop_{bno + self.crop_image_res_index}.jpg"
-        ),
-        img_crop_list[bno],
-      )
-
-    self.crop_image_res_index += bbox_num
-
-  def __call__(self, img, cls=True):
-    ori_im = img.copy()
-    # 文字检测
-    dt_boxes = self.text_detector(img)
-
-    if dt_boxes is None:
-      return None, None
-
-    img_crop_list = []
-
-    dt_boxes = sorted_boxes(dt_boxes)
-
-    # 图片裁剪
-    for bno in range(len(dt_boxes)):
-      tmp_box = copy.deepcopy(dt_boxes[bno])
-      if self.args.det_box_type == "quad":
-        img_crop = get_rotate_crop_image(ori_im, tmp_box)
-      else:
-        img_crop = get_minarea_rect_crop(ori_im, tmp_box)
-      img_crop_list.append(img_crop)
-
-    # 方向分类
-    if self.use_angle_cls and cls:
-      img_crop_list, angle_list = self.text_classifier(img_crop_list)
-
-    # 图像识别
-    rec_res = self.text_recognizer(img_crop_list)
-
-    if self.args.save_crop_res:
-      self.draw_crop_rec_res(self.args.crop_res_save_dir, img_crop_list, rec_res)
-    filter_boxes, filter_rec_res = [], []
-    for box, rec_result in zip(dt_boxes, rec_res):
-      text, score = rec_result
-      if score >= self.drop_score:
-        filter_boxes.append(box)
-        filter_rec_res.append(rec_result)
-
-    return filter_boxes, filter_rec_res
-
-
-def sorted_boxes(dt_boxes):
-  """
-  Sort text boxes in order from top to bottom, left to right
-  args:
-    dt_boxes(array):detected text boxes with shape [4, 2]
-  return:
-    sorted boxes(array) with shape [4, 2]
-  """
-  num_boxes = dt_boxes.shape[0]
-  sorted_boxes = sorted(dt_boxes, key=lambda x: (x[0][1], x[0][0]))
-  _boxes = list(sorted_boxes)
-
-  for i in range(num_boxes - 1):
-    for j in range(i, -1, -1):
-      if abs(_boxes[j + 1][0][1] - _boxes[j][0][1]) < 10 and (
-        _boxes[j + 1][0][0] < _boxes[j][0][0]
-      ):
-        tmp = _boxes[j]
-        _boxes[j] = _boxes[j + 1]
-        _boxes[j + 1] = tmp
-      else:
-        break
-  return _boxes