draw_bbox.py 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307
  1. from magic_pdf.libs.commons import fitz # PyMuPDF
  2. from magic_pdf.libs.Constants import CROSS_PAGE
  3. from magic_pdf.libs.ocr_content_type import BlockType, CategoryId, ContentType
  4. from magic_pdf.model.magic_model import MagicModel
  5. def draw_bbox_without_number(i, bbox_list, page, rgb_config, fill_config):
  6. new_rgb = []
  7. for item in rgb_config:
  8. item = float(item) / 255
  9. new_rgb.append(item)
  10. page_data = bbox_list[i]
  11. for bbox in page_data:
  12. x0, y0, x1, y1 = bbox
  13. rect_coords = fitz.Rect(x0, y0, x1, y1) # Define the rectangle
  14. if fill_config:
  15. page.draw_rect(
  16. rect_coords,
  17. color=None,
  18. fill=new_rgb,
  19. fill_opacity=0.3,
  20. width=0.5,
  21. overlay=True,
  22. ) # Draw the rectangle
  23. else:
  24. page.draw_rect(
  25. rect_coords,
  26. color=new_rgb,
  27. fill=None,
  28. fill_opacity=1,
  29. width=0.5,
  30. overlay=True,
  31. ) # Draw the rectangle
  32. def draw_bbox_with_number(i, bbox_list, page, rgb_config, fill_config):
  33. new_rgb = []
  34. for item in rgb_config:
  35. item = float(item) / 255
  36. new_rgb.append(item)
  37. page_data = bbox_list[i]
  38. for j, bbox in enumerate(page_data):
  39. x0, y0, x1, y1 = bbox
  40. rect_coords = fitz.Rect(x0, y0, x1, y1) # Define the rectangle
  41. if fill_config:
  42. page.draw_rect(
  43. rect_coords,
  44. color=None,
  45. fill=new_rgb,
  46. fill_opacity=0.3,
  47. width=0.5,
  48. overlay=True,
  49. ) # Draw the rectangle
  50. else:
  51. page.draw_rect(
  52. rect_coords,
  53. color=new_rgb,
  54. fill=None,
  55. fill_opacity=1,
  56. width=0.5,
  57. overlay=True,
  58. ) # Draw the rectangle
  59. page.insert_text(
  60. (x0, y0 + 10), str(j + 1), fontsize=10, color=new_rgb
  61. ) # Insert the index in the top left corner of the rectangle
  62. def draw_layout_bbox(pdf_info, pdf_bytes, out_path, filename):
  63. layout_bbox_list = []
  64. dropped_bbox_list = []
  65. tables_list, tables_body_list = [], []
  66. tables_caption_list, tables_footnote_list = [], []
  67. imgs_list, imgs_body_list, imgs_caption_list = [], [], []
  68. titles_list = []
  69. texts_list = []
  70. interequations_list = []
  71. for page in pdf_info:
  72. page_layout_list = []
  73. page_dropped_list = []
  74. tables, tables_body, tables_caption, tables_footnote = [], [], [], []
  75. imgs, imgs_body, imgs_caption = [], [], []
  76. titles = []
  77. texts = []
  78. interequations = []
  79. for layout in page['layout_bboxes']:
  80. page_layout_list.append(layout['layout_bbox'])
  81. layout_bbox_list.append(page_layout_list)
  82. for dropped_bbox in page['discarded_blocks']:
  83. page_dropped_list.append(dropped_bbox['bbox'])
  84. dropped_bbox_list.append(page_dropped_list)
  85. for block in page['para_blocks']:
  86. bbox = block['bbox']
  87. if block['type'] == BlockType.Table:
  88. tables.append(bbox)
  89. for nested_block in block['blocks']:
  90. bbox = nested_block['bbox']
  91. if nested_block['type'] == BlockType.TableBody:
  92. tables_body.append(bbox)
  93. elif nested_block['type'] == BlockType.TableCaption:
  94. tables_caption.append(bbox)
  95. elif nested_block['type'] == BlockType.TableFootnote:
  96. tables_footnote.append(bbox)
  97. elif block['type'] == BlockType.Image:
  98. imgs.append(bbox)
  99. for nested_block in block['blocks']:
  100. bbox = nested_block['bbox']
  101. if nested_block['type'] == BlockType.ImageBody:
  102. imgs_body.append(bbox)
  103. elif nested_block['type'] == BlockType.ImageCaption:
  104. imgs_caption.append(bbox)
  105. elif block['type'] == BlockType.Title:
  106. titles.append(bbox)
  107. elif block['type'] == BlockType.Text:
  108. texts.append(bbox)
  109. elif block['type'] == BlockType.InterlineEquation:
  110. interequations.append(bbox)
  111. tables_list.append(tables)
  112. tables_body_list.append(tables_body)
  113. tables_caption_list.append(tables_caption)
  114. tables_footnote_list.append(tables_footnote)
  115. imgs_list.append(imgs)
  116. imgs_body_list.append(imgs_body)
  117. imgs_caption_list.append(imgs_caption)
  118. titles_list.append(titles)
  119. texts_list.append(texts)
  120. interequations_list.append(interequations)
  121. pdf_docs = fitz.open('pdf', pdf_bytes)
  122. for i, page in enumerate(pdf_docs):
  123. draw_bbox_with_number(i, layout_bbox_list, page, [255, 0, 0], False)
  124. draw_bbox_without_number(i, dropped_bbox_list, page, [158, 158, 158],
  125. True)
  126. draw_bbox_without_number(i, tables_list, page, [153, 153, 0],
  127. True) # color !
  128. draw_bbox_without_number(i, tables_body_list, page, [204, 204, 0],
  129. True)
  130. draw_bbox_without_number(i, tables_caption_list, page, [255, 255, 102],
  131. True)
  132. draw_bbox_without_number(i, tables_footnote_list, page,
  133. [229, 255, 204], True)
  134. draw_bbox_without_number(i, imgs_list, page, [51, 102, 0], True)
  135. draw_bbox_without_number(i, imgs_body_list, page, [153, 255, 51], True)
  136. draw_bbox_without_number(i, imgs_caption_list, page, [102, 178, 255],
  137. True)
  138. draw_bbox_without_number(i, titles_list, page, [102, 102, 255], True)
  139. draw_bbox_without_number(i, texts_list, page, [153, 0, 76], True)
  140. draw_bbox_without_number(i, interequations_list, page, [0, 255, 0],
  141. True)
  142. # Save the PDF
  143. pdf_docs.save(f'{out_path}/{filename}_layout.pdf')
  144. def draw_span_bbox(pdf_info, pdf_bytes, out_path, filename):
  145. text_list = []
  146. inline_equation_list = []
  147. interline_equation_list = []
  148. image_list = []
  149. table_list = []
  150. dropped_list = []
  151. next_page_text_list = []
  152. next_page_inline_equation_list = []
  153. def get_span_info(span):
  154. if span['type'] == ContentType.Text:
  155. if span.get(CROSS_PAGE, False):
  156. next_page_text_list.append(span['bbox'])
  157. else:
  158. page_text_list.append(span['bbox'])
  159. elif span['type'] == ContentType.InlineEquation:
  160. if span.get(CROSS_PAGE, False):
  161. next_page_inline_equation_list.append(span['bbox'])
  162. else:
  163. page_inline_equation_list.append(span['bbox'])
  164. elif span['type'] == ContentType.InterlineEquation:
  165. page_interline_equation_list.append(span['bbox'])
  166. elif span['type'] == ContentType.Image:
  167. page_image_list.append(span['bbox'])
  168. elif span['type'] == ContentType.Table:
  169. page_table_list.append(span['bbox'])
  170. for page in pdf_info:
  171. page_text_list = []
  172. page_inline_equation_list = []
  173. page_interline_equation_list = []
  174. page_image_list = []
  175. page_table_list = []
  176. page_dropped_list = []
  177. # 将跨页的span放到移动到下一页的列表中
  178. if len(next_page_text_list) > 0:
  179. page_text_list.extend(next_page_text_list)
  180. next_page_text_list.clear()
  181. if len(next_page_inline_equation_list) > 0:
  182. page_inline_equation_list.extend(next_page_inline_equation_list)
  183. next_page_inline_equation_list.clear()
  184. # 构造dropped_list
  185. for block in page['discarded_blocks']:
  186. if block['type'] == BlockType.Discarded:
  187. for line in block['lines']:
  188. for span in line['spans']:
  189. page_dropped_list.append(span['bbox'])
  190. dropped_list.append(page_dropped_list)
  191. # 构造其余useful_list
  192. for block in page['para_blocks']:
  193. if block['type'] in [
  194. BlockType.Text,
  195. BlockType.Title,
  196. BlockType.InterlineEquation,
  197. ]:
  198. for line in block['lines']:
  199. for span in line['spans']:
  200. get_span_info(span)
  201. elif block['type'] in [BlockType.Image, BlockType.Table]:
  202. for sub_block in block['blocks']:
  203. for line in sub_block['lines']:
  204. for span in line['spans']:
  205. get_span_info(span)
  206. text_list.append(page_text_list)
  207. inline_equation_list.append(page_inline_equation_list)
  208. interline_equation_list.append(page_interline_equation_list)
  209. image_list.append(page_image_list)
  210. table_list.append(page_table_list)
  211. pdf_docs = fitz.open('pdf', pdf_bytes)
  212. for i, page in enumerate(pdf_docs):
  213. # 获取当前页面的数据
  214. draw_bbox_without_number(i, text_list, page, [255, 0, 0], False)
  215. draw_bbox_without_number(i, inline_equation_list, page, [0, 255, 0],
  216. False)
  217. draw_bbox_without_number(i, interline_equation_list, page, [0, 0, 255],
  218. False)
  219. draw_bbox_without_number(i, image_list, page, [255, 204, 0], False)
  220. draw_bbox_without_number(i, table_list, page, [204, 0, 255], False)
  221. draw_bbox_without_number(i, dropped_list, page, [158, 158, 158], False)
  222. # Save the PDF
  223. pdf_docs.save(f'{out_path}/{filename}_spans.pdf')
  224. def drow_model_bbox(model_list: list, pdf_bytes, out_path, filename):
  225. dropped_bbox_list = []
  226. tables_body_list, tables_caption_list, tables_footnote_list = [], [], []
  227. imgs_body_list, imgs_caption_list = [], []
  228. titles_list = []
  229. texts_list = []
  230. interequations_list = []
  231. pdf_docs = fitz.open('pdf', pdf_bytes)
  232. magic_model = MagicModel(model_list, pdf_docs)
  233. for i in range(len(model_list)):
  234. page_dropped_list = []
  235. tables_body, tables_caption, tables_footnote = [], [], []
  236. imgs_body, imgs_caption = [], []
  237. titles = []
  238. texts = []
  239. interequations = []
  240. page_info = magic_model.get_model_list(i)
  241. layout_dets = page_info['layout_dets']
  242. for layout_det in layout_dets:
  243. bbox = layout_det['bbox']
  244. if layout_det['category_id'] == CategoryId.Text:
  245. texts.append(bbox)
  246. elif layout_det['category_id'] == CategoryId.Title:
  247. titles.append(bbox)
  248. elif layout_det['category_id'] == CategoryId.TableBody:
  249. tables_body.append(bbox)
  250. elif layout_det['category_id'] == CategoryId.TableCaption:
  251. tables_caption.append(bbox)
  252. elif layout_det['category_id'] == CategoryId.TableFootnote:
  253. tables_footnote.append(bbox)
  254. elif layout_det['category_id'] == CategoryId.ImageBody:
  255. imgs_body.append(bbox)
  256. elif layout_det['category_id'] == CategoryId.ImageCaption:
  257. imgs_caption.append(bbox)
  258. elif layout_det[
  259. 'category_id'] == CategoryId.InterlineEquation_YOLO:
  260. interequations.append(bbox)
  261. elif layout_det['category_id'] == CategoryId.Abandon:
  262. page_dropped_list.append(bbox)
  263. tables_body_list.append(tables_body)
  264. tables_caption_list.append(tables_caption)
  265. tables_footnote_list.append(tables_footnote)
  266. imgs_body_list.append(imgs_body)
  267. imgs_caption_list.append(imgs_caption)
  268. titles_list.append(titles)
  269. texts_list.append(texts)
  270. interequations_list.append(interequations)
  271. dropped_bbox_list.append(page_dropped_list)
  272. for i, page in enumerate(pdf_docs):
  273. draw_bbox_with_number(i, dropped_bbox_list, page, [158, 158, 158],
  274. True) # color !
  275. draw_bbox_with_number(i, tables_body_list, page, [204, 204, 0], True)
  276. draw_bbox_with_number(i, tables_caption_list, page, [255, 255, 102],
  277. True)
  278. draw_bbox_with_number(i, tables_footnote_list, page, [229, 255, 204],
  279. True)
  280. draw_bbox_with_number(i, imgs_body_list, page, [153, 255, 51], True)
  281. draw_bbox_with_number(i, imgs_caption_list, page, [102, 178, 255],
  282. True)
  283. draw_bbox_with_number(i, titles_list, page, [102, 102, 255], True)
  284. draw_bbox_with_number(i, texts_list, page, [153, 0, 76], True)
  285. draw_bbox_with_number(i, interequations_list, page, [0, 255, 0], True)
  286. # Save the PDF
  287. pdf_docs.save(f'{out_path}/{filename}_model.pdf')