result_v2.py 29 KB

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