utils.py 24 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668
  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. __all__ = [
  15. "get_sub_regions_ocr_res",
  16. "get_show_color",
  17. "sorted_layout_boxes",
  18. "update_layout_order_config_block_index",
  19. ]
  20. import re
  21. from copy import deepcopy
  22. from typing import Dict, List, Optional, Tuple, Union
  23. import numpy as np
  24. from PIL import Image
  25. from ..components import convert_points_to_boxes
  26. from ..ocr.result import OCRResult
  27. from .xycut_enhanced import calculate_projection_iou
  28. def get_overlap_boxes_idx(src_boxes: np.ndarray, ref_boxes: np.ndarray) -> List:
  29. """
  30. Get the indices of source boxes that overlap with reference boxes based on a specified threshold.
  31. Args:
  32. src_boxes (np.ndarray): A 2D numpy array of source bounding boxes.
  33. ref_boxes (np.ndarray): A 2D numpy array of reference bounding boxes.
  34. Returns:
  35. match_idx_list (list): A list of indices of source boxes that overlap with reference boxes.
  36. """
  37. match_idx_list = []
  38. src_boxes_num = len(src_boxes)
  39. if src_boxes_num > 0 and len(ref_boxes) > 0:
  40. for rno in range(len(ref_boxes)):
  41. ref_box = ref_boxes[rno]
  42. x1 = np.maximum(ref_box[0], src_boxes[:, 0])
  43. y1 = np.maximum(ref_box[1], src_boxes[:, 1])
  44. x2 = np.minimum(ref_box[2], src_boxes[:, 2])
  45. y2 = np.minimum(ref_box[3], src_boxes[:, 3])
  46. pub_w = x2 - x1
  47. pub_h = y2 - y1
  48. match_idx = np.where((pub_w > 3) & (pub_h > 3))[0]
  49. match_idx_list.extend(match_idx)
  50. return match_idx_list
  51. def get_sub_regions_ocr_res(
  52. overall_ocr_res: OCRResult,
  53. object_boxes: List,
  54. flag_within: bool = True,
  55. return_match_idx: bool = False,
  56. ) -> OCRResult:
  57. """
  58. Filters OCR results to only include text boxes within specified object boxes based on a flag.
  59. Args:
  60. overall_ocr_res (OCRResult): The original OCR result containing all text boxes.
  61. object_boxes (list): A list of bounding boxes for the objects of interest.
  62. flag_within (bool): If True, only include text boxes within the object boxes. If False, exclude text boxes within the object boxes.
  63. return_match_idx (bool): If True, return the list of matching indices.
  64. Returns:
  65. OCRResult: A filtered OCR result containing only the relevant text boxes.
  66. """
  67. sub_regions_ocr_res = {}
  68. sub_regions_ocr_res["rec_polys"] = []
  69. sub_regions_ocr_res["rec_texts"] = []
  70. sub_regions_ocr_res["rec_scores"] = []
  71. sub_regions_ocr_res["rec_boxes"] = []
  72. overall_text_boxes = overall_ocr_res["rec_boxes"]
  73. match_idx_list = get_overlap_boxes_idx(overall_text_boxes, object_boxes)
  74. match_idx_list = list(set(match_idx_list))
  75. for box_no in range(len(overall_text_boxes)):
  76. if flag_within:
  77. if box_no in match_idx_list:
  78. flag_match = True
  79. else:
  80. flag_match = False
  81. else:
  82. if box_no not in match_idx_list:
  83. flag_match = True
  84. else:
  85. flag_match = False
  86. if flag_match:
  87. sub_regions_ocr_res["rec_polys"].append(
  88. overall_ocr_res["rec_polys"][box_no]
  89. )
  90. sub_regions_ocr_res["rec_texts"].append(
  91. overall_ocr_res["rec_texts"][box_no]
  92. )
  93. sub_regions_ocr_res["rec_scores"].append(
  94. overall_ocr_res["rec_scores"][box_no]
  95. )
  96. sub_regions_ocr_res["rec_boxes"].append(
  97. overall_ocr_res["rec_boxes"][box_no]
  98. )
  99. for key in ["rec_polys", "rec_scores", "rec_boxes"]:
  100. sub_regions_ocr_res[key] = np.array(sub_regions_ocr_res[key])
  101. return (
  102. (sub_regions_ocr_res, match_idx_list)
  103. if return_match_idx
  104. else sub_regions_ocr_res
  105. )
  106. def sorted_layout_boxes(res, w):
  107. """
  108. Sort text boxes in order from top to bottom, left to right
  109. Args:
  110. res: List of dictionaries containing layout information.
  111. w: Width of image.
  112. Returns:
  113. List of dictionaries containing sorted layout information.
  114. """
  115. num_boxes = len(res)
  116. if num_boxes == 1:
  117. return res
  118. # Sort on the y axis first or sort it on the x axis
  119. sorted_boxes = sorted(res, key=lambda x: (x["block_bbox"][1], x["block_bbox"][0]))
  120. _boxes = list(sorted_boxes)
  121. new_res = []
  122. res_left = []
  123. res_right = []
  124. i = 0
  125. while True:
  126. if i >= num_boxes:
  127. break
  128. # Check that the bbox is on the left
  129. elif (
  130. _boxes[i]["block_bbox"][0] < w / 4
  131. and _boxes[i]["block_bbox"][2] < 3 * w / 5
  132. ):
  133. res_left.append(_boxes[i])
  134. i += 1
  135. elif _boxes[i]["block_bbox"][0] > 2 * w / 5:
  136. res_right.append(_boxes[i])
  137. i += 1
  138. else:
  139. new_res += res_left
  140. new_res += res_right
  141. new_res.append(_boxes[i])
  142. res_left = []
  143. res_right = []
  144. i += 1
  145. res_left = sorted(res_left, key=lambda x: (x["block_bbox"][1]))
  146. res_right = sorted(res_right, key=lambda x: (x["block_bbox"][1]))
  147. if res_left:
  148. new_res += res_left
  149. if res_right:
  150. new_res += res_right
  151. return new_res
  152. def _calculate_overlap_area_div_minbox_area_ratio(
  153. bbox1: Union[list, tuple],
  154. bbox2: Union[list, tuple],
  155. ) -> float:
  156. """
  157. Calculate the ratio of the overlap area between bbox1 and bbox2
  158. to the area of the smaller bounding box.
  159. Args:
  160. bbox1 (list or tuple): Coordinates of the first bounding box [x_min, y_min, x_max, y_max].
  161. bbox2 (list or tuple): Coordinates of the second bounding box [x_min, y_min, x_max, y_max].
  162. Returns:
  163. float: The ratio of the overlap area to the area of the smaller bounding box.
  164. """
  165. bbox1 = list(map(int, bbox1))
  166. bbox2 = list(map(int, bbox2))
  167. x_left = max(bbox1[0], bbox2[0])
  168. y_top = max(bbox1[1], bbox2[1])
  169. x_right = min(bbox1[2], bbox2[2])
  170. y_bottom = min(bbox1[3], bbox2[3])
  171. if x_right <= x_left or y_bottom <= y_top:
  172. return 0.0
  173. intersection_area = (x_right - x_left) * (y_bottom - y_top)
  174. area_bbox1 = (bbox1[2] - bbox1[0]) * (bbox1[3] - bbox1[1])
  175. area_bbox2 = (bbox2[2] - bbox2[0]) * (bbox2[3] - bbox2[1])
  176. min_box_area = min(area_bbox1, area_bbox2)
  177. if min_box_area <= 0:
  178. return 0.0
  179. return intersection_area / min_box_area
  180. def group_boxes_into_lines(ocr_rec_res, block_info, line_height_iou_threshold):
  181. rec_boxes = ocr_rec_res["boxes"]
  182. rec_texts = ocr_rec_res["rec_texts"]
  183. rec_labels = ocr_rec_res["rec_labels"]
  184. spans = list(zip(rec_boxes, rec_texts, rec_labels))
  185. spans.sort(key=lambda span: span[0][1])
  186. spans = [list(span) for span in spans]
  187. lines = []
  188. line = [spans[0]]
  189. line_region_box = spans[0][0][:]
  190. block_info.seg_start_coordinate = spans[0][0][0]
  191. block_info.seg_end_coordinate = spans[-1][0][2]
  192. # merge line
  193. for span in spans[1:]:
  194. rec_bbox = span[0]
  195. if (
  196. calculate_projection_iou(line_region_box, rec_bbox, "vertical")
  197. >= line_height_iou_threshold
  198. ):
  199. line.append(span)
  200. line_region_box[1] = min(line_region_box[1], rec_bbox[1])
  201. line_region_box[3] = max(line_region_box[3], rec_bbox[3])
  202. else:
  203. lines.append(line)
  204. line = [span]
  205. line_region_box = rec_bbox[:]
  206. lines.append(line)
  207. return lines
  208. def calculate_text_orientation(
  209. bboxes: List[List[int]], orientation_ratio: float = 1.5
  210. ) -> bool:
  211. """
  212. Calculate the orientation of the text based on the bounding boxes.
  213. Args:
  214. bboxes (list): A list of bounding boxes.
  215. orientation_ratio (float): Ratio for determining orientation. Default is 1.5.
  216. Returns:
  217. str: "horizontal" or "vertical".
  218. """
  219. bboxes = np.array(bboxes)
  220. x_min = np.min(bboxes[:, 0])
  221. x_max = np.max(bboxes[:, 2])
  222. width = x_max - x_min
  223. y_min = np.min(bboxes[:, 1])
  224. y_max = np.max(bboxes[:, 3])
  225. height = y_max - y_min
  226. return "horizontal" if width * orientation_ratio >= height else "vertical"
  227. def format_line(
  228. line: List[List[Union[List[int], str]]],
  229. block_left_coordinate: int,
  230. block_right_coordinate: int,
  231. first_line_span_limit: int = 10,
  232. last_line_span_limit: int = 10,
  233. block_label: str = "text",
  234. delimiter_map: Dict = {},
  235. ) -> None:
  236. """
  237. Format a line of text spans based on layout constraints.
  238. Args:
  239. line (list): A list of spans, where each span is a list containing a bounding box and text.
  240. block_left_coordinate (int): The minimum x-coordinate of the layout bounding box.
  241. block_right_coordinate (int): The maximum x-coordinate of the layout bounding box.
  242. first_line_span_limit (int): The limit for the number of pixels before the first span that should be considered part of the first line. Default is 10.
  243. last_line_span_limit (int): The limit for the number of pixels after the last span that should be considered part of the last line. Default is 10.
  244. block_label (str): The label associated with the entire block. Default is 'text'.
  245. Returns:
  246. None: The function modifies the line in place.
  247. """
  248. first_span = line[0]
  249. last_span = line[-1]
  250. if first_span[0][0] - block_left_coordinate > first_line_span_limit:
  251. first_span[1] = "\n" + first_span[1]
  252. if block_right_coordinate - last_span[0][2] > last_line_span_limit:
  253. last_span[1] = last_span[1] + "\n"
  254. line[0] = first_span
  255. line[-1] = last_span
  256. delim = delimiter_map.get(block_label, " ")
  257. line_text = delim.join([span[1] for span in line])
  258. if block_label != "reference":
  259. line_text = remove_extra_space(line_text)
  260. if line_text.endswith("-"):
  261. line_text = line_text[:-1]
  262. return line_text
  263. def split_boxes_if_x_contained(boxes, offset=1e-5):
  264. """
  265. Check if there is any complete containment in the x-direction
  266. between the bounding boxes and split the containing box accordingly.
  267. Args:
  268. boxes (list of lists): Each element is a list containing an ndarray of length 4, a description, and a label.
  269. offset (float): A small offset value to ensure that the split boxes are not too close to the original boxes.
  270. Returns:
  271. A new list of boxes, including split boxes, with the same `rec_text` and `label` attributes.
  272. """
  273. def is_x_contained(box_a, box_b):
  274. """Check if box_a completely contains box_b in the x-direction."""
  275. return box_a[0][0] <= box_b[0][0] and box_a[0][2] >= box_b[0][2]
  276. new_boxes = []
  277. for i in range(len(boxes)):
  278. box_a = boxes[i]
  279. is_split = False
  280. for j in range(len(boxes)):
  281. if i == j:
  282. continue
  283. box_b = boxes[j]
  284. if is_x_contained(box_a, box_b):
  285. is_split = True
  286. # Split box_a based on the x-coordinates of box_b
  287. if box_a[0][0] < box_b[0][0]:
  288. w = box_b[0][0] - offset - box_a[0][0]
  289. if w > 1:
  290. new_boxes.append(
  291. [
  292. np.array(
  293. [
  294. box_a[0][0],
  295. box_a[0][1],
  296. box_b[0][0] - offset,
  297. box_a[0][3],
  298. ]
  299. ),
  300. box_a[1],
  301. box_a[2],
  302. ]
  303. )
  304. if box_a[0][2] > box_b[0][2]:
  305. w = box_a[0][2] - box_b[0][2] + offset
  306. if w > 1:
  307. box_a = [
  308. np.array(
  309. [
  310. box_b[0][2] + offset,
  311. box_a[0][1],
  312. box_a[0][2],
  313. box_a[0][3],
  314. ]
  315. ),
  316. box_a[1],
  317. box_a[2],
  318. ]
  319. if j == len(boxes) - 1 and is_split:
  320. new_boxes.append(box_a)
  321. if not is_split:
  322. new_boxes.append(box_a)
  323. return new_boxes
  324. def remove_extra_space(input_text: str) -> str:
  325. """
  326. Process the input text to handle spaces.
  327. The function removes multiple consecutive spaces between Chinese characters and ensures that
  328. only a single space is retained between Chinese and non-Chinese characters.
  329. Args:
  330. input_text (str): The text to be processed.
  331. Returns:
  332. str: The processed text with properly formatted spaces.
  333. """
  334. # Remove spaces between Chinese characters
  335. text_without_spaces = re.sub(
  336. r"(?<=[\u4e00-\u9fff])\s+(?=[\u4e00-\u9fff])", "", input_text
  337. )
  338. # Ensure single space between Chinese and non-Chinese characters
  339. text_with_single_spaces = re.sub(
  340. r"(?<=[\u4e00-\u9fff])\s+(?=[^\u4e00-\u9fff])|(?<=[^\u4e00-\u9fff])\s+(?=[\u4e00-\u9fff])",
  341. " ",
  342. text_without_spaces,
  343. )
  344. # Reduce any remaining consecutive spaces to a single space
  345. final_text = re.sub(r"\s+", " ", text_with_single_spaces).strip()
  346. return final_text
  347. def gather_imgs(original_img, layout_det_objs):
  348. imgs_in_doc = []
  349. for det_obj in layout_det_objs:
  350. if det_obj["label"] in ("image", "chart"):
  351. x_min, y_min, x_max, y_max = list(map(int, det_obj["coordinate"]))
  352. img_path = f"imgs/img_in_table_box_{x_min}_{y_min}_{x_max}_{y_max}.jpg"
  353. img = Image.fromarray(original_img[y_min:y_max, x_min:x_max, ::-1])
  354. imgs_in_doc.append(
  355. {
  356. "path": img_path,
  357. "img": img,
  358. "coordinate": (x_min, y_min, x_max, y_max),
  359. "score": det_obj["score"],
  360. }
  361. )
  362. return imgs_in_doc
  363. def _get_minbox_if_overlap_by_ratio(
  364. bbox1: Union[List[int], Tuple[int, int, int, int]],
  365. bbox2: Union[List[int], Tuple[int, int, int, int]],
  366. ratio: float,
  367. smaller: bool = True,
  368. ) -> Optional[Union[List[int], Tuple[int, int, int, int]]]:
  369. """
  370. Determine if the overlap area between two bounding boxes exceeds a given ratio
  371. and return the smaller (or larger) bounding box based on the `smaller` flag.
  372. Args:
  373. bbox1 (Union[List[int], Tuple[int, int, int, int]]): Coordinates of the first bounding box [x_min, y_min, x_max, y_max].
  374. bbox2 (Union[List[int], Tuple[int, int, int, int]]): Coordinates of the second bounding box [x_min, y_min, x_max, y_max].
  375. ratio (float): The overlap ratio threshold.
  376. smaller (bool): If True, return the smaller bounding box; otherwise, return the larger one.
  377. Returns:
  378. Optional[Union[List[int], Tuple[int, int, int, int]]]:
  379. The selected bounding box or None if the overlap ratio is not exceeded.
  380. """
  381. # Calculate the areas of both bounding boxes
  382. area1 = (bbox1[2] - bbox1[0]) * (bbox1[3] - bbox1[1])
  383. area2 = (bbox2[2] - bbox2[0]) * (bbox2[3] - bbox2[1])
  384. # Calculate the overlap ratio using a helper function
  385. overlap_ratio = _calculate_overlap_area_div_minbox_area_ratio(bbox1, bbox2)
  386. # Check if the overlap ratio exceeds the threshold
  387. if overlap_ratio > ratio:
  388. if (area1 <= area2 and smaller) or (area1 >= area2 and not smaller):
  389. return 1
  390. else:
  391. return 2
  392. return None
  393. def remove_overlap_blocks(
  394. blocks: List[Dict[str, List[int]]], threshold: float = 0.65, smaller: bool = True
  395. ) -> Tuple[List[Dict[str, List[int]]], List[Dict[str, List[int]]]]:
  396. """
  397. Remove overlapping blocks based on a specified overlap ratio threshold.
  398. Args:
  399. blocks (List[Dict[str, List[int]]]): List of block dictionaries, each containing a 'block_bbox' key.
  400. threshold (float): Ratio threshold to determine significant overlap.
  401. smaller (bool): If True, the smaller block in overlap is removed.
  402. Returns:
  403. Tuple[List[Dict[str, List[int]]], List[Dict[str, List[int]]]]:
  404. A tuple containing the updated list of blocks and a list of dropped blocks.
  405. """
  406. dropped_indexes = set()
  407. blocks = deepcopy(blocks)
  408. # Iterate over each pair of blocks to find overlaps
  409. for i, block1 in enumerate(blocks["boxes"]):
  410. for j in range(i + 1, len(blocks["boxes"])):
  411. block2 = blocks["boxes"][j]
  412. # Skip blocks that are already marked for removal
  413. if i in dropped_indexes or j in dropped_indexes:
  414. continue
  415. # Check for overlap and determine which block to remove
  416. overlap_box_index = _get_minbox_if_overlap_by_ratio(
  417. block1["coordinate"],
  418. block2["coordinate"],
  419. threshold,
  420. smaller=smaller,
  421. )
  422. if overlap_box_index is not None:
  423. # Determine which block to remove based on overlap_box_index
  424. if overlap_box_index == 1:
  425. drop_index = i
  426. else:
  427. drop_index = j
  428. dropped_indexes.add(drop_index)
  429. # Remove marked blocks from the original list
  430. for index in sorted(dropped_indexes, reverse=True):
  431. del blocks["boxes"][index]
  432. return blocks
  433. def get_bbox_intersection(bbox1, bbox2, return_format="bbox"):
  434. """
  435. Compute the intersection of two bounding boxes, supporting both 4-coordinate and 8-coordinate formats.
  436. Args:
  437. bbox1 (tuple): The first bounding box, either in 4-coordinate format (x_min, y_min, x_max, y_max)
  438. or 8-coordinate format (x1, y1, x2, y2, x3, y3, x4, y4).
  439. bbox2 (tuple): The second bounding box in the same format as bbox1.
  440. return_format (str): The format of the output intersection, either 'bbox' or 'poly'.
  441. Returns:
  442. tuple or None: The intersection bounding box in the specified format, or None if there is no intersection.
  443. """
  444. bbox1 = np.array(bbox1)
  445. bbox2 = np.array(bbox2)
  446. # Convert both bounding boxes to rectangles
  447. rect1 = bbox1 if len(bbox1.shape) == 1 else convert_points_to_boxes([bbox1])[0]
  448. rect2 = bbox2 if len(bbox2.shape) == 1 else convert_points_to_boxes([bbox2])[0]
  449. # Calculate the intersection rectangle
  450. x_min_inter = max(rect1[0], rect2[0])
  451. y_min_inter = max(rect1[1], rect2[1])
  452. x_max_inter = min(rect1[2], rect2[2])
  453. y_max_inter = min(rect1[3], rect2[3])
  454. # Check if there is an intersection
  455. if x_min_inter >= x_max_inter or y_min_inter >= y_max_inter:
  456. return None
  457. if return_format == "bbox":
  458. return np.array([x_min_inter, y_min_inter, x_max_inter, y_max_inter])
  459. elif return_format == "poly":
  460. return np.array(
  461. [
  462. [x_min_inter, y_min_inter],
  463. [x_max_inter, y_min_inter],
  464. [x_max_inter, y_max_inter],
  465. [x_min_inter, y_max_inter],
  466. ],
  467. dtype=np.int16,
  468. )
  469. else:
  470. raise ValueError("return_format must be either 'bbox' or 'poly'.")
  471. def update_layout_order_config_block_index(
  472. config: dict, block_label: str, block_idx: int
  473. ) -> None:
  474. doc_title_labels = config["doc_title_labels"]
  475. paragraph_title_labels = config["paragraph_title_labels"]
  476. vision_labels = config["vision_labels"]
  477. vision_title_labels = config["vision_title_labels"]
  478. header_labels = config["header_labels"]
  479. unordered_labels = config["unordered_labels"]
  480. footer_labels = config["footer_labels"]
  481. text_labels = config["text_labels"]
  482. text_title_labels = doc_title_labels + paragraph_title_labels
  483. config["text_title_labels"] = text_title_labels
  484. if block_label in doc_title_labels:
  485. config["doc_title_block_idxes"].append(block_idx)
  486. if block_label in paragraph_title_labels:
  487. config["paragraph_title_block_idxes"].append(block_idx)
  488. if block_label in vision_labels:
  489. config["vision_block_idxes"].append(block_idx)
  490. if block_label in vision_title_labels:
  491. config["vision_title_block_idxes"].append(block_idx)
  492. if block_label in unordered_labels:
  493. config["unordered_block_idxes"].append(block_idx)
  494. if block_label in text_title_labels:
  495. config["text_title_block_idxes"].append(block_idx)
  496. if block_label in text_labels:
  497. config["text_block_idxes"].append(block_idx)
  498. if block_label in header_labels:
  499. config["header_block_idxes"].append(block_idx)
  500. if block_label in footer_labels:
  501. config["footer_block_idxes"].append(block_idx)
  502. def update_region_box(bbox, region_box):
  503. if region_box is None:
  504. return bbox
  505. x1, y1, x2, y2 = bbox
  506. x1_region, y1_region, x2_region, y2_region = region_box
  507. x1_region = int(min(x1, x1_region))
  508. y1_region = int(min(y1, y1_region))
  509. x2_region = int(max(x2, x2_region))
  510. y2_region = int(max(y2, y2_region))
  511. region_box = [x1_region, y1_region, x2_region, y2_region]
  512. return region_box
  513. def convert_formula_res_to_ocr_format(formula_res_list: List, ocr_res: dict):
  514. for formula_res in formula_res_list:
  515. x_min, y_min, x_max, y_max = list(map(int, formula_res["dt_polys"]))
  516. poly_points = [
  517. (x_min, y_min),
  518. (x_max, y_min),
  519. (x_max, y_max),
  520. (x_min, y_max),
  521. ]
  522. ocr_res["dt_polys"].append(poly_points)
  523. ocr_res["rec_texts"].append(f"${formula_res['rec_formula']}$")
  524. ocr_res["rec_boxes"] = np.vstack(
  525. (ocr_res["rec_boxes"], [formula_res["dt_polys"]])
  526. )
  527. ocr_res["rec_labels"].append("formula")
  528. ocr_res["rec_polys"].append(poly_points)
  529. ocr_res["rec_scores"].append(1)
  530. def caculate_bbox_area(bbox):
  531. x1, y1, x2, y2 = bbox
  532. area = abs((x2 - x1) * (y2 - y1))
  533. return area
  534. def get_show_color(label: str) -> Tuple:
  535. label_colors = {
  536. # Medium Blue (from 'titles_list')
  537. "paragraph_title": (102, 102, 255, 100),
  538. "doc_title": (255, 248, 220, 100), # Cornsilk
  539. # Light Yellow (from 'tables_caption_list')
  540. "table_title": (255, 255, 102, 100),
  541. # Sky Blue (from 'imgs_caption_list')
  542. "figure_title": (102, 178, 255, 100),
  543. "chart_title": (221, 160, 221, 100), # Plum
  544. "vision_footnote": (144, 238, 144, 100), # Light Green
  545. # Deep Purple (from 'texts_list')
  546. "text": (153, 0, 76, 100),
  547. # Bright Green (from 'interequations_list')
  548. "formula": (0, 255, 0, 100),
  549. "abstract": (255, 239, 213, 100), # Papaya Whip
  550. # Medium Green (from 'lists_list' and 'indexs_list')
  551. "content": (40, 169, 92, 100),
  552. # Neutral Gray (from 'dropped_bbox_list')
  553. "seal": (158, 158, 158, 100),
  554. # Olive Yellow (from 'tables_body_list')
  555. "table": (204, 204, 0, 100),
  556. # Bright Green (from 'imgs_body_list')
  557. "image": (153, 255, 51, 100),
  558. # Bright Green (from 'imgs_body_list')
  559. "figure": (153, 255, 51, 100),
  560. "chart": (216, 191, 216, 100), # Thistle
  561. # Pale Yellow-Green (from 'tables_footnote_list')
  562. "reference": (229, 255, 204, 100),
  563. "algorithm": (255, 250, 240, 100), # Floral White
  564. }
  565. default_color = (158, 158, 158, 100)
  566. return label_colors.get(label, default_color)