convert_tree.py 4.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183
  1. import io
  2. import cv2
  3. from PIL import Image
  4. import numpy as np
  5. from format_convert.convert_image import image_process
  6. from format_convert.utils import add_div, judge_error_code, get_table_html, sort_object
  7. class _Document:
  8. def __init__(self, doc_path):
  9. self.doc_path = doc_path
  10. # Document's child -> Page
  11. self.children = []
  12. self.error_code = None
  13. def add_child(self, child):
  14. if child.error_code is None:
  15. self.children.append(child)
  16. else:
  17. self.error_code = child.error_code
  18. def get_html(self):
  19. if self.error_code is not None:
  20. return self.error_code
  21. html_text = ""
  22. for child in self.children:
  23. # 先调用get_html才能更新error_code
  24. child_html_text = child.get_html()
  25. if child.error_code is not None:
  26. self.error_code = child.error_code
  27. return self.error_code
  28. else:
  29. html_text += child_html_text
  30. return [html_text]
  31. class _Page:
  32. def __init__(self, page, page_no):
  33. self.page = page
  34. self.page_no = page_no
  35. # Page's child -> Image, Table, Sentence
  36. self.children = []
  37. self.error_code = None
  38. # pdf对象需反向排序
  39. self.is_reverse = False
  40. # objs in tables
  41. self.in_table_objs = set()
  42. def add_child(self, child):
  43. if child.error_code is None:
  44. self.children.append(child)
  45. else:
  46. self.error_code = child.error_code
  47. def get_html(self):
  48. if self.error_code is not None:
  49. return ""
  50. html_text = ""
  51. self.children = sort_object(self.children, self.is_reverse)
  52. for child in self.children:
  53. # 先调用get_html才能更新error_code
  54. child_html_text = child.get_html()
  55. if child.error_code is not None:
  56. self.error_code = child.error_code
  57. return ""
  58. else:
  59. html_text += child_html_text
  60. return html_text
  61. class _Image:
  62. def __init__(self, content, path):
  63. self.content = content
  64. self.path = path
  65. # 来源
  66. self.is_from_pdf = False
  67. # 位置
  68. self.x = 0
  69. self.y = 0
  70. # 识别结果
  71. self.otr_result = None
  72. self.ocr_result = None
  73. # Image's child -> Table, Sentence
  74. self.children = []
  75. self.error_code = None
  76. # objs in tables
  77. self.in_table_objs = set()
  78. def add_child(self, child):
  79. if child.error_code is None:
  80. self.children.append(child)
  81. else:
  82. self.error_code = child.error_code
  83. def get_html(self):
  84. # 将Image转为Sentence,table
  85. self.convert()
  86. if self.error_code is not None:
  87. return ""
  88. html_text = ""
  89. self.children = sort_object(self.children)
  90. for child in self.children:
  91. # 先调用get_html才能更新error_code
  92. child_html_text = child.get_html()
  93. if child.error_code is not None:
  94. self.error_code = child.error_code
  95. return ""
  96. else:
  97. html_text += child_html_text
  98. return html_text
  99. def get_text(self):
  100. return
  101. def convert(self):
  102. # 二进制转numpy
  103. image_np = Image.open(io.BytesIO(self.content))
  104. image_np = cv2.cvtColor(np.asarray(image_np), cv2.COLOR_RGB2BGR)
  105. obj_list = image_process(image_np, self.path, use_ocr=True)
  106. if judge_error_code(obj_list):
  107. self.error_code = obj_list
  108. return
  109. for obj in obj_list:
  110. self.add_child(obj)
  111. class _Table:
  112. def __init__(self, content, bbox, is_html=False):
  113. self.content = content
  114. self.is_html = is_html
  115. self.bbox = bbox
  116. self.x = bbox[0]
  117. self.y = bbox[1]
  118. self.shape = (len(content), len(content[0]))
  119. self.error_code = None
  120. def get_html(self):
  121. if self.error_code is not None:
  122. return ""
  123. if self.is_html:
  124. return self.content
  125. else:
  126. # 将二维数组转为html table
  127. html_text = get_table_html(self.content)
  128. return html_text
  129. class _Sentence:
  130. def __init__(self, content, bbox, is_html=False):
  131. self.content = content
  132. self.is_html = is_html
  133. # 位置
  134. self.bbox = bbox
  135. self.x = bbox[0]
  136. self.y = bbox[1]
  137. self.error_code = None
  138. def get_html(self):
  139. if self.error_code is not None:
  140. return ""
  141. print("_Sentence", self.content, self.bbox)
  142. if self.is_html:
  143. return self.content
  144. else:
  145. return add_div(self.content)
  146. class TextBox:
  147. def __init__(self, bbox, text):
  148. self.bbox = bbox
  149. self.text = text
  150. def get_text(self):
  151. return self.text
  152. class TableLine:
  153. def __init__(self, bbox):
  154. self.bbox = bbox