transforms.py 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344
  1. # !/usr/bin/env python3
  2. # -*- coding: UTF-8 -*-
  3. ################################################################################
  4. #
  5. # Copyright (c) 2024 Baidu.com, Inc. All Rights Reserved
  6. #
  7. ################################################################################
  8. """
  9. Author: PaddlePaddle Authors
  10. """
  11. import os
  12. import os.path as osp
  13. import numpy as np
  14. from PIL import Image
  15. import cv2
  16. import paddle
  17. from ....utils import logging
  18. from ...base import BaseTransform
  19. from ...base.predictor.io.writers import ImageWriter
  20. from .keys import TableRecKeys as K
  21. __all__ = ['TableLabelDecode', 'TableMasterLabelDecode', 'SaveTableResults']
  22. class TableLabelDecode(BaseTransform):
  23. """ decode the table model outputs(probs) to character str"""
  24. def __init__(self,
  25. character_dict_type='TableAttn_ch',
  26. merge_no_span_structure=True):
  27. dict_character = []
  28. supported_dict = ['TableAttn_ch', 'TableAttn_en', 'TableMaster']
  29. if character_dict_type == 'TableAttn_ch':
  30. character_dict_name = 'table_structure_dict_ch.txt'
  31. elif character_dict_type == 'TableAttn_en':
  32. character_dict_name = 'table_structure_dict.txt'
  33. elif character_dict_type == 'TableMaster':
  34. character_dict_name = 'table_master_structure_dict.txt'
  35. else:
  36. assert False, " character_dict_type must in %s " \
  37. % supported_dict
  38. character_dict_path = osp.abspath(
  39. osp.join(osp.dirname(__file__), character_dict_name))
  40. with open(character_dict_path, "rb") as fin:
  41. lines = fin.readlines()
  42. for line in lines:
  43. line = line.decode('utf-8').strip("\n").strip("\r\n")
  44. dict_character.append(line)
  45. if merge_no_span_structure:
  46. if "<td></td>" not in dict_character:
  47. dict_character.append("<td></td>")
  48. if "<td>" in dict_character:
  49. dict_character.remove("<td>")
  50. dict_character = self.add_special_char(dict_character)
  51. self.dict = {}
  52. for i, char in enumerate(dict_character):
  53. self.dict[char] = i
  54. self.character = dict_character
  55. self.td_token = ['<td>', '<td', '<td></td>']
  56. def add_special_char(self, dict_character):
  57. """ add_special_char """
  58. self.beg_str = "sos"
  59. self.end_str = "eos"
  60. dict_character = dict_character
  61. dict_character = [self.beg_str] + dict_character + [self.end_str]
  62. return dict_character
  63. def get_ignored_tokens(self):
  64. """ get_ignored_tokens """
  65. beg_idx = self.get_beg_end_flag_idx("beg")
  66. end_idx = self.get_beg_end_flag_idx("end")
  67. return [beg_idx, end_idx]
  68. def get_beg_end_flag_idx(self, beg_or_end):
  69. """ get_beg_end_flag_idx """
  70. if beg_or_end == "beg":
  71. idx = np.array(self.dict[self.beg_str])
  72. elif beg_or_end == "end":
  73. idx = np.array(self.dict[self.end_str])
  74. else:
  75. assert False, "unsupported type %s in get_beg_end_flag_idx" \
  76. % beg_or_end
  77. return idx
  78. def apply(self, data):
  79. """ apply """
  80. shape_list = data[K.SHAPE_LIST]
  81. structure_probs = data[K.STRUCTURE_PROB]
  82. bbox_preds = data[K.LOC_PROB]
  83. if isinstance(structure_probs, paddle.Tensor):
  84. structure_probs = structure_probs.numpy()
  85. if isinstance(bbox_preds, paddle.Tensor):
  86. bbox_preds = bbox_preds.numpy()
  87. post_result = self.decode(structure_probs, bbox_preds, shape_list)
  88. structure_str_list = post_result['structure_batch_list'][0]
  89. bbox_list = post_result['bbox_batch_list'][0]
  90. structure_str_list = structure_str_list[0]
  91. structure_str_list = [
  92. '<html>', '<body>', '<table>'
  93. ] + structure_str_list + ['</table>', '</body>', '</html>']
  94. data[K.BBOX_RES] = bbox_list
  95. data[K.HTML_RES] = structure_str_list
  96. return data
  97. @classmethod
  98. def get_input_keys(cls):
  99. """ get input keys """
  100. return [K.STRUCTURE_PROB, K.LOC_PROB, K.SHAPE_LIST]
  101. @classmethod
  102. def get_output_keys(cls):
  103. """ get output keys """
  104. return [K.BBOX_RES, K.HTML_RES]
  105. def decode(self, structure_probs, bbox_preds, shape_list):
  106. """convert text-label into text-index.
  107. """
  108. ignored_tokens = self.get_ignored_tokens()
  109. end_idx = self.dict[self.end_str]
  110. structure_idx = structure_probs.argmax(axis=2)
  111. structure_probs = structure_probs.max(axis=2)
  112. structure_batch_list = []
  113. bbox_batch_list = []
  114. batch_size = len(structure_idx)
  115. for batch_idx in range(batch_size):
  116. structure_list = []
  117. bbox_list = []
  118. score_list = []
  119. for idx in range(len(structure_idx[batch_idx])):
  120. char_idx = int(structure_idx[batch_idx][idx])
  121. if idx > 0 and char_idx == end_idx:
  122. break
  123. if char_idx in ignored_tokens:
  124. continue
  125. text = self.character[char_idx]
  126. if text in self.td_token:
  127. bbox = bbox_preds[batch_idx, idx]
  128. bbox = self._bbox_decode(bbox, shape_list[batch_idx])
  129. bbox_list.append(bbox)
  130. structure_list.append(text)
  131. score_list.append(structure_probs[batch_idx, idx])
  132. structure_batch_list.append([structure_list, np.mean(score_list)])
  133. bbox_batch_list.append(np.array(bbox_list))
  134. result = {
  135. 'bbox_batch_list': bbox_batch_list,
  136. 'structure_batch_list': structure_batch_list,
  137. }
  138. return result
  139. def decode_label(self, batch):
  140. """convert text-label into text-index.
  141. """
  142. structure_idx = batch[1]
  143. gt_bbox_list = batch[2]
  144. shape_list = batch[-1]
  145. ignored_tokens = self.get_ignored_tokens()
  146. end_idx = self.dict[self.end_str]
  147. structure_batch_list = []
  148. bbox_batch_list = []
  149. batch_size = len(structure_idx)
  150. for batch_idx in range(batch_size):
  151. structure_list = []
  152. bbox_list = []
  153. for idx in range(len(structure_idx[batch_idx])):
  154. char_idx = int(structure_idx[batch_idx][idx])
  155. if idx > 0 and char_idx == end_idx:
  156. break
  157. if char_idx in ignored_tokens:
  158. continue
  159. structure_list.append(self.character[char_idx])
  160. bbox = gt_bbox_list[batch_idx][idx]
  161. if bbox.sum() != 0:
  162. bbox = self._bbox_decode(bbox, shape_list[batch_idx])
  163. bbox_list.append(bbox)
  164. structure_batch_list.append(structure_list)
  165. bbox_batch_list.append(bbox_list)
  166. result = {
  167. 'bbox_batch_list': bbox_batch_list,
  168. 'structure_batch_list': structure_batch_list,
  169. }
  170. return result
  171. def _bbox_decode(self, bbox, shape):
  172. w, h = shape[:2]
  173. bbox[0::2] *= w
  174. bbox[1::2] *= h
  175. return bbox
  176. class TableMasterLabelDecode(TableLabelDecode):
  177. """ decode the table model outputs(probs) to character str"""
  178. def __init__(self,
  179. character_dict_type='TableMaster',
  180. box_shape='pad',
  181. merge_no_span_structure=True):
  182. super(TableMasterLabelDecode, self).__init__(character_dict_type,
  183. merge_no_span_structure)
  184. self.box_shape = box_shape
  185. assert box_shape in [
  186. 'ori', 'pad'
  187. ], 'The shape used for box normalization must be ori or pad'
  188. def add_special_char(self, dict_character):
  189. """ add_special_char """
  190. self.beg_str = '<SOS>'
  191. self.end_str = '<EOS>'
  192. self.unknown_str = '<UKN>'
  193. self.pad_str = '<PAD>'
  194. dict_character = dict_character
  195. dict_character = dict_character + [
  196. self.unknown_str, self.beg_str, self.end_str, self.pad_str
  197. ]
  198. return dict_character
  199. def get_ignored_tokens(self):
  200. """ get_ignored_tokens """
  201. pad_idx = self.dict[self.pad_str]
  202. start_idx = self.dict[self.beg_str]
  203. end_idx = self.dict[self.end_str]
  204. unknown_idx = self.dict[self.unknown_str]
  205. return [start_idx, end_idx, pad_idx, unknown_idx]
  206. def _bbox_decode(self, bbox, shape):
  207. """ _bbox_decode """
  208. h, w, ratio_h, ratio_w, pad_h, pad_w = shape
  209. if self.box_shape == 'pad':
  210. h, w = pad_h, pad_w
  211. bbox[0::2] *= w
  212. bbox[1::2] *= h
  213. bbox[0::2] /= ratio_w
  214. bbox[1::2] /= ratio_h
  215. x, y, w, h = bbox
  216. x1, y1, x2, y2 = x - w // 2, y - h // 2, x + w // 2, y + h // 2
  217. bbox = np.array([x1, y1, x2, y2])
  218. return bbox
  219. class SaveTableResults(BaseTransform):
  220. """ SaveTableResults """
  221. _TABLE_RES_SUFFIX = '_bbox'
  222. _FILE_EXT = '.png'
  223. # _DEFAULT_FILE_NAME = 'table_res_out.png'
  224. def __init__(self, save_dir):
  225. super().__init__()
  226. self.save_dir = save_dir
  227. # We use pillow backend to save both numpy arrays and PIL Image objects
  228. self._writer = ImageWriter(backend='pillow')
  229. def apply(self, data):
  230. """ apply """
  231. ori_path = data[K.IM_PATH]
  232. bbox_res = data[K.BBOX_RES]
  233. file_name = os.path.basename(ori_path)
  234. file_name = self._replace_ext(file_name, self._FILE_EXT)
  235. table_res_save_path = os.path.join(self.save_dir, file_name)
  236. if len(bbox_res) > 0 and len(bbox_res[0]) == 4:
  237. vis_img = self.draw_rectangle(data[K.ORI_IM], bbox_res)
  238. else:
  239. vis_img = self.draw_bbox(data[K.ORI_IM], bbox_res)
  240. table_res_save_path = self._add_suffix(table_res_save_path,
  241. self._TABLE_RES_SUFFIX)
  242. self._write_im(table_res_save_path, vis_img)
  243. return data
  244. @classmethod
  245. def get_input_keys(cls):
  246. """ get input keys """
  247. return [K.IM_PATH, K.ORI_IM, K.BBOX_RES]
  248. @classmethod
  249. def get_output_keys(cls):
  250. """ get output keys """
  251. return []
  252. def _write_im(self, path, im):
  253. """ write image """
  254. if os.path.exists(path):
  255. logging.warning(f"{path} already exists. Overwriting it.")
  256. self._writer.write(path, im)
  257. @staticmethod
  258. def _add_suffix(path, suffix):
  259. """ _add_suffix """
  260. stem, ext = os.path.splitext(path)
  261. return stem + suffix + ext
  262. @staticmethod
  263. def _replace_ext(path, new_ext):
  264. """ _replace_ext """
  265. stem, _ = os.path.splitext(path)
  266. return stem + new_ext
  267. def draw_rectangle(self, img_path, boxes):
  268. """ draw_rectangle """
  269. boxes = np.array(boxes)
  270. img = cv2.imread(img_path)
  271. img_show = img.copy()
  272. for box in boxes.astype(int):
  273. x1, y1, x2, y2 = box
  274. cv2.rectangle(img_show, (x1, y1), (x2, y2), (255, 0, 0), 2)
  275. return img_show
  276. def draw_bbox(self, image, boxes):
  277. """ draw_bbox """
  278. for box in boxes:
  279. box = np.reshape(np.array(box), [-1, 1, 2]).astype(np.int64)
  280. image = cv2.polylines(np.array(image), [box], True, (255, 0, 0), 2)
  281. return image
  282. class PrintResult(BaseTransform):
  283. """ Print Result Transform """
  284. def apply(self, data):
  285. """ apply """
  286. logging.info("The prediction result is:")
  287. logging.info(data[K.BOXES])
  288. return data
  289. @classmethod
  290. def get_input_keys(cls):
  291. """ get input keys """
  292. return [K.BOXES]
  293. @classmethod
  294. def get_output_keys(cls):
  295. """ get output keys """
  296. return []