convert_image.py 5.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172
  1. import logging
  2. import os
  3. import sys
  4. sys.path.append(os.path.dirname(__file__) + "/../")
  5. from pdfminer.layout import LTLine
  6. import traceback
  7. import cv2
  8. from format_convert import get_memory_info
  9. from format_convert.utils import judge_error_code, add_div, LineTable, get_table_html
  10. from format_convert.table_correct import get_rotated_image
  11. from format_convert.convert_need_interface import from_otr_interface, from_ocr_interface
  12. def image_process(image_np, image_path, use_ocr=True):
  13. from format_convert.convert_tree import _Table, _Sentence
  14. logging.info("into image_preprocess")
  15. try:
  16. # 图片倾斜校正,写入原来的图片路径
  17. print("image_process", image_path)
  18. g_r_i = get_rotated_image(image_np, image_path)
  19. if g_r_i == [-1]:
  20. return [-1]
  21. # otr需要图片resize, 写入另一个路径
  22. image_np = cv2.imread(image_path)
  23. if image_np is None:
  24. return []
  25. best_h, best_w = get_best_predict_size(image_np)
  26. image_resize = cv2.resize(image_np, (best_w, best_h), interpolation=cv2.INTER_AREA)
  27. # image_resize_path = image_path[:-4] + "_resize" + image_path[-4:]
  28. image_resize_path = image_path.split(".")[0] + "_resize." + image_path.split(".")[-1]
  29. cv2.imwrite(image_resize_path, image_resize)
  30. # 调用otr模型接口
  31. with open(image_resize_path, "rb") as f:
  32. image_bytes = f.read()
  33. list_line = from_otr_interface(image_bytes)
  34. if judge_error_code(list_line):
  35. return list_line
  36. # 将resize后得到的bbox根据比例还原
  37. ratio = (image_np.shape[0]/best_h, image_np.shape[1]/best_w)
  38. for i in range(len(list_line)):
  39. point = list_line[i]
  40. list_line[i] = [int(point[0]*ratio[1]), int(point[1]*ratio[0]),
  41. int(point[2]*ratio[1]), int(point[3]*ratio[0])]
  42. # 调用ocr模型接口
  43. with open(image_path, "rb") as f:
  44. image_bytes = f.read()
  45. text_list, bbox_list = from_ocr_interface(image_bytes, True)
  46. if judge_error_code(text_list):
  47. return text_list
  48. # 调用现成方法形成表格
  49. try:
  50. from format_convert.convert_tree import TableLine
  51. list_lines = []
  52. for line in list_line:
  53. list_lines.append(LTLine(1, (line[0], line[1]), (line[2], line[3])))
  54. from format_convert.convert_tree import TextBox
  55. list_text_boxes = []
  56. for i in range(len(bbox_list)):
  57. bbox = bbox_list[i]
  58. b_text = text_list[i]
  59. list_text_boxes.append(TextBox([bbox[0][0], bbox[0][1],
  60. bbox[2][0], bbox[2][1]], b_text))
  61. lt = LineTable()
  62. tables, obj_in_table, _ = lt.recognize_table(list_text_boxes, list_lines, False)
  63. obj_list = []
  64. for table in tables:
  65. obj_list.append(_Table(table["table"], table["bbox"]))
  66. for text_box in list_text_boxes:
  67. if text_box not in obj_in_table:
  68. obj_list.append(_Sentence(text_box.get_text(), text_box.bbox))
  69. return obj_list
  70. except:
  71. traceback.print_exc()
  72. return [-8]
  73. except Exception as e:
  74. logging.info("image_preprocess error")
  75. print("image_preprocess", traceback.print_exc())
  76. return [-1]
  77. @get_memory_info.memory_decorator
  78. def picture2text(path, html=False):
  79. logging.info("into picture2text")
  80. try:
  81. # 判断图片中表格
  82. img = cv2.imread(path)
  83. if img is None:
  84. return [-3]
  85. text = image_process(img, path)
  86. if judge_error_code(text):
  87. return text
  88. if html:
  89. text = add_div(text)
  90. return [text]
  91. except Exception as e:
  92. logging.info("picture2text error!")
  93. print("picture2text", traceback.print_exc())
  94. return [-1]
  95. def get_best_predict_size(image_np, times=64):
  96. sizes = []
  97. for i in range(1, 100):
  98. if i*times <= 3000:
  99. sizes.append(i*times)
  100. sizes.sort(key=lambda x: x, reverse=True)
  101. min_len = 10000
  102. best_height = sizes[0]
  103. for height in sizes:
  104. if abs(image_np.shape[0] - height) < min_len:
  105. min_len = abs(image_np.shape[0] - height)
  106. best_height = height
  107. min_len = 10000
  108. best_width = sizes[0]
  109. for width in sizes:
  110. if abs(image_np.shape[1] - width) < min_len:
  111. min_len = abs(image_np.shape[1] - width)
  112. best_width = width
  113. return best_height, best_width
  114. class ImageConvert:
  115. def __init__(self, path, unique_type_dir):
  116. from format_convert.convert_tree import _Document
  117. self._doc = _Document(path)
  118. self.path = path
  119. self.unique_type_dir = unique_type_dir
  120. def init_package(self):
  121. # 各个包初始化
  122. try:
  123. with open(self.path, "rb") as f:
  124. self.image = f.read()
  125. except:
  126. logging.info("cannot open image!")
  127. traceback.print_exc()
  128. self._doc.error_code = [-3]
  129. def convert(self):
  130. from format_convert.convert_tree import _Page, _Image
  131. self.init_package()
  132. if self._doc.error_code is not None:
  133. return
  134. _page = _Page(None, 0)
  135. _image = _Image(self.image, self.path)
  136. _page.add_child(_image)
  137. self._doc.add_child(_page)
  138. def get_html(self):
  139. try:
  140. self.convert()
  141. except:
  142. traceback.print_exc()
  143. self._doc.error_code = [-1]
  144. if self._doc.error_code is not None:
  145. return self._doc.error_code
  146. return self._doc.get_html()