utils.py 76 KB

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