processors.py 5.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186
  1. # copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. import os
  15. import os.path as osp
  16. import re
  17. import numpy as np
  18. from PIL import Image
  19. import cv2
  20. import math
  21. import json
  22. import tempfile
  23. from tokenizers import Tokenizer as TokenizerFast
  24. from ....utils import logging
  25. class OCRReisizeNormImg:
  26. """for ocr image resize and normalization"""
  27. def __init__(self, rec_image_shape=[3, 48, 320]):
  28. super().__init__()
  29. self.rec_image_shape = rec_image_shape
  30. def resize_norm_img(self, img, max_wh_ratio):
  31. """resize and normalize the img"""
  32. imgC, imgH, imgW = self.rec_image_shape
  33. assert imgC == img.shape[2]
  34. imgW = int((imgH * max_wh_ratio))
  35. h, w = img.shape[:2]
  36. ratio = w / float(h)
  37. if math.ceil(imgH * ratio) > imgW:
  38. resized_w = imgW
  39. else:
  40. resized_w = int(math.ceil(imgH * ratio))
  41. resized_image = cv2.resize(img, (resized_w, imgH))
  42. resized_image = resized_image.astype("float32")
  43. resized_image = resized_image.transpose((2, 0, 1)) / 255
  44. resized_image -= 0.5
  45. resized_image /= 0.5
  46. padding_im = np.zeros((imgC, imgH, imgW), dtype=np.float32)
  47. padding_im[:, :, 0:resized_w] = resized_image
  48. return padding_im
  49. def __call__(self, imgs):
  50. """apply"""
  51. return [self.resize(img) for img in imgs]
  52. def resize(self, img):
  53. imgC, imgH, imgW = self.rec_image_shape
  54. max_wh_ratio = imgW / imgH
  55. h, w = img.shape[:2]
  56. wh_ratio = w * 1.0 / h
  57. max_wh_ratio = max(max_wh_ratio, wh_ratio)
  58. img = self.resize_norm_img(img, max_wh_ratio)
  59. return img
  60. class BaseRecLabelDecode:
  61. """Convert between text-label and text-index"""
  62. def __init__(self, character_str=None, use_space_char=True):
  63. super().__init__()
  64. self.reverse = False
  65. character_list = (
  66. list(character_str)
  67. if character_str is not None
  68. else list("0123456789abcdefghijklmnopqrstuvwxyz")
  69. )
  70. if use_space_char:
  71. character_list.append(" ")
  72. character_list = self.add_special_char(character_list)
  73. self.dict = {}
  74. for i, char in enumerate(character_list):
  75. self.dict[char] = i
  76. self.character = character_list
  77. def pred_reverse(self, pred):
  78. """pred_reverse"""
  79. pred_re = []
  80. c_current = ""
  81. for c in pred:
  82. if not bool(re.search("[a-zA-Z0-9 :*./%+-]", c)):
  83. if c_current != "":
  84. pred_re.append(c_current)
  85. pred_re.append(c)
  86. c_current = ""
  87. else:
  88. c_current += c
  89. if c_current != "":
  90. pred_re.append(c_current)
  91. return "".join(pred_re[::-1])
  92. def add_special_char(self, character_list):
  93. """add_special_char"""
  94. return character_list
  95. def decode(self, text_index, text_prob=None, is_remove_duplicate=False):
  96. """convert text-index into text-label."""
  97. result_list = []
  98. ignored_tokens = self.get_ignored_tokens()
  99. batch_size = len(text_index)
  100. for batch_idx in range(batch_size):
  101. selection = np.ones(len(text_index[batch_idx]), dtype=bool)
  102. if is_remove_duplicate:
  103. selection[1:] = text_index[batch_idx][1:] != text_index[batch_idx][:-1]
  104. for ignored_token in ignored_tokens:
  105. selection &= text_index[batch_idx] != ignored_token
  106. char_list = [
  107. self.character[text_id] for text_id in text_index[batch_idx][selection]
  108. ]
  109. if text_prob is not None:
  110. conf_list = text_prob[batch_idx][selection]
  111. else:
  112. conf_list = [1] * len(selection)
  113. if len(conf_list) == 0:
  114. conf_list = [0]
  115. text = "".join(char_list)
  116. if self.reverse: # for arabic rec
  117. text = self.pred_reverse(text)
  118. result_list.append((text, np.mean(conf_list).tolist()))
  119. return result_list
  120. def get_ignored_tokens(self):
  121. """get_ignored_tokens"""
  122. return [0] # for ctc blank
  123. def __call__(self, pred):
  124. """apply"""
  125. preds = np.array(pred)
  126. if isinstance(preds, tuple) or isinstance(preds, list):
  127. preds = preds[-1]
  128. preds_idx = preds.argmax(axis=-1)
  129. preds_prob = preds.max(axis=-1)
  130. text = self.decode(preds_idx, preds_prob, is_remove_duplicate=True)
  131. texts = []
  132. scores = []
  133. for t in text:
  134. texts.append(t[0])
  135. scores.append(t[1])
  136. return texts, scores
  137. class CTCLabelDecode(BaseRecLabelDecode):
  138. """Convert between text-label and text-index"""
  139. def __init__(self, character_list=None, use_space_char=True):
  140. super().__init__(character_list, use_space_char=use_space_char)
  141. def __call__(self, pred):
  142. """apply"""
  143. preds = np.array(pred[0])
  144. preds_idx = preds.argmax(axis=-1)
  145. preds_prob = preds.max(axis=-1)
  146. text = self.decode(preds_idx, preds_prob, is_remove_duplicate=True)
  147. texts = []
  148. scores = []
  149. for t in text:
  150. texts.append(t[0])
  151. scores.append(t[1])
  152. return texts, scores
  153. def add_special_char(self, character_list):
  154. """add_special_char"""
  155. character_list = ["blank"] + character_list
  156. return character_list