result_v2.py 28 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747
  1. # Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
  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. import math
  17. import re
  18. from functools import partial
  19. from pathlib import Path
  20. from typing import List
  21. import numpy as np
  22. from PIL import Image, ImageDraw, ImageFont
  23. from ....utils.fonts import PINGFANG_FONT_FILE_PATH
  24. from ...common.result import (
  25. BaseCVResult,
  26. HtmlMixin,
  27. JsonMixin,
  28. MarkdownMixin,
  29. XlsxMixin,
  30. )
  31. from .setting import BLOCK_LABEL_MAP
  32. def compile_title_pattern():
  33. # Precompiled regex pattern for matching numbering at the beginning of the title
  34. numbering_pattern = (
  35. r"(?:" + r"[1-9][0-9]*(?:\.[1-9][0-9]*)*[\.、]?|" + r"[\(\(](?:[1-9][0-9]*|["
  36. r"一二三四五六七八九十百千万亿零壹贰叁肆伍陆柒捌玖拾]+)[\)\)]|" + r"["
  37. r"一二三四五六七八九十百千万亿零壹贰叁肆伍陆柒捌玖拾]+"
  38. r"[、\.]?|" + r"(?:I|II|III|IV|V|VI|VII|VIII|IX|X)\.?" + r")"
  39. )
  40. return re.compile(r"^\s*(" + numbering_pattern + r")(\s*)(.*)$")
  41. TITLE_RE_PATTERN = compile_title_pattern()
  42. def format_title_func(block):
  43. """
  44. Normalize chapter title.
  45. Add the '#' to indicate the level of the title.
  46. If numbering exists, ensure there's exactly one space between it and the title content.
  47. If numbering does not exist, return the original title unchanged.
  48. :param title: Original chapter title string.
  49. :return: Normalized chapter title string.
  50. """
  51. title = block.content
  52. match = TITLE_RE_PATTERN.match(title)
  53. if match:
  54. numbering = match.group(1).strip()
  55. title_content = match.group(3).lstrip()
  56. # Return numbering and title content separated by one space
  57. title = numbering + " " + title_content
  58. title = title.rstrip(".")
  59. level = (
  60. title.count(
  61. ".",
  62. )
  63. + 1
  64. if "." in title
  65. else 1
  66. )
  67. return f"#{'#' * level} {title}".replace("-\n", "").replace(
  68. "\n",
  69. " ",
  70. )
  71. def format_centered_by_html(string):
  72. return (
  73. f'<div style="text-align: center;">{string}</div>'.replace(
  74. "-\n",
  75. "",
  76. ).replace("\n", " ")
  77. + "\n"
  78. )
  79. def format_text_plain_func(block):
  80. return block.content
  81. def format_image_scaled_by_html_func(block, original_image_width):
  82. img_tags = []
  83. image_path = "".join(block.image.keys())
  84. image_width = block.image[image_path].width
  85. scale = int(image_width / original_image_width * 100)
  86. img_tags.append(
  87. '<img src="{}" alt="Image" width="{}%" />'.format(
  88. image_path.replace("-\n", "").replace("\n", " "), scale
  89. ),
  90. )
  91. return "\n".join(img_tags)
  92. def format_image_plain_func(block):
  93. img_tags = []
  94. image_path = "".join(block.image.keys())
  95. img_tags.append("![]({})".format(image_path.replace("-\n", "").replace("\n", " ")))
  96. return "\n".join(img_tags)
  97. def format_chart2table_func(block):
  98. lines_list = block.content.split("\n")
  99. column_num = len(lines_list[0].split("|"))
  100. lines_list.insert(1, "|".join(["---"] * column_num))
  101. lines_list = [f"|{line}|" for line in lines_list]
  102. return "\n".join(lines_list)
  103. def simplify_table_func(table_code):
  104. return "\n" + table_code.replace("<html>", "").replace(
  105. "</html>", ""
  106. ).replace("<body>", "").replace("</body>", "")
  107. def format_first_line_func(block, templates, format_func, spliter):
  108. lines = block.content.split(spliter)
  109. for idx in range(len(lines)):
  110. line = lines[idx]
  111. if line.strip() == "":
  112. continue
  113. if line.lower() in templates:
  114. lines[idx] = format_func(line)
  115. break
  116. return spliter.join(lines)
  117. def compose_funcs(block, funcs):
  118. res = ""
  119. for func in funcs:
  120. res += func(block)
  121. return res
  122. def get_seg_flag(block: LayoutParsingBlock, prev_block: LayoutParsingBlock):
  123. seg_start_flag = True
  124. seg_end_flag = True
  125. block_box = block.bbox
  126. context_left_coordinate = block_box[0]
  127. context_right_coordinate = block_box[2]
  128. seg_start_coordinate = block.seg_start_coordinate
  129. seg_end_coordinate = block.seg_end_coordinate
  130. if prev_block is not None:
  131. prev_block_bbox = prev_block.bbox
  132. num_of_prev_lines = prev_block.num_of_lines
  133. pre_block_seg_end_coordinate = prev_block.seg_end_coordinate
  134. prev_end_space_small = (
  135. abs(prev_block_bbox[2] - pre_block_seg_end_coordinate) < 10
  136. )
  137. prev_lines_more_than_one = num_of_prev_lines > 1
  138. overlap_blocks = context_left_coordinate < prev_block_bbox[2]
  139. # update context_left_coordinate and context_right_coordinate
  140. if overlap_blocks:
  141. context_left_coordinate = min(prev_block_bbox[0], context_left_coordinate)
  142. context_right_coordinate = max(prev_block_bbox[2], context_right_coordinate)
  143. prev_end_space_small = (
  144. abs(context_right_coordinate - pre_block_seg_end_coordinate) < 10
  145. )
  146. edge_distance = 0
  147. else:
  148. edge_distance = abs(block_box[0] - prev_block_bbox[2])
  149. current_start_space_small = seg_start_coordinate - context_left_coordinate < 10
  150. if (
  151. prev_end_space_small
  152. and current_start_space_small
  153. and prev_lines_more_than_one
  154. and edge_distance < max(prev_block.width, block.width)
  155. ):
  156. seg_start_flag = False
  157. else:
  158. if seg_start_coordinate - context_left_coordinate < 10:
  159. seg_start_flag = False
  160. if context_right_coordinate - seg_end_coordinate < 10:
  161. seg_end_flag = False
  162. return seg_start_flag, seg_end_flag
  163. class LayoutParsingResultV2(BaseCVResult, HtmlMixin, XlsxMixin, MarkdownMixin):
  164. """Layout Parsing Result V2"""
  165. def __init__(self, data) -> None:
  166. """Initializes a new instance of the class with the specified data."""
  167. super().__init__(data)
  168. HtmlMixin.__init__(self)
  169. XlsxMixin.__init__(self)
  170. MarkdownMixin.__init__(self)
  171. JsonMixin.__init__(self)
  172. def _get_input_fn(self):
  173. fn = super()._get_input_fn()
  174. if (page_idx := self["page_index"]) is not None:
  175. fp = Path(fn)
  176. stem, suffix = fp.stem, fp.suffix
  177. return f"{stem}_{page_idx}{suffix}"
  178. else:
  179. return fn
  180. def _to_img(self) -> dict[str, np.ndarray]:
  181. from .utils import get_show_color
  182. res_img_dict = {}
  183. model_settings = self["model_settings"]
  184. if model_settings["use_doc_preprocessor"]:
  185. for key, value in self["doc_preprocessor_res"].img.items():
  186. res_img_dict[key] = value
  187. res_img_dict["layout_det_res"] = self["layout_det_res"].img["res"]
  188. if model_settings["use_region_detection"]:
  189. res_img_dict["region_det_res"] = self["region_det_res"].img["res"]
  190. if model_settings["use_general_ocr"] or model_settings["use_table_recognition"]:
  191. res_img_dict["overall_ocr_res"] = self["overall_ocr_res"].img["ocr_res_img"]
  192. if model_settings["use_table_recognition"] and len(self["table_res_list"]) > 0:
  193. table_cell_img = Image.fromarray(
  194. copy.deepcopy(self["doc_preprocessor_res"]["output_img"])
  195. )
  196. table_draw = ImageDraw.Draw(table_cell_img)
  197. rectangle_color = (255, 0, 0)
  198. for sno in range(len(self["table_res_list"])):
  199. table_res = self["table_res_list"][sno]
  200. cell_box_list = table_res["cell_box_list"]
  201. for box in cell_box_list:
  202. x1, y1, x2, y2 = [int(pos) for pos in box]
  203. table_draw.rectangle(
  204. [x1, y1, x2, y2], outline=rectangle_color, width=2
  205. )
  206. res_img_dict["table_cell_img"] = table_cell_img
  207. if model_settings["use_seal_recognition"] and len(self["seal_res_list"]) > 0:
  208. for sno in range(len(self["seal_res_list"])):
  209. seal_res = self["seal_res_list"][sno]
  210. seal_region_id = seal_res["seal_region_id"]
  211. sub_seal_res_dict = seal_res.img
  212. key = f"seal_res_region{seal_region_id}"
  213. res_img_dict[key] = sub_seal_res_dict["ocr_res_img"]
  214. # for layout ordering image
  215. image = Image.fromarray(self["doc_preprocessor_res"]["output_img"][:, :, ::-1])
  216. draw = ImageDraw.Draw(image, "RGBA")
  217. font_size = int(0.018 * int(image.width)) + 2
  218. font = ImageFont.truetype(PINGFANG_FONT_FILE_PATH, font_size, encoding="utf-8")
  219. parsing_result: List[LayoutParsingBlock] = self["parsing_res_list"]
  220. for block in parsing_result:
  221. bbox = block.bbox
  222. index = block.order_index
  223. label = block.label
  224. fill_color = get_show_color(label, False)
  225. draw.rectangle(bbox, fill=fill_color)
  226. if index is not None:
  227. text_position = (bbox[2] + 2, bbox[1] - font_size // 2)
  228. if int(image.width) - bbox[2] < font_size:
  229. text_position = (
  230. int(bbox[2] - font_size * 1.1),
  231. bbox[1] - font_size // 2,
  232. )
  233. draw.text(text_position, str(index), font=font, fill="red")
  234. res_img_dict["layout_order_res"] = image
  235. return res_img_dict
  236. def _to_str(self, *args, **kwargs) -> dict[str, str]:
  237. """Converts the instance's attributes to a dictionary and then to a string.
  238. Args:
  239. *args: Additional positional arguments passed to the base class method.
  240. **kwargs: Additional keyword arguments passed to the base class method.
  241. Returns:
  242. Dict[str, str]: A dictionary with the instance's attributes converted to strings.
  243. """
  244. data = {}
  245. data["input_path"] = self["input_path"]
  246. data["page_index"] = self["page_index"]
  247. model_settings = self["model_settings"]
  248. data["model_settings"] = model_settings
  249. if self["model_settings"]["use_doc_preprocessor"]:
  250. data["doc_preprocessor_res"] = self["doc_preprocessor_res"].str["res"]
  251. data["layout_det_res"] = self["layout_det_res"].str["res"]
  252. if model_settings["use_general_ocr"] or model_settings["use_table_recognition"]:
  253. data["overall_ocr_res"] = self["overall_ocr_res"].str["res"]
  254. if model_settings["use_table_recognition"] and len(self["table_res_list"]) > 0:
  255. data["table_res_list"] = []
  256. for sno in range(len(self["table_res_list"])):
  257. table_res = self["table_res_list"][sno]
  258. data["table_res_list"].append(table_res.str["res"])
  259. if model_settings["use_seal_recognition"] and len(self["seal_res_list"]) > 0:
  260. data["seal_res_list"] = []
  261. for sno in range(len(self["seal_res_list"])):
  262. seal_res = self["seal_res_list"][sno]
  263. data["seal_res_list"].append(seal_res.str["res"])
  264. if (
  265. model_settings["use_formula_recognition"]
  266. and len(self["formula_res_list"]) > 0
  267. ):
  268. data["formula_res_list"] = []
  269. for sno in range(len(self["formula_res_list"])):
  270. formula_res = self["formula_res_list"][sno]
  271. data["formula_res_list"].append(formula_res.str["res"])
  272. return JsonMixin._to_str(data, *args, **kwargs)
  273. def _to_json(self, *args, **kwargs) -> dict[str, str]:
  274. """
  275. Converts the object's data to a JSON dictionary.
  276. Args:
  277. *args: Positional arguments passed to the JsonMixin._to_json method.
  278. **kwargs: Keyword arguments passed to the JsonMixin._to_json method.
  279. Returns:
  280. Dict[str, str]: A dictionary containing the object's data in JSON format.
  281. """
  282. data = {}
  283. data["input_path"] = self["input_path"]
  284. data["page_index"] = self["page_index"]
  285. model_settings = self["model_settings"]
  286. data["model_settings"] = model_settings
  287. parsing_res_list = self["parsing_res_list"]
  288. parsing_res_list = [
  289. {
  290. "block_label": parsing_res.label,
  291. "block_content": parsing_res.content,
  292. "block_bbox": parsing_res.bbox,
  293. }
  294. for parsing_res in parsing_res_list
  295. ]
  296. data["parsing_res_list"] = parsing_res_list
  297. if self["model_settings"]["use_doc_preprocessor"]:
  298. data["doc_preprocessor_res"] = self["doc_preprocessor_res"].json["res"]
  299. data["layout_det_res"] = self["layout_det_res"].json["res"]
  300. if model_settings["use_general_ocr"] or model_settings["use_table_recognition"]:
  301. data["overall_ocr_res"] = self["overall_ocr_res"].json["res"]
  302. if model_settings["use_table_recognition"] and len(self["table_res_list"]) > 0:
  303. data["table_res_list"] = []
  304. for sno in range(len(self["table_res_list"])):
  305. table_res = self["table_res_list"][sno]
  306. data["table_res_list"].append(table_res.json["res"])
  307. if model_settings["use_seal_recognition"] and len(self["seal_res_list"]) > 0:
  308. data["seal_res_list"] = []
  309. for sno in range(len(self["seal_res_list"])):
  310. seal_res = self["seal_res_list"][sno]
  311. data["seal_res_list"].append(seal_res.json["res"])
  312. if (
  313. model_settings["use_formula_recognition"]
  314. and len(self["formula_res_list"]) > 0
  315. ):
  316. data["formula_res_list"] = []
  317. for sno in range(len(self["formula_res_list"])):
  318. formula_res = self["formula_res_list"][sno]
  319. data["formula_res_list"].append(formula_res.json["res"])
  320. return JsonMixin._to_json(data, *args, **kwargs)
  321. def _to_html(self) -> dict[str, str]:
  322. """Converts the prediction to its corresponding HTML representation.
  323. Returns:
  324. Dict[str, str]: The str type HTML representation result.
  325. """
  326. model_settings = self["model_settings"]
  327. res_html_dict = {}
  328. if model_settings["use_table_recognition"] and len(self["table_res_list"]) > 0:
  329. for sno in range(len(self["table_res_list"])):
  330. table_res = self["table_res_list"][sno]
  331. table_region_id = table_res["table_region_id"]
  332. key = f"table_{table_region_id}"
  333. res_html_dict[key] = table_res.html["pred"]
  334. return res_html_dict
  335. def _to_xlsx(self) -> dict[str, str]:
  336. """Converts the prediction HTML to an XLSX file path.
  337. Returns:
  338. Dict[str, str]: The str type XLSX representation result.
  339. """
  340. model_settings = self["model_settings"]
  341. res_xlsx_dict = {}
  342. if model_settings["use_table_recognition"] and len(self["table_res_list"]) > 0:
  343. for sno in range(len(self["table_res_list"])):
  344. table_res = self["table_res_list"][sno]
  345. table_region_id = table_res["table_region_id"]
  346. key = f"table_{table_region_id}"
  347. res_xlsx_dict[key] = table_res.xlsx["pred"]
  348. return res_xlsx_dict
  349. def _to_markdown(self, pretty_markdown=True) -> dict:
  350. """
  351. Save the parsing result to a Markdown file.
  352. Args:
  353. pretty_markdown (Optional[bool]): wheather to pretting markdown by HTML, default by True.
  354. Returns:
  355. Dict
  356. """
  357. original_image_width = self["doc_preprocessor_res"]["output_img"].shape[1]
  358. if pretty_markdown:
  359. format_text_func = lambda block: format_centered_by_html(
  360. format_text_plain_func(block)
  361. )
  362. format_image_func = lambda block: format_centered_by_html(
  363. format_image_scaled_by_html_func(
  364. block,
  365. original_image_width=original_image_width,
  366. )
  367. )
  368. format_table = lambda block: "\n" + format_text_func(block)
  369. else:
  370. format_text_func = lambda block: block.content
  371. format_image_func = format_image_plain_func
  372. format_table = lambda block: simplify_table_func("\n" + block.content)
  373. if self["model_settings"].get("use_chart_recognition", False):
  374. format_chart_func = format_chart2table_func
  375. else:
  376. format_chart_func = format_image_func
  377. if self["model_settings"].get("use_seal_recognition", False):
  378. format_seal_func = lambda block: "\n".join(
  379. [format_image_func(block), format_text_func(block)]
  380. )
  381. else:
  382. format_seal_func = format_image_func
  383. handle_funcs_dict = {
  384. "paragraph_title": format_title_func,
  385. "abstract_title": format_title_func,
  386. "reference_title": format_title_func,
  387. "content_title": format_title_func,
  388. "doc_title": lambda block: f"# {block.content}".replace(
  389. "-\n",
  390. "",
  391. ).replace("\n", " "),
  392. "table_title": format_text_func,
  393. "figure_title": format_text_func,
  394. "chart_title": format_text_func,
  395. "text": lambda block: block.content.replace("\n\n", "\n").replace(
  396. "\n", "\n\n"
  397. ),
  398. "abstract": partial(
  399. format_first_line_func,
  400. templates=["摘要", "abstract"],
  401. format_func=lambda l: f"## {l}\n",
  402. spliter=" ",
  403. ),
  404. "content": lambda block: block.content.replace("-\n", " \n").replace(
  405. "\n", " \n"
  406. ),
  407. "image": format_image_func,
  408. "chart": format_chart_func,
  409. "formula": lambda block: f"$${block.content}$$",
  410. "table": format_table,
  411. "reference": partial(
  412. format_first_line_func,
  413. templates=["参考文献", "references"],
  414. format_func=lambda l: f"## {l}",
  415. spliter="\n",
  416. ),
  417. "algorithm": lambda block: block.content.strip("\n"),
  418. "seal": format_seal_func,
  419. }
  420. markdown_content = ""
  421. last_label = None
  422. seg_start_flag = None
  423. seg_end_flag = None
  424. prev_block = None
  425. page_first_element_seg_start_flag = None
  426. page_last_element_seg_end_flag = None
  427. for block in self["parsing_res_list"]:
  428. seg_start_flag, seg_end_flag = get_seg_flag(block, prev_block)
  429. label = block.label
  430. page_first_element_seg_start_flag = (
  431. seg_start_flag
  432. if (page_first_element_seg_start_flag is None)
  433. else page_first_element_seg_start_flag
  434. )
  435. handle_func = handle_funcs_dict.get(label, None)
  436. if handle_func:
  437. prev_block = block
  438. if label == last_label == "text" and seg_start_flag == False:
  439. markdown_content += handle_func(block)
  440. else:
  441. markdown_content += (
  442. "\n\n" + handle_func(block)
  443. if markdown_content
  444. else handle_func(block)
  445. )
  446. last_label = label
  447. page_last_element_seg_end_flag = seg_end_flag
  448. markdown_info = {
  449. "markdown_texts": markdown_content,
  450. "page_continuation_flags": (
  451. page_first_element_seg_start_flag,
  452. page_last_element_seg_end_flag,
  453. ),
  454. }
  455. markdown_info["markdown_images"] = {}
  456. for img in self["imgs_in_doc"]:
  457. markdown_info["markdown_images"][img["path"]] = img["img"]
  458. return markdown_info
  459. class LayoutParsingBlock:
  460. def __init__(self, label, bbox, content="") -> None:
  461. self.label = label
  462. self.order_label = None
  463. self.bbox = list(map(int, bbox))
  464. self.content = content
  465. self.seg_start_coordinate = float("inf")
  466. self.seg_end_coordinate = float("-inf")
  467. self.width = bbox[2] - bbox[0]
  468. self.height = bbox[3] - bbox[1]
  469. self.area = self.width * self.height
  470. self.num_of_lines = 1
  471. self.image = None
  472. self.index = None
  473. self.order_index = None
  474. self.text_line_width = 1
  475. self.text_line_height = 1
  476. self.direction = self.get_bbox_direction()
  477. self.child_blocks = []
  478. self.update_direction_info()
  479. def __str__(self) -> str:
  480. return f"{self.__dict__}"
  481. def __repr__(self) -> str:
  482. _str = f"\n\n#################\nindex:\t{self.index}\nlabel:\t{self.label}\nregion_label:\t{self.order_label}\nbbox:\t{self.bbox}\ncontent:\t{self.content}\n#################"
  483. return _str
  484. def to_dict(self) -> dict:
  485. return self.__dict__
  486. def update_direction_info(self) -> None:
  487. if self.direction == "horizontal":
  488. self.secondary_direction = "vertical"
  489. self.short_side_length = self.height
  490. self.long_side_length = self.width
  491. self.start_coordinate = self.bbox[0]
  492. self.end_coordinate = self.bbox[2]
  493. self.secondary_direction_start_coordinate = self.bbox[1]
  494. self.secondary_direction_end_coordinate = self.bbox[3]
  495. else:
  496. self.secondary_direction = "horizontal"
  497. self.short_side_length = self.width
  498. self.long_side_length = self.height
  499. self.start_coordinate = self.bbox[1]
  500. self.end_coordinate = self.bbox[3]
  501. self.secondary_direction_start_coordinate = self.bbox[0]
  502. self.secondary_direction_end_coordinate = self.bbox[2]
  503. def append_child_block(self, child_block: LayoutParsingBlock) -> None:
  504. if not self.child_blocks:
  505. self.ori_bbox = self.bbox.copy()
  506. x1, y1, x2, y2 = self.bbox
  507. x1_child, y1_child, x2_child, y2_child = child_block.bbox
  508. union_bbox = (
  509. min(x1, x1_child),
  510. min(y1, y1_child),
  511. max(x2, x2_child),
  512. max(y2, y2_child),
  513. )
  514. self.bbox = union_bbox
  515. self.update_direction_info()
  516. child_blocks = [child_block]
  517. if child_block.child_blocks:
  518. child_blocks.extend(child_block.get_child_blocks())
  519. self.child_blocks.extend(child_blocks)
  520. def get_child_blocks(self) -> list:
  521. self.bbox = self.ori_bbox
  522. child_blocks = self.child_blocks.copy()
  523. self.child_blocks = []
  524. return child_blocks
  525. def get_centroid(self) -> tuple:
  526. x1, y1, x2, y2 = self.bbox
  527. centroid = ((x1 + x2) / 2, (y1 + y2) / 2)
  528. return centroid
  529. def get_bbox_direction(self, direction_ratio: float = 1.0) -> bool:
  530. """
  531. Determine if a bounding box is horizontal or vertical.
  532. Args:
  533. bbox (List[float]): Bounding box [x_min, y_min, x_max, y_max].
  534. direction_ratio (float): Ratio for determining direction. Default is 1.0.
  535. Returns:
  536. str: "horizontal" or "vertical".
  537. """
  538. return (
  539. "horizontal" if self.width * direction_ratio >= self.height else "vertical"
  540. )
  541. class LayoutParsingRegion:
  542. def __init__(
  543. self, bbox, blocks: List[LayoutParsingBlock] = [], image_shape=None
  544. ) -> None:
  545. self.bbox = bbox
  546. self.block_map = {}
  547. self.direction = "horizontal"
  548. self.calculate_bbox_metrics(image_shape)
  549. self.doc_title_block_idxes = []
  550. self.paragraph_title_block_idxes = []
  551. self.vision_block_idxes = []
  552. self.unordered_block_idxes = []
  553. self.vision_title_block_idxes = []
  554. self.normal_text_block_idxes = []
  555. self.header_block_idxes = []
  556. self.footer_block_idxes = []
  557. self.text_line_width = 20
  558. self.text_line_height = 10
  559. self.init_region_info_from_layout(blocks)
  560. self.init_direction_info()
  561. def init_region_info_from_layout(self, blocks: List[LayoutParsingBlock]):
  562. horizontal_normal_text_block_num = 0
  563. text_line_height_list = []
  564. text_line_width_list = []
  565. for idx, block in enumerate(blocks):
  566. self.block_map[idx] = block
  567. block.index = idx
  568. if block.label in BLOCK_LABEL_MAP["header_labels"]:
  569. self.header_block_idxes.append(idx)
  570. elif block.label in BLOCK_LABEL_MAP["doc_title_labels"]:
  571. self.doc_title_block_idxes.append(idx)
  572. elif block.label in BLOCK_LABEL_MAP["paragraph_title_labels"]:
  573. self.paragraph_title_block_idxes.append(idx)
  574. elif block.label in BLOCK_LABEL_MAP["vision_labels"]:
  575. self.vision_block_idxes.append(idx)
  576. elif block.label in BLOCK_LABEL_MAP["vision_title_labels"]:
  577. self.vision_title_block_idxes.append(idx)
  578. elif block.label in BLOCK_LABEL_MAP["footer_labels"]:
  579. self.footer_block_idxes.append(idx)
  580. elif block.label in BLOCK_LABEL_MAP["unordered_labels"]:
  581. self.unordered_block_idxes.append(idx)
  582. else:
  583. self.normal_text_block_idxes.append(idx)
  584. text_line_height_list.append(block.text_line_height)
  585. text_line_width_list.append(block.text_line_width)
  586. if block.direction == "horizontal":
  587. horizontal_normal_text_block_num += 1
  588. self.direction = (
  589. "horizontal"
  590. if horizontal_normal_text_block_num
  591. >= len(self.normal_text_block_idxes) * 0.5
  592. else "vertical"
  593. )
  594. self.text_line_width = (
  595. np.mean(text_line_width_list) if text_line_width_list else 20
  596. )
  597. self.text_line_height = (
  598. np.mean(text_line_height_list) if text_line_height_list else 10
  599. )
  600. def init_direction_info(self):
  601. if self.direction == "horizontal":
  602. self.direction_start_index = 0
  603. self.direction_end_index = 2
  604. self.secondary_direction_start_index = 1
  605. self.secondary_direction_end_index = 3
  606. self.secondary_direction = "vertical"
  607. else:
  608. self.direction_start_index = 1
  609. self.direction_end_index = 3
  610. self.secondary_direction_start_index = 0
  611. self.secondary_direction_end_index = 2
  612. self.secondary_direction = "horizontal"
  613. self.direction_center_coordinate = (
  614. self.bbox[self.direction_start_index] + self.bbox[self.direction_end_index]
  615. ) / 2
  616. self.secondary_direction_center_coordinate = (
  617. self.bbox[self.secondary_direction_start_index]
  618. + self.bbox[self.secondary_direction_end_index]
  619. ) / 2
  620. def calculate_bbox_metrics(self, image_shape):
  621. x1, y1, x2, y2 = self.bbox
  622. width = x2 - x1
  623. image_height, image_width = image_shape
  624. x_center, y_center = (x1 + x2) / 2, (y1 + y2) / 2
  625. self.euclidean_distance = math.sqrt(((x1) ** 2 + (y1) ** 2))
  626. self.center_euclidean_distance = math.sqrt(((x_center) ** 2 + (y_center) ** 2))
  627. self.angle_rad = math.atan2(y_center, x_center)
  628. self.weighted_distance = (
  629. y1 + width + (x1 // (image_width // 10)) * (image_width // 10) * 1.5
  630. )
  631. def sort_normal_blocks(self, blocks):
  632. if self.direction == "horizontal":
  633. blocks.sort(
  634. key=lambda x: (
  635. x.bbox[1] // self.text_line_height,
  636. x.bbox[0] // self.text_line_width,
  637. x.bbox[1] ** 2 + x.bbox[0] ** 2,
  638. ),
  639. )
  640. else:
  641. blocks.sort(
  642. key=lambda x: (
  643. -x.bbox[0] // self.text_line_width,
  644. x.bbox[1] // self.text_line_height,
  645. -(x.bbox[2] ** 2 + x.bbox[1] ** 2),
  646. ),
  647. )
  648. def sort(self):
  649. from .xycut_enhanced import xycut_enhanced
  650. return xycut_enhanced(self)