vis_utils.py 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310
  1. from libs.commons import fitz
  2. import os
  3. from loguru import logger
  4. from layout.bbox_sort import CONTENT_TYPE_IDX
  5. def draw_bbox_on_page(raw_pdf_doc: fitz.Document, paras_dict:dict, save_path: str):
  6. """
  7. 在page上画出bbox,保存到save_path
  8. """
  9. # 检查文件是否存在
  10. is_new_pdf = False
  11. if os.path.exists(save_path):
  12. # 打开现有的 PDF 文件
  13. doc = fitz.open(save_path)
  14. else:
  15. # 创建一个新的空白 PDF 文件
  16. is_new_pdf = True
  17. doc = fitz.open('')
  18. color_map = {
  19. 'image': fitz.pdfcolor["yellow"],
  20. 'text': fitz.pdfcolor['blue'],
  21. "table": fitz.pdfcolor['green']
  22. }
  23. for k, v in paras_dict.items():
  24. page_idx = v['page_idx']
  25. width = raw_pdf_doc[page_idx].rect.width
  26. height = raw_pdf_doc[page_idx].rect.height
  27. new_page = doc.new_page(width=width, height=height)
  28. shape = new_page.new_shape()
  29. for order, block in enumerate(v['preproc_blocks']):
  30. rect = fitz.Rect(block['bbox'])
  31. shape = new_page.new_shape()
  32. shape.draw_rect(rect)
  33. shape.finish(color=None, fill=color_map['text'], fill_opacity=0.2)
  34. shape.finish()
  35. shape.commit()
  36. for img in v['images']:
  37. # 原始box画上去
  38. rect = fitz.Rect(img['bbox'])
  39. shape = new_page.new_shape()
  40. shape.draw_rect(rect)
  41. shape.finish(color=None, fill=fitz.pdfcolor['yellow'])
  42. shape.finish()
  43. shape.commit()
  44. for img in v['image_backup']:
  45. # 原始box画上去
  46. rect = fitz.Rect(img['bbox'])
  47. shape = new_page.new_shape()
  48. shape.draw_rect(rect)
  49. shape.finish(color=fitz.pdfcolor['yellow'], fill=None)
  50. shape.finish()
  51. shape.commit()
  52. for tb in v['droped_text_block']:
  53. # 原始box画上去
  54. rect = fitz.Rect(tb['bbox'])
  55. shape = new_page.new_shape()
  56. shape.draw_rect(rect)
  57. shape.finish(color=None, fill=fitz.pdfcolor['black'], fill_opacity=0.4)
  58. shape.finish()
  59. shape.commit()
  60. # TODO table
  61. for tb in v['tables']:
  62. rect = fitz.Rect(tb['bbox'])
  63. shape = new_page.new_shape()
  64. shape.draw_rect(rect)
  65. shape.finish(color=None, fill=fitz.pdfcolor['green'], fill_opacity=0.2)
  66. shape.finish()
  67. shape.commit()
  68. parent_dir = os.path.dirname(save_path)
  69. if not os.path.exists(parent_dir):
  70. os.makedirs(parent_dir)
  71. if is_new_pdf:
  72. doc.save(save_path)
  73. else:
  74. doc.saveIncr()
  75. doc.close()
  76. def debug_show_bbox(raw_pdf_doc: fitz.Document, page_idx: int, bboxes: list, droped_bboxes:list, expect_drop_bboxes:list, save_path: str, expected_page_id:int):
  77. """
  78. 以覆盖的方式写个临时的pdf,用于debug
  79. """
  80. if page_idx!=expected_page_id:
  81. return
  82. if os.path.exists(save_path):
  83. # 删除已经存在的文件
  84. os.remove(save_path)
  85. # 创建一个新的空白 PDF 文件
  86. doc = fitz.open('')
  87. width = raw_pdf_doc[page_idx].rect.width
  88. height = raw_pdf_doc[page_idx].rect.height
  89. new_page = doc.new_page(width=width, height=height)
  90. shape = new_page.new_shape()
  91. for bbox in bboxes:
  92. # 原始box画上去
  93. rect = fitz.Rect(*bbox[0:4])
  94. shape = new_page.new_shape()
  95. shape.draw_rect(rect)
  96. shape.finish(color=fitz.pdfcolor['red'], fill=fitz.pdfcolor['blue'], fill_opacity=0.2)
  97. shape.finish()
  98. shape.commit()
  99. for bbox in droped_bboxes:
  100. # 原始box画上去
  101. rect = fitz.Rect(*bbox[0:4])
  102. shape = new_page.new_shape()
  103. shape.draw_rect(rect)
  104. shape.finish(color=None, fill=fitz.pdfcolor['yellow'], fill_opacity=0.2)
  105. shape.finish()
  106. shape.commit()
  107. for bbox in expect_drop_bboxes:
  108. # 原始box画上去
  109. rect = fitz.Rect(*bbox[0:4])
  110. shape = new_page.new_shape()
  111. shape.draw_rect(rect)
  112. shape.finish(color=fitz.pdfcolor['red'], fill=None)
  113. shape.finish()
  114. shape.commit()
  115. # shape.insert_textbox(fitz.Rect(200, 0, 600, 20), f"total bboxes: {len(bboxes)}", fontname="helv", fontsize=12,
  116. # color=(0, 0, 0))
  117. # shape.finish(color=fitz.pdfcolor['black'])
  118. # shape.commit()
  119. parent_dir = os.path.dirname(save_path)
  120. if not os.path.exists(parent_dir):
  121. os.makedirs(parent_dir)
  122. doc.save(save_path)
  123. doc.close()
  124. def debug_show_page(page, bboxes1: list,bboxes2: list,bboxes3: list,):
  125. save_path = "./tmp/debug.pdf"
  126. if os.path.exists(save_path):
  127. # 删除已经存在的文件
  128. os.remove(save_path)
  129. # 创建一个新的空白 PDF 文件
  130. doc = fitz.open('')
  131. width = page.rect.width
  132. height = page.rect.height
  133. new_page = doc.new_page(width=width, height=height)
  134. shape = new_page.new_shape()
  135. for bbox in bboxes1:
  136. # 原始box画上去
  137. rect = fitz.Rect(*bbox[0:4])
  138. shape = new_page.new_shape()
  139. shape.draw_rect(rect)
  140. shape.finish(color=fitz.pdfcolor['red'], fill=fitz.pdfcolor['blue'], fill_opacity=0.2)
  141. shape.finish()
  142. shape.commit()
  143. for bbox in bboxes2:
  144. # 原始box画上去
  145. rect = fitz.Rect(*bbox[0:4])
  146. shape = new_page.new_shape()
  147. shape.draw_rect(rect)
  148. shape.finish(color=None, fill=fitz.pdfcolor['yellow'], fill_opacity=0.2)
  149. shape.finish()
  150. shape.commit()
  151. for bbox in bboxes3:
  152. # 原始box画上去
  153. rect = fitz.Rect(*bbox[0:4])
  154. shape = new_page.new_shape()
  155. shape.draw_rect(rect)
  156. shape.finish(color=fitz.pdfcolor['red'], fill=None)
  157. shape.finish()
  158. shape.commit()
  159. parent_dir = os.path.dirname(save_path)
  160. if not os.path.exists(parent_dir):
  161. os.makedirs(parent_dir)
  162. doc.save(save_path)
  163. doc.close()
  164. def draw_layout_bbox_on_page(raw_pdf_doc: fitz.Document, paras_dict:dict, header, footer, pdf_path: str):
  165. """
  166. 在page上画出bbox,保存到save_path
  167. """
  168. # 检查文件是否存在
  169. is_new_pdf = False
  170. if os.path.exists(pdf_path):
  171. # 打开现有的 PDF 文件
  172. doc = fitz.open(pdf_path)
  173. else:
  174. # 创建一个新的空白 PDF 文件
  175. is_new_pdf = True
  176. doc = fitz.open('')
  177. for k, v in paras_dict.items():
  178. page_idx = v['page_idx']
  179. layouts = v['layout_bboxes']
  180. page = doc[page_idx]
  181. shape = page.new_shape()
  182. for order, layout in enumerate(layouts):
  183. border_offset = 1
  184. rect_box = layout['layout_bbox']
  185. layout_label = layout['layout_label']
  186. fill_color = fitz.pdfcolor['pink'] if layout_label=='U' else None
  187. rect_box = [rect_box[0]+1, rect_box[1]-border_offset, rect_box[2]-1, rect_box[3]+border_offset]
  188. rect = fitz.Rect(*rect_box)
  189. shape.draw_rect(rect)
  190. shape.finish(color=fitz.pdfcolor['red'], fill=fill_color, fill_opacity=0.4)
  191. """
  192. draw order text on layout box
  193. """
  194. font_size = 10
  195. shape.insert_text((rect_box[0] + 1, rect_box[1] + font_size), f"{order}", fontsize=font_size, color=(0, 0, 0))
  196. """画上footer header"""
  197. if header:
  198. shape.draw_rect(fitz.Rect(header))
  199. shape.finish(color=None, fill=fitz.pdfcolor['black'], fill_opacity=0.2)
  200. if footer:
  201. shape.draw_rect(fitz.Rect(footer))
  202. shape.finish(color=None, fill=fitz.pdfcolor['black'], fill_opacity=0.2)
  203. shape.commit()
  204. if is_new_pdf:
  205. doc.save(pdf_path)
  206. else:
  207. doc.saveIncr()
  208. doc.close()
  209. @DeprecationWarning
  210. def draw_layout_on_page(raw_pdf_doc: fitz.Document, page_idx: int, page_layout: list, pdf_path: str):
  211. """
  212. 把layout的box用红色边框花在pdf_path的page_idx上
  213. """
  214. def draw(shape, layout, fill_color=fitz.pdfcolor['pink']):
  215. border_offset = 1
  216. rect_box = layout['layout_bbox']
  217. layout_label = layout['layout_label']
  218. sub_layout = layout['sub_layout']
  219. if len(sub_layout)==0:
  220. fill_color = fill_color if layout_label=='U' else None
  221. rect_box = [rect_box[0]+1, rect_box[1]-border_offset, rect_box[2]-1, rect_box[3]+border_offset]
  222. rect = fitz.Rect(*rect_box)
  223. shape.draw_rect(rect)
  224. shape.finish(color=fitz.pdfcolor['red'], fill=fill_color, fill_opacity=0.2)
  225. # if layout_label=='U':
  226. # bad_boxes = layout.get("bad_boxes", [])
  227. # for bad_box in bad_boxes:
  228. # rect = fitz.Rect(*bad_box)
  229. # shape.draw_rect(rect)
  230. # shape.finish(color=fitz.pdfcolor['red'], fill=fitz.pdfcolor['red'], fill_opacity=0.2)
  231. # else:
  232. # rect = fitz.Rect(*rect_box)
  233. # shape.draw_rect(rect)
  234. # shape.finish(color=fitz.pdfcolor['blue'])
  235. for sub_layout in sub_layout:
  236. draw(shape, sub_layout)
  237. shape.commit()
  238. # 检查文件是否存在
  239. is_new_pdf = False
  240. if os.path.exists(pdf_path):
  241. # 打开现有的 PDF 文件
  242. doc = fitz.open(pdf_path)
  243. else:
  244. # 创建一个新的空白 PDF 文件
  245. is_new_pdf = True
  246. doc = fitz.open('')
  247. page = doc[page_idx]
  248. shape = page.new_shape()
  249. for order, layout in enumerate(page_layout):
  250. draw(shape, layout, fitz.pdfcolor['yellow'])
  251. # shape.insert_textbox(fitz.Rect(200, 0, 600, 20), f"total bboxes: {len(layout)}", fontname="helv", fontsize=12,
  252. # color=(0, 0, 0))
  253. # shape.finish(color=fitz.pdfcolor['black'])
  254. # shape.commit()
  255. parent_dir = os.path.dirname(pdf_path)
  256. if not os.path.exists(parent_dir):
  257. os.makedirs(parent_dir)
  258. if is_new_pdf:
  259. doc.save(pdf_path)
  260. else:
  261. doc.saveIncr()
  262. doc.close()