convert_image.py 32 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888
  1. # encoding=utf8
  2. import copy
  3. import inspect
  4. import io
  5. import logging
  6. import os
  7. import sys
  8. import time
  9. import requests
  10. import numpy as np
  11. from PIL import Image
  12. sys.path.append(os.path.dirname(__file__) + "/../")
  13. from pdfminer.layout import LTLine
  14. import traceback
  15. import cv2
  16. from isr.pre_process import count_red_pixel
  17. from format_convert.utils import judge_error_code, add_div, LineTable, get_table_html, get_logger, log, \
  18. memory_decorator, pil_resize, np2bytes, ocr_cant_read
  19. from format_convert.convert_need_interface import from_otr_interface, from_ocr_interface, from_gpu_interface_redis, \
  20. from_idc_interface, from_isr_interface
  21. from format_convert.table_correct import get_rotated_image
  22. def image_process(image_np, image_path, is_from_pdf=False, is_from_docx=False, use_ocr=True):
  23. from format_convert.convert_tree import _Table, _Sentence
  24. def get_cluster(t_list, b_list, axis):
  25. zip_list = list(zip(t_list, b_list))
  26. if len(zip_list) == 0:
  27. return t_list, b_list
  28. if len(zip_list[0]) > 0:
  29. zip_list.sort(key=lambda x: x[1][axis][1])
  30. cluster_list = []
  31. margin = 5
  32. for text, bbox in zip_list:
  33. _find = 0
  34. for cluster in cluster_list:
  35. if abs(cluster[1] - bbox[axis][1]) <= margin:
  36. cluster[0].append([text, bbox])
  37. cluster[1] = bbox[axis][1]
  38. _find = 1
  39. break
  40. if not _find:
  41. cluster_list.append([[[text, bbox]], bbox[axis][1]])
  42. new_text_list = []
  43. new_bbox_list = []
  44. for cluster in cluster_list:
  45. # print("=============convert_image")
  46. # print("cluster_list", cluster)
  47. center_y = 0
  48. for text, bbox in cluster[0]:
  49. center_y += bbox[axis][1]
  50. center_y = int(center_y / len(cluster[0]))
  51. for text, bbox in cluster[0]:
  52. bbox[axis][1] = center_y
  53. new_text_list.append(text)
  54. new_bbox_list.append(bbox)
  55. # print("cluster_list", cluster)
  56. return new_text_list, new_bbox_list
  57. def merge_textbox(textbox_list, in_objs):
  58. delete_obj = []
  59. threshold = 5
  60. textbox_list.sort(key=lambda x:x.bbox[0])
  61. for k in range(len(textbox_list)):
  62. tb1 = textbox_list[k]
  63. if tb1 not in in_objs and tb1 not in delete_obj:
  64. for m in range(k+1, len(textbox_list)):
  65. tb2 = textbox_list[m]
  66. if tb2 in in_objs:
  67. continue
  68. if abs(tb1.bbox[1]-tb2.bbox[1]) <= threshold \
  69. and abs(tb1.bbox[3]-tb2.bbox[3]) <= threshold:
  70. if tb1.bbox[0] <= tb2.bbox[0]:
  71. tb1.text = tb1.text + tb2.text
  72. else:
  73. tb1.text = tb2.text + tb1.text
  74. tb1.bbox[0] = min(tb1.bbox[0], tb2.bbox[0])
  75. tb1.bbox[2] = max(tb1.bbox[2], tb2.bbox[2])
  76. delete_obj.append(tb2)
  77. for _obj in delete_obj:
  78. if _obj in textbox_list:
  79. textbox_list.remove(_obj)
  80. return textbox_list
  81. def idc_process(_image_np):
  82. # 图片倾斜校正,写入原来的图片路径
  83. # print("image_process", image_path)
  84. # g_r_i = get_rotated_image(_image_np, image_path)
  85. # if judge_error_code(g_r_i):
  86. # if is_from_docx:
  87. # return []
  88. # else:
  89. # return g_r_i
  90. # _image_np = cv2.imread(image_path)
  91. # if _image_np is None:
  92. # return []
  93. # return _image_np
  94. # if _image_np is None:
  95. # return []
  96. # idc模型实现图片倾斜校正
  97. h, w = get_best_predict_size2(_image_np, 1080)
  98. image_resize = pil_resize(_image_np, h, w)
  99. # image_resize_path = image_path.split(".")[0] + "_resize_idc." + image_path.split(".")[-1]
  100. # cv2.imwrite(image_resize_path, image_resize)
  101. # with open(image_resize_path, "rb") as f:
  102. # image_bytes = f.read()
  103. image_bytes = np2bytes(image_resize)
  104. angle = from_idc_interface(image_bytes)
  105. if judge_error_code(angle):
  106. if is_from_docx:
  107. return []
  108. else:
  109. return angle
  110. # 根据角度旋转
  111. image_pil = Image.fromarray(_image_np)
  112. _image_np = np.array(image_pil.rotate(angle, expand=1))
  113. # 写入
  114. # idc_path = image_path.split(".")[0] + "_idc." + image_path.split(".")[-1]
  115. # cv2.imwrite(idc_path, image_np)
  116. return _image_np
  117. def isr_process(_image_np):
  118. log("isr_process image shape " + str(_image_np.shape))
  119. image_np_copy = copy.deepcopy(_image_np)
  120. # isr模型去除印章
  121. _isr_time = time.time()
  122. if count_red_pixel(_image_np):
  123. # 红色像素达到一定值才过模型
  124. image_bytes = np2bytes(_image_np)
  125. _image_np = from_isr_interface(image_bytes)
  126. if judge_error_code(_image_np):
  127. if is_from_docx:
  128. return []
  129. else:
  130. return _image_np
  131. # [1]代表检测不到印章,直接返回
  132. if isinstance(_image_np, list) and _image_np == [1]:
  133. log("no seals detected!")
  134. _image_np = image_np_copy
  135. log("isr total time "+str(time.time()-_isr_time))
  136. return _image_np
  137. def ocr_process(_image_np, _threshold=2048):
  138. log("ocr_process image shape " + str(_image_np.shape))
  139. # ocr图片过大内存溢出,需resize
  140. # 大图按比例缩小,小图维持不变;若统一拉伸成固定大小如1024会爆显存
  141. ratio = (1, 1)
  142. if _image_np.shape[0] > _threshold or _image_np.shape[1] > _threshold:
  143. best_h, best_w = get_best_predict_size2(_image_np, _threshold)
  144. _image_np = pil_resize(_image_np, best_h, best_w)
  145. log("ocr_process image resize " + str(_image_np.shape))
  146. ratio = (image_np.shape[0]/best_h, image_np.shape[1]/best_w)
  147. # 大图片ocr加锁,防止爆显存
  148. # if _image_np.shape[0] >= 1024 and _image_np.shape[1] >= 1024:
  149. # file_lock = True
  150. # else:
  151. # file_lock = False
  152. # 调用ocr模型接口
  153. image_bytes = np2bytes(_image_np)
  154. text_list, bbox_list = from_ocr_interface(image_bytes, is_table=True)
  155. if judge_error_code(text_list):
  156. return text_list, text_list
  157. for i in range(len(bbox_list)):
  158. point = bbox_list[i]
  159. bbox_list[i] = [[int(point[0][0]*ratio[0]), int(point[0][1]*ratio[1])],
  160. [int(point[1][0]*ratio[0]), int(point[1][1]*ratio[1])],
  161. [int(point[2][0]*ratio[0]), int(point[2][1]*ratio[1])],
  162. [int(point[3][0]*ratio[0]), int(point[3][1]*ratio[1])]]
  163. return text_list, bbox_list
  164. def otr_process(_image_np):
  165. log("otr_process image shape " + str(_image_np.shape))
  166. # otr模型识别表格,需要图片resize成模型所需大小, 写入另一个路径
  167. best_h, best_w = get_best_predict_size(_image_np)
  168. image_resize = pil_resize(_image_np, best_h, best_w)
  169. # image_resize_path = image_path.split(".")[0] + "_resize_otr." + image_path.split(".")[-1]
  170. # cv2.imwrite(image_resize_path, image_resize)
  171. # 调用otr模型接口
  172. # with open(image_resize_path, "rb") as f:
  173. # image_bytes = f.read()
  174. image_bytes = np2bytes(image_resize)
  175. list_line = from_otr_interface(image_bytes, is_from_pdf)
  176. if judge_error_code(list_line):
  177. if is_from_docx:
  178. return []
  179. else:
  180. return list_line
  181. # otr resize后得到的bbox根据比例还原
  182. start_time = time.time()
  183. ratio = (_image_np.shape[0]/best_h, _image_np.shape[1]/best_w)
  184. for i in range(len(list_line)):
  185. point = list_line[i]
  186. list_line[i] = [int(point[0]*ratio[1]), int(point[1]*ratio[0]),
  187. int(point[2]*ratio[1]), int(point[3]*ratio[0])]
  188. log("otr resize bbox recover " + str(time.time()-start_time))
  189. return list_line
  190. def table_process(list_line, text_list, bbox_list):
  191. # 调用现成方法形成表格
  192. try:
  193. from format_convert.convert_tree import TableLine
  194. list_lines = []
  195. for line in list_line:
  196. list_lines.append(LTLine(1, (line[0], line[1]), (line[2], line[3])))
  197. from format_convert.convert_tree import TextBox
  198. list_text_boxes = []
  199. for i in range(len(bbox_list)):
  200. bbox = bbox_list[i]
  201. b_text = text_list[i]
  202. list_text_boxes.append(TextBox([bbox[0][0], bbox[0][1],
  203. bbox[2][0], bbox[2][1]], b_text))
  204. # for _textbox in list_text_boxes:
  205. # print("==",_textbox.get_text())
  206. lt = LineTable()
  207. # print('text_list', text_list)
  208. # print('bbox_list', bbox_list)
  209. # print('list_line', list_line)
  210. tables, obj_in_table, _ = lt.recognize_table(list_text_boxes, list_lines, False)
  211. # 合并同一行textbox
  212. list_text_boxes = merge_textbox(list_text_boxes, obj_in_table)
  213. return list_text_boxes, tables, obj_in_table
  214. except:
  215. traceback.print_exc()
  216. return [-8], [-8], [-8]
  217. log("into image_preprocess")
  218. try:
  219. if image_np is None:
  220. return []
  221. if image_np.shape[0] <= 20 or image_np.shape[1] <= 20:
  222. return []
  223. # 判断是否需要长图分割
  224. slice_flag = need_image_slice(image_np)
  225. log("need_image_slice " + str(slice_flag) + " " + str(image_np.shape))
  226. idc_flag = False
  227. image_np_list = [image_np]
  228. if slice_flag:
  229. # 方向分类
  230. image_np = idc_process(image_np)
  231. idc_flag = True
  232. if isinstance(image_np, list):
  233. return image_np
  234. # 再判断
  235. if need_image_slice(image_np):
  236. # 长图分割
  237. image_np_list = image_slice_new(image_np)
  238. if len(image_np_list) < 1:
  239. log("image_slice failed!")
  240. image_np_list = [image_np]
  241. # return [-10]
  242. all_obj_list = []
  243. _add_y = 0
  244. for image_np in image_np_list:
  245. # print("sub image shape", image_np.shape)
  246. # 整体分辨率限制
  247. threshold = 2048
  248. if image_np.shape[0] > threshold or image_np.shape[1] > threshold:
  249. h, w = get_best_predict_size2(image_np, threshold=threshold)
  250. log("global image resize " + str(image_np.shape[:2]) + " -> " + str(h) + "," + str(w))
  251. image_np = pil_resize(image_np, h, w)
  252. # 印章去除
  253. image_np = isr_process(image_np)
  254. if isinstance(image_np, list):
  255. return image_np
  256. # 文字识别
  257. text_list, box_list = ocr_process(image_np)
  258. if judge_error_code(text_list):
  259. return text_list
  260. # 判断ocr识别是否正确
  261. if ocr_cant_read(text_list, box_list) and not idc_flag:
  262. # 方向分类
  263. image_np = idc_process(image_np)
  264. # cv2.imshow("idc_process", image_np)
  265. # cv2.waitKey(0)
  266. if isinstance(image_np, list):
  267. return image_np
  268. # 文字识别
  269. text_list1, box_list_1 = ocr_process(image_np)
  270. if judge_error_code(text_list1):
  271. return text_list1
  272. # 比较字数
  273. # print("ocr process", len("".join(text_list)), len("".join(text_list1)))
  274. if len("".join(text_list)) < len("".join(text_list1)):
  275. text_list = text_list1
  276. box_list = box_list_1
  277. # 表格识别
  278. line_list = otr_process(image_np)
  279. if judge_error_code(line_list):
  280. return line_list
  281. # 表格生成
  282. text_box_list, table_list, obj_in_table_list = table_process(line_list, text_list, box_list)
  283. if judge_error_code(table_list):
  284. return table_list
  285. # 对象生成
  286. obj_list = []
  287. for table in table_list:
  288. obj_list.append(_Table(table["table"], table["bbox"]))
  289. for text_box in text_box_list:
  290. if text_box not in obj_in_table_list:
  291. obj_list.append(_Sentence(text_box.get_text(), text_box.bbox))
  292. # 修正y
  293. if len(image_np_list) > 1:
  294. list_y = []
  295. for obj in obj_list:
  296. obj.y += _add_y
  297. list_y.append(obj.y)
  298. if len(list_y) > 0:
  299. _add_y = max(list_y)
  300. # 合并
  301. all_obj_list += obj_list
  302. return all_obj_list
  303. except Exception as e:
  304. log("image_preprocess error")
  305. traceback.print_exc()
  306. return [-1]
  307. @memory_decorator
  308. def picture2text(path, html=False):
  309. log("into picture2text")
  310. try:
  311. # 判断图片中表格
  312. img = cv2.imread(path)
  313. if img is None:
  314. return [-3]
  315. text = image_process(img, path)
  316. if judge_error_code(text):
  317. return text
  318. if html:
  319. text = add_div(text)
  320. return [text]
  321. except Exception as e:
  322. log("picture2text error!")
  323. print("picture2text", traceback.print_exc())
  324. return [-1]
  325. def get_best_predict_size(image_np, times=64):
  326. sizes = []
  327. for i in range(1, 100):
  328. if i*times <= 1300:
  329. sizes.append(i*times)
  330. sizes.sort(key=lambda x: x, reverse=True)
  331. min_len = 10000
  332. best_height = sizes[0]
  333. for height in sizes:
  334. if abs(image_np.shape[0] - height) < min_len:
  335. min_len = abs(image_np.shape[0] - height)
  336. best_height = height
  337. min_len = 10000
  338. best_width = sizes[0]
  339. for width in sizes:
  340. if abs(image_np.shape[1] - width) < min_len:
  341. min_len = abs(image_np.shape[1] - width)
  342. best_width = width
  343. return best_height, best_width
  344. def get_best_predict_size2(image_np, threshold=3000):
  345. h, w = image_np.shape[:2]
  346. scale = threshold / max(h, w)
  347. h = int(h * scale)
  348. w = int(w * scale)
  349. return h, w
  350. def image_slice(image_np):
  351. """
  352. slice the image if the height is to large
  353. :return:
  354. """
  355. _sum = np.average(image_np, axis=1)
  356. list_white_line = []
  357. list_ave = list(_sum)
  358. for _i in range(len(list_ave)):
  359. if (list_ave[_i] > 250).all():
  360. list_white_line.append(_i)
  361. set_white_line = set(list_white_line)
  362. width = image_np.shape[1]
  363. height = image_np.shape[0]
  364. list_images = []
  365. _begin = 0
  366. _end = 0
  367. while 1:
  368. if _end > height:
  369. break
  370. _end += width
  371. while 1:
  372. if _begin in set_white_line:
  373. break
  374. if _begin > height:
  375. break
  376. _begin += 1
  377. _image = image_np[_begin:_end, ...]
  378. list_images.append(_image)
  379. _begin = _end
  380. log("image_slice into %d parts" % (len(list_images)))
  381. return list_images
  382. def image_slice_new(image_np):
  383. """
  384. 长图分割
  385. :return:
  386. """
  387. height, width = image_np.shape[:2]
  388. image_origin = copy.deepcopy(image_np)
  389. # 去除黑边
  390. image_np = remove_black_border(image_np)
  391. # 1. 转化成灰度图
  392. image_np = cv2.cvtColor(image_np, cv2.COLOR_BGR2GRAY)
  393. # 2. 二值化
  394. ret, binary = cv2.threshold(image_np, 125, 255, cv2.THRESH_BINARY_INV)
  395. # 3. 膨胀和腐蚀操作的核函数
  396. kernal = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
  397. # 4. 膨胀一次,让轮廓突出
  398. dilation = cv2.dilate(binary, kernal, iterations=1)
  399. # dilation = np.add(np.int0(np.full(dilation.shape, 255)), -1 * np.int0(dilation))
  400. # dilation = np.uint8(dilation)
  401. # cv2.namedWindow("dilation", 0)
  402. # cv2.resizeWindow("dilation", 1000, 800)
  403. # cv2.imshow("dilation", dilation)
  404. # cv2.waitKey(0)
  405. # cv2.imwrite("error.jpg", dilation)
  406. # 按行求平均
  407. width_avg = np.average(np.float32(dilation), axis=1)
  408. zero_index = np.where(width_avg == 0.)[0]
  409. # print(height, width)
  410. # print(width_avg)
  411. # print(width_avg.shape)
  412. # print(zero_index)
  413. # print(zero_index.shape)
  414. # zero_index.sort(key=lambda x: x)
  415. # 截取范围内寻找分割点
  416. max_distance = int(width / 2)
  417. image_list = []
  418. last_h = 0
  419. for i in range(height // width + 1):
  420. h = last_h + width
  421. # 前后的分割点
  422. zero_h_after = zero_index[np.where(zero_index >= h)]
  423. zero_h_before = zero_index[np.where(zero_index <= h)]
  424. # print("last_h, h", last_h, h)
  425. # print("last_h, h", last_h, h)
  426. # print(zero_index.shape)
  427. # print("zero_h_after.shape", zero_h_after.shape)
  428. if zero_h_after.shape[0] == 0:
  429. # 最后一截
  430. last_image = image_origin[last_h:, :, :]
  431. if last_image.shape[0] <= max_distance:
  432. image_list[-1] = np.concatenate([image_list[-1], last_image], axis=0)
  433. else:
  434. image_list.append(last_image)
  435. break
  436. # 分割点距离不能太远
  437. cut_h = zero_h_after.tolist()[0]
  438. print("cut_h", cut_h)
  439. if abs(h - cut_h) <= max_distance:
  440. image_list.append(image_origin[last_h:cut_h, :, :])
  441. last_h = cut_h
  442. # 后面找不到往前找
  443. else:
  444. cut_h = zero_h_before.tolist()[-1]
  445. if abs(cut_h - h) <= max_distance:
  446. image_list.append(image_origin[last_h:cut_h, :, :])
  447. last_h = cut_h
  448. # i = 0
  449. # for im in image_list:
  450. # print(im.shape)
  451. # cv2.imwrite("error" + str(i) + ".jpg", im)
  452. # i += 1
  453. # cv2.namedWindow("im", 0)
  454. # cv2.resizeWindow("im", 1000, 800)
  455. # cv2.imshow("im", im)
  456. # cv2.waitKey(0)
  457. log("image_slice into %d parts" % (len(image_list)))
  458. return image_list
  459. def need_image_slice(image_np):
  460. h, w = image_np.shape[:2]
  461. # if h > 3000 and w < 2000:
  462. # return True
  463. if 2. <= h / w and w >= 100:
  464. return True
  465. return False
  466. def remove_black_border(img_np):
  467. try:
  468. # 阈值
  469. threshold = 100
  470. # 转换为灰度图像
  471. gray = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY)
  472. # 获取图片尺寸
  473. h, w = gray.shape[:2]
  474. # 无法区分黑色区域超过一半的情况
  475. rowc = gray[:, int(1/2*w)]
  476. colc = gray[int(1/2*h), :]
  477. rowflag = np.argwhere(rowc > threshold)
  478. colflag = np.argwhere(colc > threshold)
  479. left, bottom, right, top = rowflag[0, 0], colflag[-1, 0], rowflag[-1, 0], colflag[0, 0]
  480. # cv2.imshow('remove_black_border', img_np[left:right, top:bottom, :])
  481. # cv2.waitKey()
  482. return img_np[left:right, top:bottom, :]
  483. except:
  484. return img_np
  485. class ImageConvert:
  486. def __init__(self, path, unique_type_dir):
  487. from format_convert.convert_tree import _Document
  488. self._doc = _Document(path)
  489. self.path = path
  490. self.unique_type_dir = unique_type_dir
  491. def init_package(self):
  492. # 各个包初始化
  493. try:
  494. with open(self.path, "rb") as f:
  495. self.image = f.read()
  496. except:
  497. log("cannot open image!")
  498. traceback.print_exc()
  499. self._doc.error_code = [-3]
  500. def convert(self):
  501. from format_convert.convert_tree import _Page, _Image
  502. self.init_package()
  503. if self._doc.error_code is not None:
  504. return
  505. _page = _Page(None, 0)
  506. _image = _Image(self.image, self.path)
  507. _page.add_child(_image)
  508. self._doc.add_child(_page)
  509. def get_html(self):
  510. try:
  511. self.convert()
  512. except:
  513. traceback.print_exc()
  514. self._doc.error_code = [-1]
  515. if self._doc.error_code is not None:
  516. return self._doc.error_code
  517. return self._doc.get_html()
  518. def image_process_old(image_np, image_path, is_from_pdf=False, is_from_docx=False, use_ocr=True):
  519. from format_convert.convert_tree import _Table, _Sentence
  520. def get_cluster(t_list, b_list, axis):
  521. zip_list = list(zip(t_list, b_list))
  522. if len(zip_list) == 0:
  523. return t_list, b_list
  524. if len(zip_list[0]) > 0:
  525. zip_list.sort(key=lambda x: x[1][axis][1])
  526. cluster_list = []
  527. margin = 5
  528. for text, bbox in zip_list:
  529. _find = 0
  530. for cluster in cluster_list:
  531. if abs(cluster[1] - bbox[axis][1]) <= margin:
  532. cluster[0].append([text, bbox])
  533. cluster[1] = bbox[axis][1]
  534. _find = 1
  535. break
  536. if not _find:
  537. cluster_list.append([[[text, bbox]], bbox[axis][1]])
  538. new_text_list = []
  539. new_bbox_list = []
  540. for cluster in cluster_list:
  541. # print("=============convert_image")
  542. # print("cluster_list", cluster)
  543. center_y = 0
  544. for text, bbox in cluster[0]:
  545. center_y += bbox[axis][1]
  546. center_y = int(center_y / len(cluster[0]))
  547. for text, bbox in cluster[0]:
  548. bbox[axis][1] = center_y
  549. new_text_list.append(text)
  550. new_bbox_list.append(bbox)
  551. # print("cluster_list", cluster)
  552. return new_text_list, new_bbox_list
  553. def merge_textbox(textbox_list, in_objs):
  554. delete_obj = []
  555. threshold = 5
  556. textbox_list.sort(key=lambda x:x.bbox[0])
  557. for k in range(len(textbox_list)):
  558. tb1 = textbox_list[k]
  559. if tb1 not in in_objs and tb1 not in delete_obj:
  560. for m in range(k+1, len(textbox_list)):
  561. tb2 = textbox_list[m]
  562. if tb2 in in_objs:
  563. continue
  564. if abs(tb1.bbox[1]-tb2.bbox[1]) <= threshold \
  565. and abs(tb1.bbox[3]-tb2.bbox[3]) <= threshold:
  566. if tb1.bbox[0] <= tb2.bbox[0]:
  567. tb1.text = tb1.text + tb2.text
  568. else:
  569. tb1.text = tb2.text + tb1.text
  570. tb1.bbox[0] = min(tb1.bbox[0], tb2.bbox[0])
  571. tb1.bbox[2] = max(tb1.bbox[2], tb2.bbox[2])
  572. delete_obj.append(tb2)
  573. for _obj in delete_obj:
  574. if _obj in textbox_list:
  575. textbox_list.remove(_obj)
  576. return textbox_list
  577. log("into image_preprocess")
  578. try:
  579. if image_np is None:
  580. return []
  581. # 整体分辨率限制
  582. if image_np.shape[0] > 2000 or image_np.shape[1] > 2000:
  583. h, w = get_best_predict_size2(image_np, threshold=2000)
  584. log("global image resize " + str(image_np.shape[:2]) + " -> " + str(h) + "," + str(w))
  585. image_np = pil_resize(image_np, h, w)
  586. # 图片倾斜校正,写入原来的图片路径
  587. # print("image_process", image_path)
  588. g_r_i = get_rotated_image(image_np, image_path)
  589. if judge_error_code(g_r_i):
  590. if is_from_docx:
  591. return []
  592. else:
  593. return g_r_i
  594. image_np = cv2.imread(image_path)
  595. image_np_copy = copy.deepcopy(image_np)
  596. if image_np is None:
  597. return []
  598. # if image_np is None:
  599. # return []
  600. #
  601. # # idc模型实现图片倾斜校正
  602. # image_resize = pil_resize(image_np, 640, 640)
  603. # image_resize_path = image_path.split(".")[0] + "_resize_idc." + image_path.split(".")[-1]
  604. # cv2.imwrite(image_resize_path, image_resize)
  605. #
  606. # with open(image_resize_path, "rb") as f:
  607. # image_bytes = f.read()
  608. # angle = from_idc_interface(image_bytes)
  609. # if judge_error_code(angle):
  610. # if is_from_docx:
  611. # return []
  612. # else:
  613. # return angle
  614. # # 根据角度旋转
  615. # image_pil = Image.fromarray(image_np)
  616. # image_np = np.array(image_pil.rotate(angle, expand=1))
  617. # # 写入
  618. # idc_path = image_path.split(".")[0] + "_idc." + image_path.split(".")[-1]
  619. # cv2.imwrite(idc_path, image_np)
  620. # isr模型去除印章
  621. _isr_time = time.time()
  622. if count_red_pixel(image_np):
  623. # 红色像素达到一定值才过模型
  624. with open(image_path, "rb") as f:
  625. image_bytes = f.read()
  626. image_np = from_isr_interface(image_bytes)
  627. if judge_error_code(image_np):
  628. if is_from_docx:
  629. return []
  630. else:
  631. return image_np
  632. # [1]代表检测不到印章,直接返回
  633. if isinstance(image_np, list) and image_np == [1]:
  634. log("no seals detected!")
  635. image_np = image_np_copy
  636. else:
  637. isr_path = image_path.split(".")[0] + "_isr." + image_path.split(".")[-1]
  638. cv2.imwrite(isr_path, image_np)
  639. log("isr total time "+str(time.time()-_isr_time))
  640. # otr模型识别表格,需要图片resize成模型所需大小, 写入另一个路径
  641. best_h, best_w = get_best_predict_size(image_np)
  642. # image_resize = cv2.resize(image_np, (best_w, best_h), interpolation=cv2.INTER_AREA)
  643. image_resize = pil_resize(image_np, best_h, best_w)
  644. image_resize_path = image_path.split(".")[0] + "_resize_otr." + image_path.split(".")[-1]
  645. cv2.imwrite(image_resize_path, image_resize)
  646. # 调用otr模型接口
  647. with open(image_resize_path, "rb") as f:
  648. image_bytes = f.read()
  649. list_line = from_otr_interface(image_bytes, is_from_pdf)
  650. if judge_error_code(list_line):
  651. return list_line
  652. # # 预处理
  653. # if is_from_pdf:
  654. # prob = 0.2
  655. # else:
  656. # prob = 0.5
  657. # with open(image_resize_path, "rb") as f:
  658. # image_bytes = f.read()
  659. # img_new, inputs = table_preprocess(image_bytes, prob)
  660. # if type(img_new) is list and judge_error_code(img_new):
  661. # return img_new
  662. # log("img_new.shape " + str(img_new.shape))
  663. #
  664. # # 调用模型运行接口
  665. # _dict = {"inputs": inputs, "md5": _global.get("md5")}
  666. # result = from_gpu_interface(_dict, model_type="otr", predictor_type="")
  667. # if judge_error_code(result):
  668. # logging.error("from_gpu_interface failed! " + str(result))
  669. # raise requests.exceptions.RequestException
  670. #
  671. # pred = result.get("preds")
  672. # gpu_time = result.get("gpu_time")
  673. # log("otr model predict time " + str(gpu_time))
  674. #
  675. # # # 解压numpy
  676. # # decompressed_array = io.BytesIO()
  677. # # decompressed_array.write(pred)
  678. # # decompressed_array.seek(0)
  679. # # pred = np.load(decompressed_array, allow_pickle=True)['arr_0']
  680. # # log("inputs.shape" + str(pred.shape))
  681. #
  682. # 调用gpu共享内存处理
  683. # _dict = {"inputs": inputs, "md5": _global.get("md5")}
  684. # result = from_gpu_share_memory(_dict, model_type="otr", predictor_type="")
  685. # if judge_error_code(result):
  686. # logging.error("from_gpu_interface failed! " + str(result))
  687. # raise requests.exceptions.RequestException
  688. #
  689. # pred = result.get("preds")
  690. # gpu_time = result.get("gpu_time")
  691. # log("otr model predict time " + str(gpu_time))
  692. #
  693. # # 后处理
  694. # list_line = table_postprocess(img_new, pred, prob)
  695. # log("len(list_line) " + str(len(list_line)))
  696. # if judge_error_code(list_line):
  697. # return list_line
  698. # otr resize后得到的bbox根据比例还原
  699. start_time = time.time()
  700. ratio = (image_np.shape[0]/best_h, image_np.shape[1]/best_w)
  701. for i in range(len(list_line)):
  702. point = list_line[i]
  703. list_line[i] = [int(point[0]*ratio[1]), int(point[1]*ratio[0]),
  704. int(point[2]*ratio[1]), int(point[3]*ratio[0])]
  705. log("otr resize bbox recover " + str(time.time()-start_time))
  706. # ocr图片过大内存溢出,需resize
  707. start_time = time.time()
  708. threshold = 3000
  709. ocr_resize_flag = 0
  710. if image_np.shape[0] >= threshold or image_np.shape[1] >= threshold:
  711. ocr_resize_flag = 1
  712. best_h, best_w = get_best_predict_size2(image_np, threshold)
  713. # image_resize = cv2.resize(image_np, (best_w, best_h), interpolation=cv2.INTER_AREA)
  714. image_resize = pil_resize(image_np, best_h, best_w)
  715. log("ocr_process image resize " + str(image_resize.shape))
  716. image_resize_path = image_path.split(".")[0] + "_resize_ocr." + image_path.split(".")[-1]
  717. cv2.imwrite(image_resize_path, image_resize)
  718. log("ocr resize before " + str(time.time()-start_time))
  719. # 调用ocr模型接口
  720. with open(image_resize_path, "rb") as f:
  721. image_bytes = f.read()
  722. text_list, bbox_list = from_ocr_interface(image_bytes, is_table=True)
  723. if judge_error_code(text_list):
  724. return text_list
  725. # # PaddleOCR内部包括预处理,调用模型运行接口,后处理
  726. # paddle_ocr = PaddleOCR(use_angle_cls=True, lang="ch")
  727. # results = paddle_ocr.ocr(image_resize, det=True, rec=True, cls=True)
  728. # # 循环每张图片识别结果
  729. # text_list = []
  730. # bbox_list = []
  731. # for line in results:
  732. # # print("ocr_interface line", line)
  733. # text_list.append(line[-1][0])
  734. # bbox_list.append(line[0])
  735. # if len(text_list) == 0:
  736. # return []
  737. # ocr resize后的bbox还原
  738. if ocr_resize_flag:
  739. ratio = (image_np.shape[0]/best_h, image_np.shape[1]/best_w)
  740. else:
  741. ratio = (1, 1)
  742. for i in range(len(bbox_list)):
  743. point = bbox_list[i]
  744. bbox_list[i] = [[int(point[0][0]*ratio[1]), int(point[0][1]*ratio[0])],
  745. [int(point[1][0]*ratio[1]), int(point[1][1]*ratio[0])],
  746. [int(point[2][0]*ratio[1]), int(point[2][1]*ratio[0])],
  747. [int(point[3][0]*ratio[1]), int(point[3][1]*ratio[0])]]
  748. # 调用现成方法形成表格
  749. try:
  750. from format_convert.convert_tree import TableLine
  751. list_lines = []
  752. for line in list_line:
  753. list_lines.append(LTLine(1, (line[0], line[1]), (line[2], line[3])))
  754. from format_convert.convert_tree import TextBox
  755. list_text_boxes = []
  756. for i in range(len(bbox_list)):
  757. bbox = bbox_list[i]
  758. b_text = text_list[i]
  759. list_text_boxes.append(TextBox([bbox[0][0], bbox[0][1],
  760. bbox[2][0], bbox[2][1]], b_text))
  761. # for _textbox in list_text_boxes:
  762. # print("==",_textbox.get_text())
  763. lt = LineTable()
  764. tables, obj_in_table, _ = lt.recognize_table(list_text_boxes, list_lines, False)
  765. # 合并同一行textbox
  766. list_text_boxes = merge_textbox(list_text_boxes, obj_in_table)
  767. obj_list = []
  768. for table in tables:
  769. obj_list.append(_Table(table["table"], table["bbox"]))
  770. for text_box in list_text_boxes:
  771. if text_box not in obj_in_table:
  772. obj_list.append(_Sentence(text_box.get_text(), text_box.bbox))
  773. return obj_list
  774. except:
  775. traceback.print_exc()
  776. return [-8]
  777. except Exception as e:
  778. log("image_preprocess error")
  779. traceback.print_exc()
  780. return [-1]
  781. if __name__ == "__main__":
  782. image_slice_new(cv2.imread("C:/Users/Administrator/Desktop/test_image/error28.jpg"))