utils.py 67 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487148814891490149114921493149414951496149714981499150015011502150315041505150615071508150915101511151215131514151515161517151815191520152115221523152415251526152715281529153015311532153315341535153615371538153915401541154215431544154515461547154815491550155115521553155415551556155715581559156015611562156315641565156615671568156915701571157215731574157515761577157815791580158115821583158415851586158715881589159015911592159315941595159615971598159916001601160216031604160516061607160816091610161116121613161416151616161716181619162016211622162316241625162616271628162916301631163216331634163516361637163816391640164116421643164416451646164716481649165016511652165316541655165616571658165916601661166216631664166516661667166816691670167116721673167416751676167716781679168016811682168316841685168616871688168916901691169216931694169516961697169816991700170117021703170417051706170717081709171017111712171317141715171617171718171917201721172217231724172517261727172817291730173117321733173417351736173717381739174017411742174317441745174617471748174917501751175217531754175517561757175817591760176117621763176417651766176717681769177017711772177317741775177617771778177917801781178217831784178517861787178817891790179117921793179417951796179717981799180018011802180318041805180618071808180918101811181218131814181518161817181818191820182118221823182418251826182718281829183018311832183318341835183618371838
  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. __all__ = [
  15. "get_sub_regions_ocr_res",
  16. "get_layout_ordering",
  17. "recursive_img_array2path",
  18. "get_show_color",
  19. ]
  20. import numpy as np
  21. import copy
  22. import cv2
  23. import uuid
  24. from pathlib import Path
  25. from typing import List
  26. from ..ocr.result import OCRResult
  27. from ...models_new.object_detection.result import DetResult
  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. list: A list of indices of source boxes that overlap with any reference box.
  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, object_boxes: List, flag_within: bool = True
  53. ) -> OCRResult:
  54. """
  55. Filters OCR results to only include text boxes within specified object boxes based on a flag.
  56. Args:
  57. overall_ocr_res (OCRResult): The original OCR result containing all text boxes.
  58. object_boxes (list): A list of bounding boxes for the objects of interest.
  59. flag_within (bool): If True, only include text boxes within the object boxes. If False, exclude text boxes within the object boxes.
  60. Returns:
  61. OCRResult: A filtered OCR result containing only the relevant text boxes.
  62. """
  63. sub_regions_ocr_res = {}
  64. sub_regions_ocr_res["rec_polys"] = []
  65. sub_regions_ocr_res["rec_texts"] = []
  66. sub_regions_ocr_res["rec_scores"] = []
  67. sub_regions_ocr_res["rec_boxes"] = []
  68. overall_text_boxes = overall_ocr_res["rec_boxes"]
  69. match_idx_list = get_overlap_boxes_idx(overall_text_boxes, object_boxes)
  70. match_idx_list = list(set(match_idx_list))
  71. for box_no in range(len(overall_text_boxes)):
  72. if flag_within:
  73. if box_no in match_idx_list:
  74. flag_match = True
  75. else:
  76. flag_match = False
  77. else:
  78. if box_no not in match_idx_list:
  79. flag_match = True
  80. else:
  81. flag_match = False
  82. if flag_match:
  83. sub_regions_ocr_res["rec_polys"].append(
  84. overall_ocr_res["rec_polys"][box_no]
  85. )
  86. sub_regions_ocr_res["rec_texts"].append(
  87. overall_ocr_res["rec_texts"][box_no]
  88. )
  89. sub_regions_ocr_res["rec_scores"].append(
  90. overall_ocr_res["rec_scores"][box_no]
  91. )
  92. sub_regions_ocr_res["rec_boxes"].append(
  93. overall_ocr_res["rec_boxes"][box_no]
  94. )
  95. return sub_regions_ocr_res
  96. def _calculate_iou(box1, box2):
  97. """
  98. Calculate Intersection over Union (IoU) between two bounding boxes.
  99. Args:
  100. box1, box2: Lists or tuples representing bounding boxes [x_min, y_min, x_max, y_max].
  101. Returns:
  102. float: The IoU value.
  103. """
  104. box1 = list(map(int, box1))
  105. box2 = list(map(int, box2))
  106. x1_min, y1_min, x1_max, y1_max = box1
  107. x2_min, y2_min, x2_max, y2_max = box2
  108. inter_x_min = max(x1_min, x2_min)
  109. inter_y_min = max(y1_min, y2_min)
  110. inter_x_max = min(x1_max, x2_max)
  111. inter_y_max = min(y1_max, y2_max)
  112. if inter_x_max <= inter_x_min or inter_y_max <= inter_y_min:
  113. return 0.0
  114. inter_area = (inter_x_max - inter_x_min) * (inter_y_max - inter_y_min)
  115. box1_area = (x1_max - x1_min) * (y1_max - y1_min)
  116. box2_area = (x2_max - x2_min) * (y2_max - y2_min)
  117. min_area = min(box1_area, box2_area)
  118. if min_area <= 0:
  119. return 0.0
  120. iou = inter_area / min_area
  121. return iou
  122. def _whether_y_overlap_exceeds_threshold(bbox1, bbox2, overlap_ratio_threshold=0.6):
  123. """
  124. Determines whether the vertical overlap between two bounding boxes exceeds a given threshold.
  125. Args:
  126. bbox1 (tuple): The first bounding box defined as (left, top, right, bottom).
  127. bbox2 (tuple): The second bounding box defined as (left, top, right, bottom).
  128. overlap_ratio_threshold (float): The threshold ratio to determine if the overlap is significant.
  129. Defaults to 0.6.
  130. Returns:
  131. bool: True if the vertical overlap divided by the minimum height of the two bounding boxes
  132. exceeds the overlap_ratio_threshold, otherwise False.
  133. """
  134. _, y1_0, _, y1_1 = bbox1
  135. _, y2_0, _, y2_1 = bbox2
  136. overlap = max(0, min(y1_1, y2_1) - max(y1_0, y2_0))
  137. min_height = min(y1_1 - y1_0, y2_1 - y2_0)
  138. return (overlap / min_height) > overlap_ratio_threshold
  139. def _sort_box_by_y_projection(layout_bbox, ocr_res, line_height_iou_threshold=0.7):
  140. """
  141. Sorts OCR results based on their spatial arrangement, grouping them into lines and blocks.
  142. Args:
  143. layout_bbox (tuple): A tuple representing the layout bounding box, defined as (left, top, right, bottom).
  144. ocr_res (dict): A dictionary containing OCR results with the following keys:
  145. - "boxes": A list of bounding boxes, each defined as [left, top, right, bottom].
  146. - "rec_texts": A corresponding list of recognized text strings for each box.
  147. line_height_iou_threshold (float): The threshold for determining whether two boxes belong to
  148. the same line based on their vertical overlap. Defaults to 0.7.
  149. Returns:
  150. dict: A dictionary with the same structure as `ocr_res`, but with boxes and texts sorted
  151. and grouped into lines and blocks.
  152. """
  153. assert (
  154. ocr_res["boxes"] and ocr_res["rec_texts"]
  155. ), "OCR results must contain 'boxes' and 'rec_texts'"
  156. boxes = ocr_res["boxes"]
  157. rec_texts = ocr_res["rec_texts"]
  158. x_min, _, x_max, _ = layout_bbox
  159. spans = list(zip(boxes, rec_texts))
  160. spans.sort(key=lambda span: span[0][1])
  161. spans = [list(span) for span in spans]
  162. lines = []
  163. current_line = [spans[0]]
  164. current_y0, current_y1 = spans[0][0][1], spans[0][0][3]
  165. for span in spans[1:]:
  166. y0, y1 = span[0][1], span[0][3]
  167. if _whether_y_overlap_exceeds_threshold(
  168. (0, current_y0, 0, current_y1),
  169. (0, y0, 0, y1),
  170. line_height_iou_threshold,
  171. ):
  172. current_line.append(span)
  173. current_y0 = min(current_y0, y0)
  174. current_y1 = max(current_y1, y1)
  175. else:
  176. lines.append(current_line)
  177. current_line = [span]
  178. current_y0, current_y1 = y0, y1
  179. if current_line:
  180. lines.append(current_line)
  181. for line in lines:
  182. line.sort(key=lambda span: span[0][0])
  183. first_span = line[0]
  184. end_span = line[-1]
  185. if first_span[0][0] - x_min > 15:
  186. first_span[1] = "\n" + first_span[1]
  187. if x_max - end_span[0][2] > 15:
  188. end_span[1] = end_span[1] + "\n"
  189. # Flatten lines back into a single list for boxes and texts
  190. ocr_res["boxes"] = [span[0] for line in lines for span in line]
  191. ocr_res["rec_texts"] = [span[1] + " " for line in lines for span in line]
  192. return ocr_res
  193. def get_structure_res(
  194. overall_ocr_res: OCRResult,
  195. layout_det_res: DetResult,
  196. table_res_list,
  197. ) -> OCRResult:
  198. """
  199. Extract structured information from OCR and layout detection results.
  200. Args:
  201. overall_ocr_res (OCRResult): An object containing the overall OCR results, including detected text boxes and recognized text. The structure is expected to have:
  202. - "input_img": The image on which OCR was performed.
  203. - "dt_boxes": A list of detected text box coordinates.
  204. - "rec_texts": A list of recognized text corresponding to the detected boxes.
  205. layout_det_res (DetResult): An object containing the layout detection results, including detected layout boxes and their labels. The structure is expected to have:
  206. - "boxes": A list of dictionaries with keys "coordinate" for box coordinates and "label" for the type of content.
  207. table_res_list (list): A list of table detection results, where each item is a dictionary containing:
  208. - "layout_bbox": The bounding box of the table layout.
  209. - "pred_html": The predicted HTML representation of the table.
  210. Returns:
  211. list: A list of structured boxes where each item is a dictionary containing:
  212. - "label": The label of the content (e.g., 'table', 'chart', 'image').
  213. - The label as a key with either table HTML or image data and text.
  214. - "layout_bbox": The coordinates of the layout box.
  215. """
  216. structure_boxes = []
  217. input_img = overall_ocr_res["doc_preprocessor_res"]["output_img"]
  218. for box_info in layout_det_res["boxes"]:
  219. layout_bbox = box_info["coordinate"]
  220. label = box_info["label"]
  221. rec_res = {"boxes": [], "rec_texts": [], "flag": False}
  222. seg_start_flag = True
  223. seg_end_flag = True
  224. if label == "table":
  225. for i, table_res in enumerate(table_res_list):
  226. if (
  227. _calculate_iou(
  228. layout_bbox, table_res["table_ocr_pred"]["rec_boxes"][0]
  229. )
  230. > 0.5
  231. ):
  232. structure_boxes.append(
  233. {
  234. "label": label,
  235. f"{label}": table_res["pred_html"],
  236. "layout_bbox": layout_bbox,
  237. "seg_start_flag": seg_start_flag,
  238. "seg_end_flag": seg_end_flag,
  239. },
  240. )
  241. del table_res_list[i]
  242. break
  243. else:
  244. overall_text_boxes = overall_ocr_res["rec_boxes"]
  245. for box_no in range(len(overall_text_boxes)):
  246. if _calculate_iou(layout_bbox, overall_text_boxes[box_no]) > 0.5:
  247. rec_res["boxes"].append(overall_text_boxes[box_no])
  248. rec_res["rec_texts"].append(
  249. overall_ocr_res["rec_texts"][box_no],
  250. )
  251. rec_res["flag"] = True
  252. if rec_res["flag"]:
  253. rec_res = _sort_box_by_y_projection(layout_bbox, rec_res, 0.7)
  254. rec_res_first_bbox = rec_res["boxes"][0]
  255. rec_res_end_bbox = rec_res["boxes"][-1]
  256. if rec_res_first_bbox[0] - layout_bbox[0] < 20:
  257. seg_start_flag = False
  258. if layout_bbox[2] - rec_res_end_bbox[2] < 20:
  259. seg_end_flag = False
  260. if label == "formula":
  261. rec_res["rec_texts"] = [
  262. rec_res_text.replace("$", "")
  263. for rec_res_text in rec_res["rec_texts"]
  264. ]
  265. if label in ["chart", "image"]:
  266. structure_boxes.append(
  267. {
  268. "label": label,
  269. f"{label}": {
  270. "img": input_img[
  271. int(layout_bbox[1]) : int(layout_bbox[3]),
  272. int(layout_bbox[0]) : int(layout_bbox[2]),
  273. ],
  274. },
  275. "layout_bbox": layout_bbox,
  276. "seg_start_flag": seg_start_flag,
  277. "seg_end_flag": seg_end_flag,
  278. },
  279. )
  280. else:
  281. structure_boxes.append(
  282. {
  283. "label": label,
  284. f"{label}": "".join(rec_res["rec_texts"]),
  285. "layout_bbox": layout_bbox,
  286. "seg_start_flag": seg_start_flag,
  287. "seg_end_flag": seg_end_flag,
  288. },
  289. )
  290. return structure_boxes
  291. def _projection_by_bboxes(boxes: np.ndarray, axis: int) -> np.ndarray:
  292. """
  293. Generate a 1D projection histogram from bounding boxes along a specified axis.
  294. Args:
  295. boxes: A (N, 4) array of bounding boxes defined by [x_min, y_min, x_max, y_max].
  296. axis: Axis for projection; 0 for horizontal (x-axis), 1 for vertical (y-axis).
  297. Returns:
  298. A 1D numpy array representing the projection histogram based on bounding box intervals.
  299. """
  300. assert axis in [0, 1]
  301. max_length = np.max(boxes[:, axis::2])
  302. projection = np.zeros(max_length, dtype=int)
  303. # Increment projection histogram over the interval defined by each bounding box
  304. for start, end in boxes[:, axis::2]:
  305. projection[start:end] += 1
  306. return projection
  307. def _split_projection_profile(arr_values: np.ndarray, min_value: float, min_gap: float):
  308. """
  309. Split the projection profile into segments based on specified thresholds.
  310. Args:
  311. arr_values: 1D array representing the projection profile.
  312. min_value: Minimum value threshold to consider a profile segment significant.
  313. min_gap: Minimum gap width to consider a separation between segments.
  314. Returns:
  315. A tuple of start and end indices for each segment that meets the criteria.
  316. """
  317. # Identify indices where the projection exceeds the minimum value
  318. significant_indices = np.where(arr_values > min_value)[0]
  319. if not len(significant_indices):
  320. return
  321. # Calculate gaps between significant indices
  322. index_diffs = significant_indices[1:] - significant_indices[:-1]
  323. gap_indices = np.where(index_diffs > min_gap)[0]
  324. # Determine start and end indices of segments
  325. segment_starts = np.insert(
  326. significant_indices[gap_indices + 1],
  327. 0,
  328. significant_indices[0],
  329. )
  330. segment_ends = np.append(
  331. significant_indices[gap_indices],
  332. significant_indices[-1] + 1,
  333. )
  334. return segment_starts, segment_ends
  335. def _recursive_yx_cut(boxes: np.ndarray, indices: List[int], res: List[int], min_gap=1):
  336. """
  337. Recursively project and segment bounding boxes, starting with Y-axis and followed by X-axis.
  338. Args:
  339. boxes: A (N, 4) array representing bounding boxes.
  340. indices: List of indices indicating the original position of boxes.
  341. res: List to store indices of the final segmented bounding boxes.
  342. """
  343. assert len(boxes) == len(indices)
  344. # Sort by y_min for Y-axis projection
  345. y_sorted_indices = boxes[:, 1].argsort()
  346. y_sorted_boxes = boxes[y_sorted_indices]
  347. y_sorted_indices = np.array(indices)[y_sorted_indices]
  348. # Perform Y-axis projection
  349. y_projection = _projection_by_bboxes(boxes=y_sorted_boxes, axis=1)
  350. y_intervals = _split_projection_profile(y_projection, 0, 1)
  351. if not y_intervals:
  352. return
  353. # Process each segment defined by Y-axis projection
  354. for y_start, y_end in zip(*y_intervals):
  355. # Select boxes within the current y interval
  356. y_interval_indices = (y_start <= y_sorted_boxes[:, 1]) & (
  357. y_sorted_boxes[:, 1] < y_end
  358. )
  359. y_boxes_chunk = y_sorted_boxes[y_interval_indices]
  360. y_indices_chunk = y_sorted_indices[y_interval_indices]
  361. # Sort by x_min for X-axis projection
  362. x_sorted_indices = y_boxes_chunk[:, 0].argsort()
  363. x_sorted_boxes_chunk = y_boxes_chunk[x_sorted_indices]
  364. x_sorted_indices_chunk = y_indices_chunk[x_sorted_indices]
  365. # Perform X-axis projection
  366. x_projection = _projection_by_bboxes(boxes=x_sorted_boxes_chunk, axis=0)
  367. x_intervals = _split_projection_profile(x_projection, 0, min_gap)
  368. if not x_intervals:
  369. continue
  370. # If X-axis cannot be further segmented, add current indices to results
  371. if len(x_intervals[0]) == 1:
  372. res.extend(x_sorted_indices_chunk)
  373. continue
  374. # Recursively process each segment defined by X-axis projection
  375. for x_start, x_end in zip(*x_intervals):
  376. x_interval_indices = (x_start <= x_sorted_boxes_chunk[:, 0]) & (
  377. x_sorted_boxes_chunk[:, 0] < x_end
  378. )
  379. _recursive_yx_cut(
  380. x_sorted_boxes_chunk[x_interval_indices],
  381. x_sorted_indices_chunk[x_interval_indices],
  382. res,
  383. )
  384. def _recursive_xy_cut(boxes: np.ndarray, indices: List[int], res: List[int], min_gap=1):
  385. """
  386. Recursively performs X-axis projection followed by Y-axis projection to segment bounding boxes.
  387. Args:
  388. boxes: A (N, 4) array representing bounding boxes with [x_min, y_min, x_max, y_max].
  389. indices: A list of indices representing the position of boxes in the original data.
  390. res: A list to store indices of bounding boxes that meet the criteria.
  391. """
  392. # Ensure boxes and indices have the same length
  393. assert len(boxes) == len(indices)
  394. # Sort by x_min to prepare for X-axis projection
  395. x_sorted_indices = boxes[:, 0].argsort()
  396. x_sorted_boxes = boxes[x_sorted_indices]
  397. x_sorted_indices = np.array(indices)[x_sorted_indices]
  398. # Perform X-axis projection
  399. x_projection = _projection_by_bboxes(boxes=x_sorted_boxes, axis=0)
  400. x_intervals = _split_projection_profile(x_projection, 0, 1)
  401. if not x_intervals:
  402. return
  403. # Process each segment defined by X-axis projection
  404. for x_start, x_end in zip(*x_intervals):
  405. # Select boxes within the current x interval
  406. x_interval_indices = (x_start <= x_sorted_boxes[:, 0]) & (
  407. x_sorted_boxes[:, 0] < x_end
  408. )
  409. x_boxes_chunk = x_sorted_boxes[x_interval_indices]
  410. x_indices_chunk = x_sorted_indices[x_interval_indices]
  411. # Sort selected boxes by y_min to prepare for Y-axis projection
  412. y_sorted_indices = x_boxes_chunk[:, 1].argsort()
  413. y_sorted_boxes_chunk = x_boxes_chunk[y_sorted_indices]
  414. y_sorted_indices_chunk = x_indices_chunk[y_sorted_indices]
  415. # Perform Y-axis projection
  416. y_projection = _projection_by_bboxes(boxes=y_sorted_boxes_chunk, axis=1)
  417. y_intervals = _split_projection_profile(y_projection, 0, min_gap)
  418. if not y_intervals:
  419. continue
  420. # If Y-axis cannot be further segmented, add current indices to results
  421. if len(y_intervals[0]) == 1:
  422. res.extend(y_sorted_indices_chunk)
  423. continue
  424. # Recursively process each segment defined by Y-axis projection
  425. for y_start, y_end in zip(*y_intervals):
  426. y_interval_indices = (y_start <= y_sorted_boxes_chunk[:, 1]) & (
  427. y_sorted_boxes_chunk[:, 1] < y_end
  428. )
  429. _recursive_xy_cut(
  430. y_sorted_boxes_chunk[y_interval_indices],
  431. y_sorted_indices_chunk[y_interval_indices],
  432. res,
  433. )
  434. def sort_by_xycut(block_bboxes, direction=0, min_gap=1):
  435. block_bboxes = np.asarray(block_bboxes).astype(int)
  436. res = []
  437. if direction == 1:
  438. _recursive_yx_cut(
  439. block_bboxes,
  440. np.arange(
  441. len(block_bboxes),
  442. ),
  443. res,
  444. min_gap,
  445. )
  446. else:
  447. _recursive_xy_cut(
  448. block_bboxes,
  449. np.arange(
  450. len(block_bboxes),
  451. ),
  452. res,
  453. min_gap,
  454. )
  455. return res
  456. def _img_array2path(data, save_path):
  457. """
  458. Save an image array to disk and return the file path.
  459. Args:
  460. data (np.ndarray): An image represented as a numpy array.
  461. save_path (str or Path): The base path where images should be saved.
  462. Returns:
  463. str: The relative path of the saved image file.
  464. """
  465. if isinstance(data, np.ndarray) and data.ndim == 3:
  466. # Generate a unique filename using UUID
  467. img_name = f"image_{uuid.uuid4().hex}.png"
  468. img_path = Path(save_path) / "imgs" / img_name
  469. img_path.parent.mkdir(
  470. parents=True,
  471. exist_ok=True,
  472. ) # Ensure the directory exists
  473. cv2.imwrite(str(img_path), data)
  474. return f"imgs/{img_name}"
  475. else:
  476. return ValueError
  477. def recursive_img_array2path(data, save_path, labels=[]):
  478. """
  479. Process a dictionary or list to save image arrays to disk and replace them with file paths.
  480. Args:
  481. data (dict or list): The data structure that may contain image arrays.
  482. save_path (str or Path): The base path where images should be saved.
  483. """
  484. if isinstance(data, dict):
  485. for k, v in data.items():
  486. if k in labels and isinstance(v, np.ndarray) and v.ndim == 3:
  487. data[k] = _img_array2path(v, save_path)
  488. else:
  489. recursive_img_array2path(v, save_path, labels)
  490. elif isinstance(data, list):
  491. for item in data:
  492. recursive_img_array2path(item, save_path, labels)
  493. def _calculate_overlap_area_2_minbox_area_ratio(bbox1, bbox2):
  494. """
  495. Calculate the ratio of the overlap area between bbox1 and bbox2
  496. to the area of the smaller bounding box.
  497. Args:
  498. bbox1 (list or tuple): Coordinates of the first bounding box [x_min, y_min, x_max, y_max].
  499. bbox2 (list or tuple): Coordinates of the second bounding box [x_min, y_min, x_max, y_max].
  500. Returns:
  501. float: The ratio of the overlap area to the area of the smaller bounding box.
  502. """
  503. x_left = max(bbox1[0], bbox2[0])
  504. y_top = max(bbox1[1], bbox2[1])
  505. x_right = min(bbox1[2], bbox2[2])
  506. y_bottom = min(bbox1[3], bbox2[3])
  507. if x_right <= x_left or y_bottom <= y_top:
  508. return 0.0
  509. # Calculate the area of the overlap
  510. intersection_area = (x_right - x_left) * (y_bottom - y_top)
  511. # Calculate the areas of both bounding boxes
  512. area_bbox1 = (bbox1[2] - bbox1[0]) * (bbox1[3] - bbox1[1])
  513. area_bbox2 = (bbox2[2] - bbox2[0]) * (bbox2[3] - bbox2[1])
  514. # Determine the minimum non-zero box area
  515. min_box_area = min(area_bbox1, area_bbox2)
  516. # Avoid division by zero in case of zero-area boxes
  517. if min_box_area == 0:
  518. return 0.0
  519. return intersection_area / min_box_area
  520. def _get_minbox_if_overlap_by_ratio(bbox1, bbox2, ratio, smaller=True):
  521. """
  522. Determine if the overlap area between two bounding boxes exceeds a given ratio
  523. and return the smaller (or larger) bounding box based on the `smaller` flag.
  524. Args:
  525. bbox1 (list or tuple): Coordinates of the first bounding box [x_min, y_min, x_max, y_max].
  526. bbox2 (list or tuple): Coordinates of the second bounding box [x_min, y_min, x_max, y_max].
  527. ratio (float): The overlap ratio threshold.
  528. smaller (bool): If True, return the smaller bounding box; otherwise, return the larger one.
  529. Returns:
  530. list or tuple: The selected bounding box or None if the overlap ratio is not exceeded.
  531. """
  532. # Calculate the areas of both bounding boxes
  533. area1 = (bbox1[2] - bbox1[0]) * (bbox1[3] - bbox1[1])
  534. area2 = (bbox2[2] - bbox2[0]) * (bbox2[3] - bbox2[1])
  535. # Calculate the overlap ratio using a helper function
  536. overlap_ratio = _calculate_overlap_area_2_minbox_area_ratio(bbox1, bbox2)
  537. # Check if the overlap ratio exceeds the threshold
  538. if overlap_ratio > ratio:
  539. if (area1 <= area2 and smaller) or (area1 >= area2 and not smaller):
  540. return 1
  541. else:
  542. return 2
  543. return None
  544. def _remove_overlap_blocks(blocks, threshold=0.65, smaller=True):
  545. """
  546. Remove overlapping blocks based on a specified overlap ratio threshold.
  547. Args:
  548. blocks (list): List of block dictionaries, each containing a 'layout_bbox' key.
  549. threshold (float): Ratio threshold to determine significant overlap.
  550. smaller (bool): If True, the smaller block in overlap is removed.
  551. Returns:
  552. tuple: A tuple containing the updated list of blocks and a list of dropped blocks.
  553. """
  554. dropped_blocks = []
  555. dropped_indexes = []
  556. # Iterate over each pair of blocks to find overlaps
  557. for i in range(len(blocks)):
  558. block1 = blocks[i]
  559. for j in range(i + 1, len(blocks)):
  560. block2 = blocks[j]
  561. # Skip blocks that are already marked for removal
  562. if i in dropped_indexes or j in dropped_indexes:
  563. continue
  564. # Check for overlap and determine which block to remove
  565. overlap_box_index = _get_minbox_if_overlap_by_ratio(
  566. block1["layout_bbox"],
  567. block2["layout_bbox"],
  568. threshold,
  569. smaller=smaller,
  570. )
  571. if overlap_box_index is not None:
  572. if overlap_box_index == 1:
  573. block_to_remove = block1
  574. drop_index = i
  575. else:
  576. block_to_remove = block2
  577. drop_index = j
  578. if drop_index not in dropped_indexes:
  579. dropped_indexes.append(drop_index)
  580. dropped_blocks.append(block_to_remove)
  581. dropped_indexes.sort()
  582. for i in reversed(dropped_indexes):
  583. del blocks[i]
  584. return blocks, dropped_blocks
  585. def _text_median_width(blocks):
  586. widths = [
  587. block["layout_bbox"][2] - block["layout_bbox"][0]
  588. for block in blocks
  589. if block["label"] in ["text"]
  590. ]
  591. return np.median(widths) if widths else float("inf")
  592. def _get_layout_property(blocks, median_width, no_mask_labels, threshold=0.8):
  593. """
  594. Determine the layout (single or double column) of text blocks.
  595. Args:
  596. blocks (list): List of block dictionaries containing 'label' and 'layout_bbox'.
  597. median_width (float): Median width of text blocks.
  598. threshold (float): Threshold for determining layout overlap.
  599. Returns:
  600. list: Updated list of blocks with layout information.
  601. """
  602. blocks.sort(
  603. key=lambda x: (
  604. x["layout_bbox"][0],
  605. (x["layout_bbox"][2] - x["layout_bbox"][0]),
  606. ),
  607. )
  608. check_single_layout = {}
  609. page_min_x, page_max_x = float("inf"), 0
  610. double_label_height = 0
  611. double_label_area = 0
  612. single_label_area = 0
  613. for i, block in enumerate(blocks):
  614. page_min_x = min(page_min_x, block["layout_bbox"][0])
  615. page_max_x = max(page_max_x, block["layout_bbox"][2])
  616. page_width = page_max_x - page_min_x
  617. for i, block in enumerate(blocks):
  618. if block["label"] not in no_mask_labels:
  619. continue
  620. x_min_i, _, x_max_i, _ = block["layout_bbox"]
  621. layout_length = x_max_i - x_min_i
  622. cover_count, cover_with_threshold_count = 0, 0
  623. match_block_with_threshold_indexes = []
  624. for j, other_block in enumerate(blocks):
  625. if i == j or other_block["label"] not in no_mask_labels:
  626. continue
  627. x_min_j, _, x_max_j, _ = other_block["layout_bbox"]
  628. x_match_min, x_match_max = max(
  629. x_min_i,
  630. x_min_j,
  631. ), min(x_max_i, x_max_j)
  632. match_block_iou = (x_match_max - x_match_min) / (x_max_j - x_min_j)
  633. if match_block_iou > 0:
  634. cover_count += 1
  635. if match_block_iou > threshold:
  636. cover_with_threshold_count += 1
  637. match_block_with_threshold_indexes.append(
  638. (j, match_block_iou),
  639. )
  640. x_min_i = x_match_max
  641. if x_min_i >= x_max_i:
  642. break
  643. if (
  644. layout_length > median_width * 1.3
  645. and (cover_with_threshold_count >= 2 or cover_count >= 2)
  646. ) or layout_length > 0.6 * page_width:
  647. # if layout_length > median_width * 1.3 and (cover_with_threshold_count >= 2):
  648. block["layout"] = "double"
  649. double_label_height += block["layout_bbox"][3] - block["layout_bbox"][1]
  650. double_label_area += (block["layout_bbox"][2] - block["layout_bbox"][0]) * (
  651. block["layout_bbox"][3] - block["layout_bbox"][1]
  652. )
  653. else:
  654. block["layout"] = "single"
  655. check_single_layout[i] = match_block_with_threshold_indexes
  656. # Check single-layout block
  657. for i, single_layout in check_single_layout.items():
  658. if single_layout:
  659. index, match_iou = single_layout[-1]
  660. if match_iou > 0.9 and blocks[index]["layout"] == "double":
  661. blocks[i]["layout"] = "double"
  662. double_label_height += (
  663. blocks[i]["layout_bbox"][3] - blocks[i]["layout_bbox"][1]
  664. )
  665. double_label_area += (
  666. blocks[i]["layout_bbox"][2] - blocks[i]["layout_bbox"][0]
  667. ) * (blocks[i]["layout_bbox"][3] - blocks[i]["layout_bbox"][1])
  668. else:
  669. single_label_area += (
  670. blocks[i]["layout_bbox"][2] - blocks[i]["layout_bbox"][0]
  671. ) * (blocks[i]["layout_bbox"][3] - blocks[i]["layout_bbox"][1])
  672. return blocks, (double_label_area > single_label_area)
  673. def _get_bbox_direction(input_bbox, ratio=1):
  674. """
  675. Determine if a bounding box is horizontal or vertical.
  676. Args:
  677. input_bbox (list): Bounding box [x_min, y_min, x_max, y_max].
  678. ratio (float): Ratio for determining orientation.
  679. Returns:
  680. bool: True if horizontal, False if vertical.
  681. """
  682. return (input_bbox[2] - input_bbox[0]) * ratio >= (input_bbox[3] - input_bbox[1])
  683. def _get_projection_iou(input_bbox, match_bbox, is_horizontal=True):
  684. """
  685. Calculate the IoU of lines between two bounding boxes.
  686. Args:
  687. input_bbox (list): First bounding box [x_min, y_min, x_max, y_max].
  688. match_bbox (list): Second bounding box [x_min, y_min, x_max, y_max].
  689. is_horizontal (bool): Whether to compare horizontally or vertically.
  690. Returns:
  691. float: Line IoU.
  692. """
  693. if is_horizontal:
  694. x_match_min = max(input_bbox[0], match_bbox[0])
  695. x_match_max = min(input_bbox[2], match_bbox[2])
  696. return (x_match_max - x_match_min) / (input_bbox[2] - input_bbox[0])
  697. else:
  698. y_match_min = max(input_bbox[1], match_bbox[1])
  699. y_match_max = min(input_bbox[3], match_bbox[3])
  700. return (y_match_max - y_match_min) / (input_bbox[3] - input_bbox[1])
  701. def _get_sub_category(blocks, title_labels):
  702. """
  703. Determine the layout of title and text blocks.
  704. Args:
  705. blocks (list): List of block dictionaries.
  706. title_labels (list): List of labels considered as titles.
  707. Returns:
  708. list: Updated list of blocks with title-text layout information.
  709. """
  710. sub_title_labels = ["paragraph_title"]
  711. vision_labels = ["image", "table", "chart", "figure"]
  712. for i, block1 in enumerate(blocks):
  713. if block1.get("title_text") is None:
  714. block1["title_text"] = []
  715. if block1.get("sub_title") is None:
  716. block1["sub_title"] = []
  717. if block1.get("vision_footnote") is None:
  718. block1["vision_footnote"] = []
  719. if block1.get("sub_label") is None:
  720. block1["sub_label"] = block1["label"]
  721. if (
  722. block1["label"] not in title_labels
  723. and block1["label"] not in sub_title_labels
  724. and block1["label"] not in vision_labels
  725. ):
  726. continue
  727. bbox1 = block1["layout_bbox"]
  728. x1, y1, x2, y2 = bbox1
  729. is_horizontal_1 = _get_bbox_direction(block1["layout_bbox"])
  730. left_up_title_text_distance = float("inf")
  731. left_up_title_text_index = -1
  732. left_up_title_text_direction = None
  733. right_down_title_text_distance = float("inf")
  734. right_down_title_text_index = -1
  735. right_down_title_text_direction = None
  736. for j, block2 in enumerate(blocks):
  737. if i == j:
  738. continue
  739. bbox2 = block2["layout_bbox"]
  740. x1_prime, y1_prime, x2_prime, y2_prime = bbox2
  741. is_horizontal_2 = _get_bbox_direction(bbox2)
  742. match_block_iou = _get_projection_iou(
  743. bbox2,
  744. bbox1,
  745. is_horizontal_1,
  746. )
  747. def distance_(is_horizontal, is_left_up):
  748. if is_horizontal:
  749. if is_left_up:
  750. return (y1 - y2_prime + 2) // 5 + x1_prime / 5000
  751. else:
  752. return (y1_prime - y2 + 2) // 5 + x1_prime / 5000
  753. else:
  754. if is_left_up:
  755. return (x1 - x2_prime + 2) // 5 + y1_prime / 5000
  756. else:
  757. return (x1_prime - x2 + 2) // 5 + y1_prime / 5000
  758. block_iou_threshold = 0.1
  759. if block1["label"] in sub_title_labels:
  760. match_block_iou = _calculate_overlap_area_2_minbox_area_ratio(
  761. bbox2,
  762. bbox1,
  763. )
  764. block_iou_threshold = 0.7
  765. if is_horizontal_1:
  766. if match_block_iou >= block_iou_threshold:
  767. left_up_distance = distance_(True, True)
  768. right_down_distance = distance_(True, False)
  769. if (
  770. y2_prime <= y1
  771. and left_up_distance <= left_up_title_text_distance
  772. ):
  773. left_up_title_text_distance = left_up_distance
  774. left_up_title_text_index = j
  775. left_up_title_text_direction = is_horizontal_2
  776. elif (
  777. y1_prime > y2
  778. and right_down_distance < right_down_title_text_distance
  779. ):
  780. right_down_title_text_distance = right_down_distance
  781. right_down_title_text_index = j
  782. right_down_title_text_direction = is_horizontal_2
  783. else:
  784. if match_block_iou >= block_iou_threshold:
  785. left_up_distance = distance_(False, True)
  786. right_down_distance = distance_(False, False)
  787. if (
  788. x2_prime <= x1
  789. and left_up_distance <= left_up_title_text_distance
  790. ):
  791. left_up_title_text_distance = left_up_distance
  792. left_up_title_text_index = j
  793. left_up_title_text_direction = is_horizontal_2
  794. elif (
  795. x1_prime > x2
  796. and right_down_distance < right_down_title_text_distance
  797. ):
  798. right_down_title_text_distance = right_down_distance
  799. right_down_title_text_index = j
  800. right_down_title_text_direction = is_horizontal_2
  801. height = bbox1[3] - bbox1[1]
  802. width = bbox1[2] - bbox1[0]
  803. title_text_weight = [0.8, 0.8]
  804. # title_text_weight = [2, 2]
  805. title_text = []
  806. sub_title = []
  807. vision_footnote = []
  808. def get_sub_category_(
  809. title_text_direction,
  810. title_text_index,
  811. label,
  812. is_left_up=True,
  813. ):
  814. direction_ = [1, 3] if is_left_up else [2, 4]
  815. if (
  816. title_text_direction == is_horizontal_1
  817. and title_text_index != -1
  818. and (label == "text" or label == "paragraph_title")
  819. ):
  820. bbox2 = blocks[title_text_index]["layout_bbox"]
  821. if is_horizontal_1:
  822. height1 = bbox2[3] - bbox2[1]
  823. width1 = bbox2[2] - bbox2[0]
  824. if label == "text":
  825. if (
  826. _nearest_edge_distance(bbox1, bbox2)[0] <= 15
  827. and block1["label"] in vision_labels
  828. and width1 < width
  829. and height1 < 0.5 * height
  830. ):
  831. blocks[title_text_index]["sub_label"] = "vision_footnote"
  832. vision_footnote.append(bbox2)
  833. elif (
  834. height1 < height * title_text_weight[0]
  835. and (width1 < width or width1 > 1.5 * width)
  836. and block1["label"] in title_labels
  837. ):
  838. blocks[title_text_index]["sub_label"] = "title_text"
  839. title_text.append((direction_[0], bbox2))
  840. elif (
  841. label == "paragraph_title"
  842. and block1["label"] in sub_title_labels
  843. ):
  844. sub_title.append(bbox2)
  845. else:
  846. height1 = bbox2[3] - bbox2[1]
  847. width1 = bbox2[2] - bbox2[0]
  848. if label == "text":
  849. if (
  850. _nearest_edge_distance(bbox1, bbox2)[0] <= 15
  851. and block1["label"] in vision_labels
  852. and height1 < height
  853. and width1 < 0.5 * width
  854. ):
  855. blocks[title_text_index]["sub_label"] = "vision_footnote"
  856. vision_footnote.append(bbox2)
  857. elif (
  858. width1 < width * title_text_weight[1]
  859. and block1["label"] in title_labels
  860. ):
  861. blocks[title_text_index]["sub_label"] = "title_text"
  862. title_text.append((direction_[1], bbox2))
  863. elif (
  864. label == "paragraph_title"
  865. and block1["label"] in sub_title_labels
  866. ):
  867. sub_title.append(bbox2)
  868. if (
  869. is_horizontal_1
  870. and abs(left_up_title_text_distance - right_down_title_text_distance) * 5
  871. > height
  872. ) or (
  873. not is_horizontal_1
  874. and abs(left_up_title_text_distance - right_down_title_text_distance) * 5
  875. > width
  876. ):
  877. if left_up_title_text_distance < right_down_title_text_distance:
  878. get_sub_category_(
  879. left_up_title_text_direction,
  880. left_up_title_text_index,
  881. blocks[left_up_title_text_index]["label"],
  882. True,
  883. )
  884. else:
  885. get_sub_category_(
  886. right_down_title_text_direction,
  887. right_down_title_text_index,
  888. blocks[right_down_title_text_index]["label"],
  889. False,
  890. )
  891. else:
  892. get_sub_category_(
  893. left_up_title_text_direction,
  894. left_up_title_text_index,
  895. blocks[left_up_title_text_index]["label"],
  896. True,
  897. )
  898. get_sub_category_(
  899. right_down_title_text_direction,
  900. right_down_title_text_index,
  901. blocks[right_down_title_text_index]["label"],
  902. False,
  903. )
  904. if block1["label"] in title_labels:
  905. if blocks[i].get("title_text") == []:
  906. blocks[i]["title_text"] = title_text
  907. if block1["label"] in sub_title_labels:
  908. if blocks[i].get("sub_title") == []:
  909. blocks[i]["sub_title"] = sub_title
  910. if block1["label"] in vision_labels:
  911. if blocks[i].get("vision_footnote") == []:
  912. blocks[i]["vision_footnote"] = vision_footnote
  913. return blocks
  914. def get_layout_ordering(data, no_mask_labels=[], already_sorted=False):
  915. """
  916. Process layout parsing results to remove overlapping bounding boxes
  917. and assign an ordering index based on their positions.
  918. Modifies:
  919. The 'parsing_result' list in 'layout_parsing_result' by adding an 'index' to each block.
  920. """
  921. if already_sorted:
  922. return data
  923. title_text_labels = ["doc_title"]
  924. title_labels = ["doc_title", "paragraph_title"]
  925. vision_labels = ["image", "table", "seal", "chart", "figure"]
  926. vision_title_labels = ["table_title", "chart_title", "figure_title"]
  927. parsing_result = data["sub_blocks"]
  928. parsing_result, _ = _remove_overlap_blocks(
  929. parsing_result,
  930. threshold=0.5,
  931. smaller=True,
  932. )
  933. parsing_result = _get_sub_category(parsing_result, title_text_labels)
  934. doc_flag = False
  935. median_width = _text_median_width(parsing_result)
  936. parsing_result, projection_direction = _get_layout_property(
  937. parsing_result,
  938. median_width,
  939. no_mask_labels=no_mask_labels,
  940. threshold=0.3,
  941. )
  942. # Convert bounding boxes to float and remove overlaps
  943. (
  944. double_text_blocks,
  945. title_text_blocks,
  946. title_blocks,
  947. vision_blocks,
  948. vision_title_blocks,
  949. vision_footnote_blocks,
  950. other_blocks,
  951. ) = ([], [], [], [], [], [], [])
  952. drop_indexes = []
  953. for index, block in enumerate(parsing_result):
  954. label = block["sub_label"]
  955. block["layout_bbox"] = list(map(int, block["layout_bbox"]))
  956. if label == "doc_title":
  957. doc_flag = True
  958. if label in no_mask_labels:
  959. if block["layout"] == "double":
  960. double_text_blocks.append(block)
  961. drop_indexes.append(index)
  962. elif label == "title_text":
  963. title_text_blocks.append(block)
  964. drop_indexes.append(index)
  965. elif label == "vision_footnote":
  966. vision_footnote_blocks.append(block)
  967. drop_indexes.append(index)
  968. elif label in vision_title_labels:
  969. vision_title_blocks.append(block)
  970. drop_indexes.append(index)
  971. elif label in title_labels:
  972. title_blocks.append(block)
  973. drop_indexes.append(index)
  974. elif label in vision_labels:
  975. vision_blocks.append(block)
  976. drop_indexes.append(index)
  977. else:
  978. other_blocks.append(block)
  979. drop_indexes.append(index)
  980. for index in sorted(drop_indexes, reverse=True):
  981. del parsing_result[index]
  982. if len(parsing_result) > 0:
  983. # single text label
  984. if len(double_text_blocks) > len(parsing_result) or projection_direction:
  985. parsing_result.extend(title_blocks + double_text_blocks)
  986. title_blocks = []
  987. double_text_blocks = []
  988. block_bboxes = [block["layout_bbox"] for block in parsing_result]
  989. block_bboxes.sort(
  990. key=lambda x: (
  991. x[0] // max(20, median_width),
  992. x[1],
  993. ),
  994. )
  995. block_bboxes = np.array(block_bboxes)
  996. sorted_indices = sort_by_xycut(
  997. block_bboxes,
  998. direction=1,
  999. min_gap=1,
  1000. )
  1001. else:
  1002. block_bboxes = [block["layout_bbox"] for block in parsing_result]
  1003. block_bboxes.sort(key=lambda x: (x[0] // 20, x[1]))
  1004. block_bboxes = np.array(block_bboxes)
  1005. sorted_indices = sort_by_xycut(
  1006. block_bboxes,
  1007. direction=0,
  1008. min_gap=20,
  1009. )
  1010. sorted_boxes = block_bboxes[sorted_indices].tolist()
  1011. for block in parsing_result:
  1012. block["index"] = sorted_boxes.index(block["layout_bbox"]) + 1
  1013. block["sub_index"] = sorted_boxes.index(block["layout_bbox"]) + 1
  1014. def nearest_match_(input_blocks, distance_type="manhattan", is_add_index=True):
  1015. for block in input_blocks:
  1016. bbox = block["layout_bbox"]
  1017. min_distance = float("inf")
  1018. min_distance_config = [
  1019. [float("inf"), float("inf")],
  1020. float("inf"),
  1021. float("inf"),
  1022. ] # for double text
  1023. nearest_gt_index = 0
  1024. for match_block in parsing_result:
  1025. match_bbox = match_block["layout_bbox"]
  1026. if distance_type == "nearest_iou_edge_distance":
  1027. distance, min_distance_config = _nearest_iou_edge_distance(
  1028. bbox,
  1029. match_bbox,
  1030. block["sub_label"],
  1031. vision_labels=vision_labels,
  1032. no_mask_labels=no_mask_labels,
  1033. median_width=median_width,
  1034. title_labels=title_labels,
  1035. title_text=block["title_text"],
  1036. sub_title=block["sub_title"],
  1037. min_distance_config=min_distance_config,
  1038. tolerance_len=10,
  1039. )
  1040. elif distance_type == "title_text":
  1041. if (
  1042. match_block["label"] in title_labels + ["abstract"]
  1043. and match_block["title_text"] != []
  1044. ):
  1045. iou_left_up = _calculate_overlap_area_2_minbox_area_ratio(
  1046. bbox,
  1047. match_block["title_text"][0][1],
  1048. )
  1049. iou_right_down = _calculate_overlap_area_2_minbox_area_ratio(
  1050. bbox,
  1051. match_block["title_text"][-1][1],
  1052. )
  1053. iou = 1 - max(iou_left_up, iou_right_down)
  1054. distance = _manhattan_distance(bbox, match_bbox) * iou
  1055. else:
  1056. distance = float("inf")
  1057. elif distance_type == "manhattan":
  1058. distance = _manhattan_distance(bbox, match_bbox)
  1059. elif distance_type == "vision_footnote":
  1060. if (
  1061. match_block["label"] in vision_labels
  1062. and match_block["vision_footnote"] != []
  1063. ):
  1064. iou_left_up = _calculate_overlap_area_2_minbox_area_ratio(
  1065. bbox,
  1066. match_block["vision_footnote"][0],
  1067. )
  1068. iou_right_down = _calculate_overlap_area_2_minbox_area_ratio(
  1069. bbox,
  1070. match_block["vision_footnote"][-1],
  1071. )
  1072. iou = 1 - max(iou_left_up, iou_right_down)
  1073. distance = _manhattan_distance(bbox, match_bbox) * iou
  1074. else:
  1075. distance = float("inf")
  1076. elif distance_type == "vision_body":
  1077. if (
  1078. match_block["label"] in vision_title_labels
  1079. and block["vision_footnote"] != []
  1080. ):
  1081. iou_left_up = _calculate_overlap_area_2_minbox_area_ratio(
  1082. match_bbox,
  1083. block["vision_footnote"][0],
  1084. )
  1085. iou_right_down = _calculate_overlap_area_2_minbox_area_ratio(
  1086. match_bbox,
  1087. block["vision_footnote"][-1],
  1088. )
  1089. iou = 1 - max(iou_left_up, iou_right_down)
  1090. distance = _manhattan_distance(bbox, match_bbox) * iou
  1091. else:
  1092. distance = float("inf")
  1093. else:
  1094. raise NotImplementedError
  1095. if distance < min_distance:
  1096. min_distance = distance
  1097. if is_add_index:
  1098. nearest_gt_index = match_block.get("index", 999)
  1099. else:
  1100. nearest_gt_index = match_block.get("sub_index", 999)
  1101. if is_add_index:
  1102. block["index"] = nearest_gt_index
  1103. else:
  1104. block["sub_index"] = nearest_gt_index
  1105. parsing_result.append(block)
  1106. # double text label
  1107. double_text_blocks.sort(
  1108. key=lambda x: (
  1109. x["layout_bbox"][1] // 10,
  1110. x["layout_bbox"][0] // median_width,
  1111. x["layout_bbox"][1] ** 2 + x["layout_bbox"][0] ** 2,
  1112. ),
  1113. )
  1114. nearest_match_(
  1115. double_text_blocks,
  1116. distance_type="nearest_iou_edge_distance",
  1117. )
  1118. parsing_result.sort(
  1119. key=lambda x: (x["index"], x["layout_bbox"][1], x["layout_bbox"][0]),
  1120. )
  1121. for idx, block in enumerate(parsing_result):
  1122. block["index"] = idx + 1
  1123. block["sub_index"] = idx + 1
  1124. # title label
  1125. title_blocks.sort(
  1126. key=lambda x: (
  1127. x["layout_bbox"][1] // 10,
  1128. x["layout_bbox"][0] // median_width,
  1129. x["layout_bbox"][1] ** 2 + x["layout_bbox"][0] ** 2,
  1130. ),
  1131. )
  1132. nearest_match_(title_blocks, distance_type="nearest_iou_edge_distance")
  1133. if doc_flag:
  1134. # text_sort_labels = ["doc_title","paragraph_title","abstract"]
  1135. text_sort_labels = ["doc_title"]
  1136. text_label_priority = {
  1137. label: priority for priority, label in enumerate(text_sort_labels)
  1138. }
  1139. doc_titles = []
  1140. for i, block in enumerate(parsing_result):
  1141. if block["label"] == "doc_title":
  1142. doc_titles.append(
  1143. (i, block["layout_bbox"][1], block["layout_bbox"][0]),
  1144. )
  1145. doc_titles.sort(key=lambda x: (x[1], x[2]))
  1146. first_doc_title_index = doc_titles[0][0]
  1147. parsing_result[first_doc_title_index]["index"] = 1
  1148. parsing_result.sort(
  1149. key=lambda x: (
  1150. x["index"],
  1151. text_label_priority.get(x["label"], 9999),
  1152. x["layout_bbox"][1],
  1153. x["layout_bbox"][0],
  1154. ),
  1155. )
  1156. else:
  1157. parsing_result.sort(
  1158. key=lambda x: (
  1159. x["index"],
  1160. x["layout_bbox"][1],
  1161. x["layout_bbox"][0],
  1162. ),
  1163. )
  1164. for idx, block in enumerate(parsing_result):
  1165. block["index"] = idx + 1
  1166. block["sub_index"] = idx + 1
  1167. # title-text label
  1168. nearest_match_(title_text_blocks, distance_type="title_text")
  1169. text_sort_labels = ["doc_title", "paragraph_title", "title_text"]
  1170. text_label_priority = {
  1171. label: priority for priority, label in enumerate(text_sort_labels)
  1172. }
  1173. parsing_result.sort(
  1174. key=lambda x: (
  1175. x["index"],
  1176. text_label_priority.get(x["sub_label"], 9999),
  1177. x["layout_bbox"][1],
  1178. x["layout_bbox"][0],
  1179. ),
  1180. )
  1181. for idx, block in enumerate(parsing_result):
  1182. block["index"] = idx + 1
  1183. block["sub_index"] = idx + 1
  1184. # image,figure,chart,seal label
  1185. nearest_match_(
  1186. vision_title_blocks,
  1187. distance_type="nearest_iou_edge_distance",
  1188. is_add_index=False,
  1189. )
  1190. parsing_result.sort(
  1191. key=lambda x: (
  1192. x["sub_index"],
  1193. x["layout_bbox"][1],
  1194. x["layout_bbox"][0],
  1195. ),
  1196. )
  1197. for idx, block in enumerate(parsing_result):
  1198. block["sub_index"] = idx + 1
  1199. # image,figure,chart,seal label
  1200. nearest_match_(
  1201. vision_blocks,
  1202. distance_type="nearest_iou_edge_distance",
  1203. is_add_index=False,
  1204. )
  1205. parsing_result.sort(
  1206. key=lambda x: (
  1207. x["sub_index"],
  1208. x["layout_bbox"][1],
  1209. x["layout_bbox"][0],
  1210. ),
  1211. )
  1212. for idx, block in enumerate(parsing_result):
  1213. block["sub_index"] = idx + 1
  1214. # vision footnote label
  1215. nearest_match_(
  1216. vision_footnote_blocks,
  1217. distance_type="vision_footnote",
  1218. is_add_index=False,
  1219. )
  1220. text_label_priority = {"vision_footnote": 9999}
  1221. parsing_result.sort(
  1222. key=lambda x: (
  1223. x["sub_index"],
  1224. text_label_priority.get(x["sub_label"], 0),
  1225. x["layout_bbox"][1],
  1226. x["layout_bbox"][0],
  1227. ),
  1228. )
  1229. for idx, block in enumerate(parsing_result):
  1230. block["sub_index"] = idx + 1
  1231. # header、footnote、header_image... label
  1232. nearest_match_(other_blocks, distance_type="manhattan", is_add_index=False)
  1233. return data
  1234. def _generate_input_data(parsing_result):
  1235. """
  1236. The evaluation input data is generated based on the parsing results.
  1237. :param parsing_result: A list containing the results of the layout parsing
  1238. :return: A formatted list of input data
  1239. """
  1240. input_data = [
  1241. {
  1242. "block_bbox": block["block_bbox"],
  1243. "sub_indices": [],
  1244. "sub_bboxes": [],
  1245. }
  1246. for block in parsing_result
  1247. ]
  1248. for block_index, block in enumerate(parsing_result):
  1249. sub_blocks = block["sub_blocks"]
  1250. get_layout_ordering(
  1251. block_index=block_index,
  1252. no_mask_labels=[
  1253. "text",
  1254. "formula",
  1255. "algorithm",
  1256. "reference",
  1257. "content",
  1258. "abstract",
  1259. ],
  1260. )
  1261. for sub_block in sub_blocks:
  1262. input_data[block_index]["sub_bboxes"].append(
  1263. list(map(int, sub_block["layout_bbox"])),
  1264. )
  1265. input_data[block_index]["sub_indices"].append(
  1266. int(sub_block["index"]),
  1267. )
  1268. return input_data
  1269. def _manhattan_distance(point1, point2, weight_x=1, weight_y=1):
  1270. return weight_x * abs(point1[0] - point2[0]) + weight_y * abs(point1[1] - point2[1])
  1271. def _calculate_horizontal_distance(
  1272. input_bbox,
  1273. match_bbox,
  1274. height,
  1275. disperse,
  1276. title_text,
  1277. ):
  1278. """
  1279. Calculate the horizontal distance between two bounding boxes, considering title text adjustments.
  1280. Args:
  1281. input_bbox (list): The bounding box coordinates [x1, y1, x2, y2] of the input object.
  1282. match_bbox (list): The bounding box coordinates [x1', y1', x2', y2'] of the object to match against.
  1283. height (int): The height of the input bounding box used for normalization.
  1284. disperse (int): The dispersion factor used to normalize the horizontal distance.
  1285. title_text (list): A list of tuples containing title text information and their bounding box coordinates.
  1286. Format: [(position_indicator, [x1, y1, x2, y2]), ...].
  1287. Returns:
  1288. float: The calculated horizontal distance taking into account the title text adjustments.
  1289. """
  1290. x1, y1, x2, y2 = input_bbox
  1291. x1_prime, y1_prime, x2_prime, y2_prime = match_bbox
  1292. if y2 < y1_prime:
  1293. if title_text and title_text[-1][0] == 2:
  1294. y2 += title_text[-1][1][3] - title_text[-1][1][1]
  1295. distance1 = (y1_prime - y2) * 0.5
  1296. else:
  1297. if title_text and title_text[0][0] == 1:
  1298. y1 -= title_text[0][1][3] - title_text[0][1][1]
  1299. distance1 = y1 - y2_prime
  1300. return (
  1301. abs(x2_prime - x1) // disperse + distance1 // height + distance1 / 5000
  1302. ) # if page max size == 5000
  1303. def _calculate_vertical_distance(input_bbox, match_bbox, width, disperse, title_text):
  1304. """
  1305. Calculate the vertical distance between two bounding boxes, considering title text adjustments.
  1306. Args:
  1307. input_bbox (list): The bounding box coordinates [x1, y1, x2, y2] of the input object.
  1308. match_bbox (list): The bounding box coordinates [x1', y1', x2', y2'] of the object to match against.
  1309. width (int): The width of the input bounding box used for normalization.
  1310. disperse (int): The dispersion factor used to normalize the vertical distance.
  1311. title_text (list): A list of tuples containing title text information and their bounding box coordinates.
  1312. Format: [(position_indicator, [x1, y1, x2, y2]), ...].
  1313. Returns:
  1314. float: The calculated vertical distance taking into account the title text adjustments.
  1315. """
  1316. x1, y1, x2, y2 = input_bbox
  1317. x1_prime, y1_prime, x2_prime, y2_prime = match_bbox
  1318. if x1 > x2_prime:
  1319. if title_text and title_text[0][0] == 3:
  1320. x1 -= title_text[0][1][2] - title_text[0][1][0]
  1321. distance2 = (x1 - x2_prime) * 0.5
  1322. else:
  1323. if title_text and title_text[-1][0] == 4:
  1324. x2 += title_text[-1][1][2] - title_text[-1][1][0]
  1325. distance2 = x1_prime - x2
  1326. return abs(y2_prime - y1) // disperse + distance2 // width + distance2 / 5000
  1327. def _nearest_edge_distance(
  1328. input_bbox,
  1329. match_bbox,
  1330. weight=[1, 1, 1, 1],
  1331. label="text",
  1332. no_mask_labels=[],
  1333. min_edge_distances_config=[],
  1334. tolerance_len=10,
  1335. ):
  1336. """
  1337. Calculate the nearest edge distance between two bounding boxes, considering directional weights.
  1338. Args:
  1339. input_bbox (list): The bounding box coordinates [x1, y1, x2, y2] of the input object.
  1340. match_bbox (list): The bounding box coordinates [x1', y1', x2', y2'] of the object to match against.
  1341. weight (list, optional): Directional weights for the edge distances [left, right, up, down]. Defaults to [1, 1, 1, 1].
  1342. label (str, optional): The label/type of the object in the bounding box (e.g., 'text'). Defaults to 'text'.
  1343. no_mask_labels (list, optional): Labels for which no masking is applied when calculating edge distances. Defaults to an empty list.
  1344. min_edge_distances_config (list, optional): Configuration for minimum edge distances [min_edge_distance_x, min_edge_distance_y].
  1345. Defaults to [float('inf'), float('inf')].
  1346. Returns:
  1347. tuple: A tuple containing:
  1348. - The calculated minimum edge distance between the bounding boxes.
  1349. - A list with the minimum edge distances in the x and y directions.
  1350. """
  1351. match_bbox_iou = _calculate_overlap_area_2_minbox_area_ratio(
  1352. input_bbox,
  1353. match_bbox,
  1354. )
  1355. if match_bbox_iou > 0 and label not in no_mask_labels:
  1356. return 0, [0, 0]
  1357. if not min_edge_distances_config:
  1358. min_edge_distances_config = [float("inf"), float("inf")]
  1359. min_edge_distance_x, min_edge_distance_y = min_edge_distances_config
  1360. x1, y1, x2, y2 = input_bbox
  1361. x1_prime, y1_prime, x2_prime, y2_prime = match_bbox
  1362. direction_num = 0
  1363. distance_x = float("inf")
  1364. distance_y = float("inf")
  1365. distance = [float("inf")] * 4
  1366. # input_bbox is to the left of match_bbox
  1367. if x2 < x1_prime:
  1368. direction_num += 1
  1369. distance[0] = x1_prime - x2
  1370. if abs(distance[0] - min_edge_distance_x) <= tolerance_len:
  1371. distance_x = min_edge_distance_x * weight[0]
  1372. else:
  1373. distance_x = distance[0] * weight[0]
  1374. # input_bbox is to the right of match_bbox
  1375. elif x1 > x2_prime:
  1376. direction_num += 1
  1377. distance[1] = x1 - x2_prime
  1378. if abs(distance[1] - min_edge_distance_x) <= tolerance_len:
  1379. distance_x = min_edge_distance_x * weight[1]
  1380. else:
  1381. distance_x = distance[1] * weight[1]
  1382. elif match_bbox_iou > 0:
  1383. distance[0] = 0
  1384. distance_x = 0
  1385. # input_bbox is above match_bbox
  1386. if y2 < y1_prime:
  1387. direction_num += 1
  1388. distance[2] = y1_prime - y2
  1389. if abs(distance[2] - min_edge_distance_y) <= tolerance_len:
  1390. distance_y = min_edge_distance_y * weight[2]
  1391. else:
  1392. distance_y = distance[2] * weight[2]
  1393. if label in no_mask_labels:
  1394. distance_y = max(0.1, distance_y) * 100
  1395. # input_bbox is below match_bbox
  1396. elif y1 > y2_prime:
  1397. direction_num += 1
  1398. distance[3] = y1 - y2_prime
  1399. if abs(distance[3] - min_edge_distance_y) <= tolerance_len:
  1400. distance_y = min_edge_distance_y * weight[3]
  1401. else:
  1402. distance_y = distance[3] * weight[3]
  1403. elif match_bbox_iou > 0:
  1404. distance[2] = 0
  1405. distance_y = 0
  1406. if direction_num == 2:
  1407. return (distance_x + distance_y), [
  1408. min(distance[0], distance[1]),
  1409. min(distance[2], distance[3]),
  1410. ]
  1411. else:
  1412. return min(distance_x, distance_y), [
  1413. min(distance[0], distance[1]),
  1414. min(distance[2], distance[3]),
  1415. ]
  1416. def _get_weights(label, horizontal):
  1417. """Define weights based on the label and orientation."""
  1418. if label == "doc_title":
  1419. return (
  1420. [1, 0.1, 0.1, 1] if horizontal else [0.2, 0.1, 1, 1]
  1421. ) # left-down , right-left
  1422. elif label in [
  1423. "paragraph_title",
  1424. "abstract",
  1425. "figure_title",
  1426. "chart_title",
  1427. "image",
  1428. "seal",
  1429. "chart",
  1430. "figure",
  1431. ]:
  1432. return [1, 1, 0.1, 1] # down
  1433. else:
  1434. return [1, 1, 1, 0.1] # up
  1435. def _nearest_iou_edge_distance(
  1436. input_bbox,
  1437. match_bbox,
  1438. label,
  1439. vision_labels,
  1440. no_mask_labels,
  1441. median_width=-1,
  1442. title_labels=[],
  1443. title_text=[],
  1444. sub_title=[],
  1445. min_distance_config=[],
  1446. tolerance_len=10,
  1447. ):
  1448. """
  1449. Calculate the nearest IOU edge distance between two bounding boxes.
  1450. Args:
  1451. input_bbox (list): The bounding box coordinates [x1, y1, x2, y2] of the input object.
  1452. match_bbox (list): The bounding box coordinates [x1', y1', x2', y2'] of the object to match against.
  1453. label (str): The label/type of the object in the bounding box (e.g., 'image', 'text', etc.).
  1454. no_mask_labels (list): Labels for which no masking is applied when calculating edge distances.
  1455. median_width (int, optional): The median width for title dispersion calculation. Defaults to -1.
  1456. title_labels (list, optional): Labels that indicate the object is a title. Defaults to an empty list.
  1457. title_text (list, optional): Text content associated with title labels. Defaults to an empty list.
  1458. sub_title (list, optional): List of subtitle bounding boxes to adjust the input_bbox. Defaults to an empty list.
  1459. min_distance_config (list, optional): Configuration for minimum distances [min_edge_distances_config, up_edge_distances_config, total_distance].
  1460. Returns:
  1461. tuple: A tuple containing the calculated distance and updated minimum distance configuration.
  1462. """
  1463. x1, y1, x2, y2 = input_bbox
  1464. x1_prime, y1_prime, x2_prime, y2_prime = match_bbox
  1465. min_edge_distances_config, up_edge_distances_config, total_distance = (
  1466. min_distance_config
  1467. )
  1468. iou_distance = 0
  1469. if label in vision_labels:
  1470. horizontal1 = horizontal2 = True
  1471. else:
  1472. horizontal1 = _get_bbox_direction(input_bbox)
  1473. horizontal2 = _get_bbox_direction(match_bbox, 3)
  1474. if (
  1475. horizontal1 != horizontal2
  1476. or _get_projection_iou(input_bbox, match_bbox, horizontal1) < 0.01
  1477. ):
  1478. iou_distance = 1
  1479. elif label == "doc_title" or (label in title_labels and title_text):
  1480. # Calculate distance for titles
  1481. disperse = max(1, median_width)
  1482. width = x2 - x1
  1483. height = y2 - y1
  1484. if horizontal1:
  1485. return (
  1486. _calculate_horizontal_distance(
  1487. input_bbox,
  1488. match_bbox,
  1489. height,
  1490. disperse,
  1491. title_text,
  1492. ),
  1493. min_distance_config,
  1494. )
  1495. else:
  1496. return (
  1497. _calculate_vertical_distance(
  1498. input_bbox,
  1499. match_bbox,
  1500. width,
  1501. disperse,
  1502. title_text,
  1503. ),
  1504. min_distance_config,
  1505. )
  1506. # Adjust input_bbox based on sub_title
  1507. if sub_title:
  1508. for sub in sub_title:
  1509. x1_, y1_, x2_, y2_ = sub
  1510. x1, y1, x2, y2 = (
  1511. min(x1, x1_),
  1512. min(
  1513. y1,
  1514. y1_,
  1515. ),
  1516. max(x2, x2_),
  1517. max(y2, y2_),
  1518. )
  1519. input_bbox = [x1, y1, x2, y2]
  1520. # Calculate edge distance
  1521. weight = _get_weights(label, horizontal1)
  1522. if label == "abstract":
  1523. tolerance_len *= 3
  1524. edge_distance, edge_distance_config = _nearest_edge_distance(
  1525. input_bbox,
  1526. match_bbox,
  1527. weight,
  1528. label=label,
  1529. no_mask_labels=no_mask_labels,
  1530. min_edge_distances_config=min_edge_distances_config,
  1531. tolerance_len=tolerance_len,
  1532. )
  1533. # Weights for combining distances
  1534. iou_edge_weight = [10**6, 10**3, 1, 0.001]
  1535. # Calculate up and left edge distances
  1536. up_edge_distance = y1_prime
  1537. left_edge_distance = x1_prime
  1538. if (
  1539. label in no_mask_labels or label == "paragraph_title" or label in vision_labels
  1540. ) and y1 > y2_prime:
  1541. up_edge_distance = -y2_prime
  1542. left_edge_distance = -x2_prime
  1543. min_up_edge_distance = up_edge_distances_config
  1544. if abs(min_up_edge_distance - up_edge_distance) <= tolerance_len:
  1545. up_edge_distance = min_up_edge_distance
  1546. # Calculate total distance
  1547. distance = (
  1548. iou_distance * iou_edge_weight[0]
  1549. + edge_distance * iou_edge_weight[1]
  1550. + up_edge_distance * iou_edge_weight[2]
  1551. + left_edge_distance * iou_edge_weight[3]
  1552. )
  1553. # Update minimum distance configuration if a smaller distance is found
  1554. if total_distance > distance:
  1555. edge_distance_config = [
  1556. min(min_edge_distances_config[0], edge_distance_config[0]),
  1557. min(min_edge_distances_config[1], edge_distance_config[1]),
  1558. ]
  1559. min_distance_config = [
  1560. edge_distance_config,
  1561. min(up_edge_distance, up_edge_distances_config),
  1562. distance,
  1563. ]
  1564. return distance, min_distance_config
  1565. def get_show_color(label):
  1566. label_colors = {
  1567. # Medium Blue (from 'titles_list')
  1568. "paragraph_title": (102, 102, 255, 100),
  1569. "doc_title": (255, 248, 220, 100), # Cornsilk
  1570. # Light Yellow (from 'tables_caption_list')
  1571. "table_title": (255, 255, 102, 100),
  1572. # Sky Blue (from 'imgs_caption_list')
  1573. "figure_title": (102, 178, 255, 100),
  1574. "chart_title": (221, 160, 221, 100), # Plum
  1575. "vision_footnote": (144, 238, 144, 100), # Light Green
  1576. # Deep Purple (from 'texts_list')
  1577. "text": (153, 0, 76, 100),
  1578. # Bright Green (from 'interequations_list')
  1579. "formula": (0, 255, 0, 100),
  1580. "abstract": (255, 239, 213, 100), # Papaya Whip
  1581. # Medium Green (from 'lists_list' and 'indexs_list')
  1582. "content": (40, 169, 92, 100),
  1583. # Neutral Gray (from 'dropped_bbox_list')
  1584. "seal": (158, 158, 158, 100),
  1585. # Olive Yellow (from 'tables_body_list')
  1586. "table": (204, 204, 0, 100),
  1587. # Bright Green (from 'imgs_body_list')
  1588. "image": (153, 255, 51, 100),
  1589. # Bright Green (from 'imgs_body_list')
  1590. "figure": (153, 255, 51, 100),
  1591. "chart": (216, 191, 216, 100), # Thistle
  1592. # Pale Yellow-Green (from 'tables_footnote_list')
  1593. "reference": (229, 255, 204, 100),
  1594. "algorithm": (255, 250, 240, 100), # Floral White
  1595. }
  1596. default_color = (158, 158, 158, 100)
  1597. return label_colors.get(label, default_color)