result_v2.py 16 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395
  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. from __future__ import annotations
  15. import copy
  16. from pathlib import Path
  17. from PIL import Image, ImageDraw
  18. from typing import Dict
  19. import cv2
  20. import re
  21. import numpy as np
  22. from PIL import Image
  23. from PIL import ImageDraw
  24. from ...common.result import (
  25. BaseCVResult,
  26. HtmlMixin,
  27. JsonMixin,
  28. MarkdownMixin,
  29. StrMixin,
  30. XlsxMixin,
  31. )
  32. from .utils import get_layout_ordering
  33. from .utils import recursive_img_array2path
  34. from .utils import get_show_color
  35. from .utils import convert_bgr2rgb
  36. class LayoutParsingResultV2(BaseCVResult, HtmlMixin, XlsxMixin, MarkdownMixin):
  37. """Layout Parsing Result V2"""
  38. def __init__(self, data) -> None:
  39. """Initializes a new instance of the class with the specified data."""
  40. super().__init__(data)
  41. HtmlMixin.__init__(self)
  42. XlsxMixin.__init__(self)
  43. MarkdownMixin.__init__(self)
  44. JsonMixin.__init__(self)
  45. self.already_sorted = False
  46. def _get_input_fn(self):
  47. fn = super()._get_input_fn()
  48. if (page_idx := self["page_index"]) is not None:
  49. fp = Path(fn)
  50. stem, suffix = fp.stem, fp.suffix
  51. return f"{stem}_{page_idx}{suffix}"
  52. else:
  53. return fn
  54. def _to_img(self) -> dict[str, np.ndarray]:
  55. res_img_dict = {}
  56. model_settings = self["model_settings"]
  57. page_index = self["page_index"]
  58. if model_settings["use_doc_preprocessor"]:
  59. for key, value in self["doc_preprocessor_res"].img.items():
  60. res_img_dict[key] = value
  61. res_img_dict["layout_det_res"] = convert_bgr2rgb(
  62. self["layout_det_res"].img["res"]
  63. )
  64. if model_settings["use_general_ocr"] or model_settings["use_table_recognition"]:
  65. res_img_dict["overall_ocr_res"] = convert_bgr2rgb(
  66. self["overall_ocr_res"].img["ocr_res_img"]
  67. )
  68. if model_settings["use_table_recognition"] and len(self["table_res_list"]) > 0:
  69. table_cell_img = Image.fromarray(
  70. copy.deepcopy(self["doc_preprocessor_res"]["output_img"])
  71. )
  72. table_draw = ImageDraw.Draw(table_cell_img)
  73. rectangle_color = (255, 0, 0)
  74. for sno in range(len(self["table_res_list"])):
  75. table_res = self["table_res_list"][sno]
  76. cell_box_list = table_res["cell_box_list"]
  77. for box in cell_box_list:
  78. x1, y1, x2, y2 = [int(pos) for pos in box]
  79. table_draw.rectangle(
  80. [x1, y1, x2, y2], outline=rectangle_color, width=2
  81. )
  82. res_img_dict["table_cell_img"] = table_cell_img
  83. if model_settings["use_seal_recognition"] and len(self["seal_res_list"]) > 0:
  84. for sno in range(len(self["seal_res_list"])):
  85. seal_res = self["seal_res_list"][sno]
  86. seal_region_id = seal_res["seal_region_id"]
  87. sub_seal_res_dict = seal_res.img
  88. key = f"seal_res_region{seal_region_id}"
  89. res_img_dict[key] = sub_seal_res_dict["ocr_res_img"]
  90. # for layout ordering image
  91. image = Image.fromarray(self["doc_preprocessor_res"]["output_img"])
  92. draw = ImageDraw.Draw(image, "RGBA")
  93. parsing_result = self["parsing_res_list"]
  94. if self.already_sorted == False:
  95. parsing_result = get_layout_ordering(
  96. parsing_result,
  97. no_mask_labels=[
  98. "text",
  99. "formula",
  100. "algorithm",
  101. "reference",
  102. "content",
  103. "abstract",
  104. ],
  105. already_sorted=self.already_sorted,
  106. )
  107. for block in parsing_result:
  108. bbox = block["block_bbox"]
  109. index = block.get("index", None)
  110. label = block["sub_label"]
  111. fill_color = get_show_color(label)
  112. draw.rectangle(bbox, fill=fill_color)
  113. if index is not None:
  114. text_position = (bbox[2] + 2, bbox[1] - 10)
  115. draw.text(text_position, str(index), fill="red")
  116. self.already_sorted = True
  117. res_img_dict["layout_order_res"] = image
  118. return res_img_dict
  119. def _to_str(self, *args, **kwargs) -> dict[str, str]:
  120. """Converts the instance's attributes to a dictionary and then to a string.
  121. Args:
  122. *args: Additional positional arguments passed to the base class method.
  123. **kwargs: Additional keyword arguments passed to the base class method.
  124. Returns:
  125. Dict[str, str]: A dictionary with the instance's attributes converted to strings.
  126. """
  127. data = {}
  128. data["input_path"] = self["input_path"]
  129. data["page_index"] = self["page_index"]
  130. model_settings = self["model_settings"]
  131. data["model_settings"] = model_settings
  132. if self["model_settings"]["use_doc_preprocessor"]:
  133. data["doc_preprocessor_res"] = self["doc_preprocessor_res"].str["res"]
  134. data["layout_det_res"] = self["layout_det_res"].str["res"]
  135. if model_settings["use_general_ocr"] or model_settings["use_table_recognition"]:
  136. data["overall_ocr_res"] = self["overall_ocr_res"].str["res"]
  137. if model_settings["use_table_recognition"] and len(self["table_res_list"]) > 0:
  138. data["table_res_list"] = []
  139. for sno in range(len(self["table_res_list"])):
  140. table_res = self["table_res_list"][sno]
  141. data["table_res_list"].append(table_res.str["res"])
  142. if model_settings["use_seal_recognition"] and len(self["seal_res_list"]) > 0:
  143. data["seal_res_list"] = []
  144. for sno in range(len(self["seal_res_list"])):
  145. seal_res = self["seal_res_list"][sno]
  146. data["seal_res_list"].append(seal_res.str["res"])
  147. if (
  148. model_settings["use_formula_recognition"]
  149. and len(self["formula_res_list"]) > 0
  150. ):
  151. data["formula_res_list"] = []
  152. for sno in range(len(self["formula_res_list"])):
  153. formula_res = self["formula_res_list"][sno]
  154. data["formula_res_list"].append(formula_res.str["res"])
  155. return JsonMixin._to_str(data, *args, **kwargs)
  156. def _to_json(self, *args, **kwargs) -> dict[str, str]:
  157. """
  158. Converts the object's data to a JSON dictionary.
  159. Args:
  160. *args: Positional arguments passed to the JsonMixin._to_json method.
  161. **kwargs: Keyword arguments passed to the JsonMixin._to_json method.
  162. Returns:
  163. Dict[str, str]: A dictionary containing the object's data in JSON format.
  164. """
  165. data = {}
  166. data["input_path"] = self["input_path"]
  167. data["page_index"] = self["page_index"]
  168. model_settings = self["model_settings"]
  169. data["model_settings"] = model_settings
  170. parsing_res_list = self["parsing_res_list"]
  171. parsing_res_list = [
  172. {
  173. "block_label": parsing_res["block_label"],
  174. "block_content": parsing_res["block_content"],
  175. "block_bbox": parsing_res["block_bbox"],
  176. }
  177. for parsing_res in parsing_res_list
  178. ]
  179. data["parsing_res_list"] = parsing_res_list
  180. if self["model_settings"]["use_doc_preprocessor"]:
  181. data["doc_preprocessor_res"] = self["doc_preprocessor_res"].json["res"]
  182. data["layout_det_res"] = self["layout_det_res"].json["res"]
  183. if model_settings["use_general_ocr"] or model_settings["use_table_recognition"]:
  184. data["overall_ocr_res"] = self["overall_ocr_res"].json["res"]
  185. if model_settings["use_table_recognition"] and len(self["table_res_list"]) > 0:
  186. data["table_res_list"] = []
  187. for sno in range(len(self["table_res_list"])):
  188. table_res = self["table_res_list"][sno]
  189. data["table_res_list"].append(table_res.json["res"])
  190. if model_settings["use_seal_recognition"] and len(self["seal_res_list"]) > 0:
  191. data["seal_res_list"] = []
  192. for sno in range(len(self["seal_res_list"])):
  193. seal_res = self["seal_res_list"][sno]
  194. data["seal_res_list"].append(seal_res.json["res"])
  195. if (
  196. model_settings["use_formula_recognition"]
  197. and len(self["formula_res_list"]) > 0
  198. ):
  199. data["formula_res_list"] = []
  200. for sno in range(len(self["formula_res_list"])):
  201. formula_res = self["formula_res_list"][sno]
  202. data["formula_res_list"].append(formula_res.json["res"])
  203. return JsonMixin._to_json(data, *args, **kwargs)
  204. def _to_html(self) -> dict[str, str]:
  205. """Converts the prediction to its corresponding HTML representation.
  206. Returns:
  207. Dict[str, str]: The str type HTML representation result.
  208. """
  209. model_settings = self["model_settings"]
  210. res_html_dict = {}
  211. if model_settings["use_table_recognition"] and len(self["table_res_list"]) > 0:
  212. for sno in range(len(self["table_res_list"])):
  213. table_res = self["table_res_list"][sno]
  214. table_region_id = table_res["table_region_id"]
  215. key = f"table_{table_region_id}"
  216. res_html_dict[key] = table_res.html["pred"]
  217. return res_html_dict
  218. def _to_xlsx(self) -> dict[str, str]:
  219. """Converts the prediction HTML to an XLSX file path.
  220. Returns:
  221. Dict[str, str]: The str type XLSX representation result.
  222. """
  223. model_settings = self["model_settings"]
  224. res_xlsx_dict = {}
  225. if model_settings["use_table_recognition"] and len(self["table_res_list"]) > 0:
  226. for sno in range(len(self["table_res_list"])):
  227. table_res = self["table_res_list"][sno]
  228. table_region_id = table_res["table_region_id"]
  229. key = f"table_{table_region_id}"
  230. res_xlsx_dict[key] = table_res.xlsx["pred"]
  231. return res_xlsx_dict
  232. def _to_markdown(self) -> dict:
  233. """
  234. Save the parsing result to a Markdown file.
  235. Returns:
  236. Dict
  237. """
  238. parsing_result = self["parsing_res_list"]
  239. if self.already_sorted == False:
  240. parsing_result = get_layout_ordering(
  241. parsing_result,
  242. no_mask_labels=[
  243. "text",
  244. "formula",
  245. "algorithm",
  246. "reference",
  247. "content",
  248. "abstract",
  249. ],
  250. already_sorted=self.already_sorted,
  251. )
  252. self.already_sorted == True
  253. recursive_img_array2path(self["parsing_res_list"], labels=["block_image"])
  254. def _format_data(obj):
  255. def format_title(content_value):
  256. content_value = content_value.rstrip(".")
  257. level = (
  258. content_value.count(
  259. ".",
  260. )
  261. + 1
  262. if "." in content_value
  263. else 1
  264. )
  265. return f"{'#' * level} {content_value}".replace("-\n", "").replace(
  266. "\n",
  267. " ",
  268. )
  269. def format_centered_text(key):
  270. return (
  271. f'<div style="text-align: center;">{block[key]}</div>'.replace(
  272. "-\n",
  273. "",
  274. ).replace("\n", " ")
  275. + "\n"
  276. )
  277. def format_image(label):
  278. img_tags = []
  279. image_path = "".join(block[label].keys())
  280. img_tags.append(
  281. '<div style="text-align: center;"><img src="{}" alt="Image" /></div>'.format(
  282. image_path.replace("-\n", "").replace("\n", " "),
  283. ),
  284. )
  285. return "\n".join(img_tags)
  286. def format_reference():
  287. pattern = r"\s*\[\s*\d+\s*\]\s*"
  288. res = re.sub(
  289. pattern,
  290. lambda match: "\n" + match.group(),
  291. block["reference"].replace("\n", ""),
  292. )
  293. return "\n" + res
  294. def format_table():
  295. return "\n" + block["block_content"]
  296. handlers = {
  297. "paragraph_title": lambda: format_title(block["block_content"]),
  298. "doc_title": lambda: f"# {block['block_content']}".replace(
  299. "-\n",
  300. "",
  301. ).replace("\n", " "),
  302. "table_title": lambda: format_centered_text("block_content"),
  303. "figure_title": lambda: format_centered_text("block_content"),
  304. "chart_title": lambda: format_centered_text("block_content"),
  305. "text": lambda: block["block_content"]
  306. .replace("-\n", " ")
  307. .replace("\n", " "),
  308. # 'number': lambda: str(block['number']),
  309. "abstract": lambda: block["block_content"]
  310. .replace("-\n", " ")
  311. .replace("\n", " "),
  312. "content": lambda: block["block_content"]
  313. .replace("-\n", " ")
  314. .replace("\n", " "),
  315. "image": lambda: format_image("block_image"),
  316. "chart": lambda: format_image("block_image"),
  317. "formula": lambda: f"$${block['block_content']}$$",
  318. "table": format_table,
  319. # "reference": format_reference,
  320. "reference": lambda: block["block_content"],
  321. "algorithm": lambda: block["block_content"].strip("\n"),
  322. "seal": lambda: format_image("block_content"),
  323. }
  324. parsing_res_list = obj["parsing_res_list"]
  325. markdown_content = ""
  326. last_label = None
  327. seg_start_flag = None
  328. seg_end_flag = None
  329. for block in sorted(
  330. parsing_res_list,
  331. key=lambda x: x.get("sub_index", 999),
  332. ):
  333. label = block.get("block_label")
  334. seg_start_flag = block.get("seg_start_flag")
  335. handler = handlers.get(label)
  336. if handler:
  337. if (
  338. label == last_label == "text"
  339. and seg_start_flag == seg_end_flag == False
  340. ):
  341. markdown_content += " " + handler()
  342. else:
  343. markdown_content += "\n\n" + handler()
  344. last_label = label
  345. seg_end_flag = block.get("seg_end_flag")
  346. return markdown_content
  347. markdown_info = dict()
  348. markdown_info["markdown_texts"] = _format_data(self)
  349. markdown_info["markdown_images"] = dict()
  350. for block in self["parsing_res_list"]:
  351. if block["block_label"] in ["image", "chart"]:
  352. image_path, image_value = next(iter(block["block_image"].items()))
  353. markdown_info["markdown_images"][image_path] = image_value
  354. return markdown_info