draw.py 5.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123
  1. import sys
  2. from libs.commons import fitz
  3. from para.commons import *
  4. if sys.version_info[0] >= 3:
  5. sys.stdout.reconfigure(encoding="utf-8") # type: ignore
  6. class DrawAnnos:
  7. """
  8. This class draws annotations on the pdf file
  9. ----------------------------------------
  10. Color Code
  11. ----------------------------------------
  12. Red: (1, 0, 0)
  13. Green: (0, 1, 0)
  14. Blue: (0, 0, 1)
  15. Yellow: (1, 1, 0) - mix of red and green
  16. Cyan: (0, 1, 1) - mix of green and blue
  17. Magenta: (1, 0, 1) - mix of red and blue
  18. White: (1, 1, 1) - red, green and blue full intensity
  19. Black: (0, 0, 0) - no color component whatsoever
  20. Gray: (0.5, 0.5, 0.5) - equal and medium intensity of red, green and blue color components
  21. Orange: (1, 0.65, 0) - maximum intensity of red, medium intensity of green, no blue component
  22. """
  23. def __init__(self) -> None:
  24. pass
  25. def __is_nested_list(self, lst):
  26. """
  27. This function returns True if the given list is a nested list of any degree.
  28. """
  29. if isinstance(lst, list):
  30. return any(self.__is_nested_list(i) for i in lst) or any(isinstance(i, list) for i in lst)
  31. return False
  32. def __valid_rect(self, bbox):
  33. # Ensure that the rectangle is not empty or invalid
  34. if isinstance(bbox[0], list):
  35. return False # It's a nested list, hence it can't be valid rect
  36. else:
  37. return bbox[0] < bbox[2] and bbox[1] < bbox[3]
  38. def __draw_nested_boxes(self, page, nested_bbox, color=(0, 1, 1)):
  39. """
  40. This function draws the nested boxes
  41. Parameters
  42. ----------
  43. page : fitz.Page
  44. page
  45. nested_bbox : list
  46. nested bbox
  47. color : tuple
  48. color, by default (0, 1, 1) # draw with cyan color for combined paragraph
  49. """
  50. if self.__is_nested_list(nested_bbox): # If it's a nested list
  51. for bbox in nested_bbox:
  52. self.__draw_nested_boxes(page, bbox, color) # Recursively call the function
  53. elif self.__valid_rect(nested_bbox): # If valid rectangle
  54. para_rect = fitz.Rect(nested_bbox)
  55. para_anno = page.add_rect_annot(para_rect)
  56. para_anno.set_colors(stroke=color) # draw with cyan color for combined paragraph
  57. para_anno.set_border(width=1)
  58. para_anno.update()
  59. def draw_annos(self, input_pdf_path, pdf_dic, output_pdf_path):
  60. pdf_doc = open_pdf(input_pdf_path)
  61. if pdf_dic is None:
  62. pdf_dic = {}
  63. if output_pdf_path is None:
  64. output_pdf_path = input_pdf_path.replace(".pdf", "_anno.pdf")
  65. for page_id, page in enumerate(pdf_doc): # type: ignore
  66. page_key = f"page_{page_id}"
  67. for ele_key, ele_data in pdf_dic[page_key].items():
  68. if ele_key == "para_blocks":
  69. para_blocks = ele_data
  70. for para_block in para_blocks:
  71. if "paras" in para_block.keys():
  72. paras = para_block["paras"]
  73. for para_key, para_content in paras.items():
  74. para_bbox = para_content["para_bbox"]
  75. # print(f"para_bbox: {para_bbox}")
  76. # print(f"is a nested list: {self.__is_nested_list(para_bbox)}")
  77. if self.__is_nested_list(para_bbox) and len(para_bbox) > 1:
  78. color = (0, 1, 1)
  79. self.__draw_nested_boxes(
  80. page, para_bbox, color
  81. ) # draw with cyan color for combined paragraph
  82. else:
  83. if self.__valid_rect(para_bbox):
  84. para_rect = fitz.Rect(para_bbox)
  85. para_anno = page.add_rect_annot(para_rect)
  86. para_anno.set_colors(stroke=(0, 1, 0)) # draw with green color for normal paragraph
  87. para_anno.set_border(width=0.5)
  88. para_anno.update()
  89. is_para_title = para_content["is_para_title"]
  90. if is_para_title:
  91. if self.__is_nested_list(para_content["para_bbox"]) and len(para_content["para_bbox"]) > 1:
  92. color = (0, 0, 1)
  93. self.__draw_nested_boxes(
  94. page, para_content["para_bbox"], color
  95. ) # draw with cyan color for combined title
  96. else:
  97. if self.__valid_rect(para_content["para_bbox"]):
  98. para_rect = fitz.Rect(para_content["para_bbox"])
  99. if self.__valid_rect(para_content["para_bbox"]):
  100. para_anno = page.add_rect_annot(para_rect)
  101. para_anno.set_colors(stroke=(0, 0, 1)) # draw with blue color for normal title
  102. para_anno.set_border(width=0.5)
  103. para_anno.update()
  104. pdf_doc.save(output_pdf_path)
  105. pdf_doc.close()