draw_bbox.py 15 KB

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