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

doc_page_extractor/layout_order.py

@@ -1,240 +0,0 @@
-import os
-import torch
-
-from typing import Generator
-from dataclasses import dataclass
-from transformers import LayoutLMv3ForTokenClassification
-
-from .types import Layout, LayoutClass
-from .layoutreader import prepare_inputs, boxes2inputs, parse_logits
-from .utils import ensure_dir
-
-
-@dataclass
-class _BBox:
-  layout_index: int
-  fragment_index: int
-  virtual: bool
-  order: int
-  value: tuple[float, float, float, float]
-
-class LayoutOrder:
-  def __init__(self, model_path: str):
-    self._model_path: str = model_path
-    self._model: LayoutLMv3ForTokenClassification | None = None
-
-  def _get_model(self) -> LayoutLMv3ForTokenClassification:
-    if self._model is None:
-      model_path = ensure_dir(self._model_path)
-      self._model = LayoutLMv3ForTokenClassification.from_pretrained(
-        pretrained_model_name_or_path="hantian/layoutreader",
-        cache_dir=model_path,
-        local_files_only=os.path.exists(os.path.join(model_path, "models--hantian--layoutreader")),
-      )
-    return self._model
-
-  def sort(self, layouts: list[Layout], size: tuple[int, int]) -> list[Layout]:
-    width, height = size
-    if width == 0 or height == 0:
-      return layouts
-
-    bbox_list = self._order_and_get_bbox_list(
-      layouts=layouts,
-      width=width,
-      height=height,
-    )
-    if bbox_list is None:
-      return layouts
-
-    return self._sort_layouts_and_fragments(layouts, bbox_list)
-
-  def _order_and_get_bbox_list(
-      self,
-      layouts: list[Layout],
-      width: int,
-      height: int,
-    ) -> list[_BBox] | None:
-
-    line_height = self._line_height(layouts)
-    bbox_list: list[_BBox] = []
-
-    for i, layout in enumerate(layouts):
-      if layout.cls == LayoutClass.PLAIN_TEXT and \
-         len(layout.fragments) > 0:
-        for j, fragment in enumerate(layout.fragments):
-          bbox_list.append(_BBox(
-            layout_index=i,
-            fragment_index=j,
-            virtual=False,
-            order=0,
-            value=fragment.rect.wrapper,
-          ))
-      else:
-        bbox_list.extend(
-          self._generate_virtual_lines(
-            layout=layout,
-            layout_index=i,
-            line_height=line_height,
-            width=width,
-            height=height,
-          ),
-        )
-
-    if len(bbox_list) > 200:
-      # https://github.com/opendatalab/MinerU/blob/980f5c8cd70f22f8c0c9b7b40eaff6f4804e6524/magic_pdf/pdf_parse_union_core_v2.py#L522
-      return None
-
-    layoutreader_size = 1000.0
-    x_scale = layoutreader_size / float(width)
-    y_scale = layoutreader_size / float(height)
-
-    for bbox in bbox_list:
-      x0, y0, x1, y1 = self._squeeze(bbox.value, width, height)
-      x0 = round(x0 * x_scale)
-      y0 = round(y0 * y_scale)
-      x1 = round(x1 * x_scale)
-      y1 = round(y1 * y_scale)
-      bbox.value = (x0, y0, x1, y1)
-
-    bbox_list.sort(key=lambda b: b.value) # 必须排序，乱序传入 layoutreader 会令它无法识别正确顺序
-    model = self._get_model()
-
-    with torch.no_grad():
-      inputs = boxes2inputs([list(bbox.value) for bbox in bbox_list])
-      inputs = prepare_inputs(inputs, model)
-      logits = model(**inputs).logits.cpu().squeeze(0)
-      orders = parse_logits(logits, len(bbox_list))
-
-    sorted_bbox_list = [bbox_list[i] for i in orders]
-    for i, bbox in enumerate(sorted_bbox_list):
-      bbox.order = i
-
-    return sorted_bbox_list
-
-  def _sort_layouts_and_fragments(self, layouts: list[Layout], bbox_list: list[_BBox]):
-    layout_bbox_list: list[list[_BBox]] = [[] for _ in range(len(layouts))]
-    for bbox in bbox_list:
-      layout_bbox_list[bbox.layout_index].append(bbox)
-
-    layouts_with_median_order: list[tuple[Layout, float]] = []
-    for layout_index, bbox_list in enumerate(layout_bbox_list):
-      layout = layouts[layout_index]
-      orders = [b.order for b in bbox_list] # virtual bbox 保证了 orders 不可能为空
-      median_order = self._median(orders)
-      layouts_with_median_order.append((layout, median_order))
-
-    for layout, bbox_list in zip(layouts, layout_bbox_list):
-      for bbox in bbox_list:
-        if not bbox.virtual:
-          layout.fragments[bbox.fragment_index].order = bbox.order
-      if all(not bbox.virtual for bbox in bbox_list):
-        layout.fragments.sort(key=lambda f: f.order)
-
-    layouts_with_median_order.sort(key=lambda x: x[1])
-    layouts = [layout for layout, _ in layouts_with_median_order]
-    next_fragment_order: int = 0
-
-    for layout in layouts:
-      for fragment in layout.fragments:
-        fragment.order = next_fragment_order
-        next_fragment_order += 1
-
-    return layouts
-
-  def _line_height(self, layouts: list[Layout]) -> float:
-    line_height: float = 0.0
-    count: int = 0
-    for layout in layouts:
-      for fragment in layout.fragments:
-        _, height = fragment.rect.size
-        line_height += height
-        count += 1
-    if count == 0:
-      return 10.0
-    return line_height / float(count)
-
-  def _generate_virtual_lines(
-        self,
-        layout: Layout,
-        layout_index: int,
-        line_height: float,
-        width: int,
-        height: int,
-      ) -> Generator[_BBox, None, None]:
-
-    # https://github.com/opendatalab/MinerU/blob/980f5c8cd70f22f8c0c9b7b40eaff6f4804e6524/magic_pdf/pdf_parse_union_core_v2.py#L451-L490
-    x0, y0, x1, y1 = layout.rect.wrapper
-    layout_height = y1 - y0
-    layout_weight = x1 - x0
-    lines = int(layout_height / line_height)
-
-    if layout_height <= line_height * 2:
-      yield _BBox(
-        layout_index=layout_index,
-        fragment_index=0,
-        virtual=True,
-        order=0,
-        value=(x0, y0, x1, y1),
-      )
-      return
-
-    elif layout_height <= height * 0.25 or \
-         width * 0.5 <= layout_weight or \
-         width * 0.25 < layout_weight:
-      if layout_weight > width * 0.4:
-        lines = 3
-      elif layout_weight <= width * 0.25:
-        if layout_height / layout_weight > 1.2:  # 细长的不分
-          yield _BBox(
-            layout_index=layout_index,
-            fragment_index=0,
-            virtual=True,
-            order=0,
-            value=(x0, y0, x1, y1),
-          )
-          return
-        else:  # 不细长的还是分成两行
-          lines = 2
-
-    lines = max(1, lines)
-    line_height = (y1 - y0) / lines
-    current_y = y0
-
-    for i in range(lines):
-      yield _BBox(
-        layout_index=layout_index,
-        fragment_index=i,
-        virtual=True,
-        order=0,
-        value=(x0, current_y, x1, current_y + line_height),
-      )
-      current_y += line_height
-
-  def _median(self, numbers: list[int]) -> float:
-    sorted_numbers = sorted(numbers)
-    n = len(sorted_numbers)
-
-    # 判断是奇数还是偶数个元素
-    if n % 2 == 1:
-      # 奇数情况，直接取中间的数
-      return float(sorted_numbers[n // 2])
-    else:
-      # 偶数情况，取中间两个数的平均值
-      mid1 = sorted_numbers[n // 2 - 1]
-      mid2 = sorted_numbers[n // 2]
-      return float((mid1 + mid2) / 2)
-
-  def _squeeze(self, bbox: _BBox, width: int, height: int) -> _BBox:
-    x0, y0, x1, y1 = bbox
-    x0 = self._squeeze_value(x0, width)
-    x1 = self._squeeze_value(x1, width)
-    y0 = self._squeeze_value(y0, height)
-    y1 = self._squeeze_value(y1, height)
-    return x0, y0, x1, y1
-
-  def _squeeze_value(self, position: float, size: int) -> float:
-    if position < 0:
-      position = 0.0
-    if position > size:
-      position = float(size)
-    return position

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/ocr.py

@@ -1,175 +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 OCRFragment
-from .rectangle import Rectangle
-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]
-
-  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_dir_path: str,
-    ):
-    self._device: Literal["cpu", "cuda"] = device
-    self._model_dir_path: str = model_dir_path
-    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 _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))
-    # 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,
-    )
-    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