table_line.py 56 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485148614871488148914901491149214931494149514961497149814991500150115021503150415051506
  1. #!/usr/bin/env python3
  2. # -*- coding: utf-8 -*-
  3. """
  4. Created on Thu Sep 9 23:11:51 2020
  5. table line detect
  6. @author: chineseocr
  7. """
  8. import copy
  9. import tensorflow as tf
  10. import keras.backend as K
  11. from tensorflow.keras.models import Model
  12. from tensorflow.keras.layers import Input, concatenate, Conv2D, MaxPooling2D, BatchNormalization, UpSampling2D
  13. from tensorflow.keras.layers import LeakyReLU
  14. from config import tableModeLinePath
  15. from utils import letterbox_image, get_table_line, adjust_lines, line_to_line, draw_boxes
  16. import numpy as np
  17. import cv2
  18. import time
  19. from utils import draw_lines
  20. def dice_coef(y_true, y_pred, smooth=1e-5):
  21. y_true_f = K.flatten(y_true)
  22. y_pred_f = K.flatten(y_pred)
  23. intersection = K.sum(y_true_f * y_pred_f)
  24. return (2. * intersection + smooth) / (K.sum(y_true_f) + K.sum(y_pred_f) + smooth)
  25. def dice_coef_loss():
  26. def dice_coef_loss_fixed(y_true, y_pred):
  27. return -dice_coef(y_true, y_pred)
  28. return dice_coef_loss_fixed
  29. def focal_loss(gamma=3., alpha=.85):
  30. def focal_loss_fixed(y_true, y_pred):
  31. pt_1 = tf.where(tf.equal(y_true, 1), y_pred, tf.ones_like(y_pred))
  32. pt_0 = tf.where(tf.equal(y_true, 0), y_pred, tf.zeros_like(y_pred))
  33. return -K.sum(alpha * K.pow(1. - pt_1, gamma) * K.log(K.epsilon()+pt_1))-K.sum((1-alpha) * K.pow( pt_0, gamma) * K.log(1. - pt_0 + K.epsilon()))
  34. return focal_loss_fixed
  35. def table_net(input_shape=(1152, 896, 3), num_classes=1):
  36. inputs = Input(shape=input_shape)
  37. # 512
  38. use_bias = False
  39. down0a = Conv2D(16, (3, 3), padding='same', use_bias=use_bias)(inputs)
  40. down0a = BatchNormalization()(down0a)
  41. down0a = LeakyReLU(alpha=0.1)(down0a)
  42. down0a = Conv2D(16, (3, 3), padding='same', use_bias=use_bias)(down0a)
  43. down0a = BatchNormalization()(down0a)
  44. down0a = LeakyReLU(alpha=0.1)(down0a)
  45. down0a_pool = MaxPooling2D((2, 2), strides=(2, 2))(down0a)
  46. # 256
  47. down0 = Conv2D(32, (3, 3), padding='same', use_bias=use_bias)(down0a_pool)
  48. down0 = BatchNormalization()(down0)
  49. down0 = LeakyReLU(alpha=0.1)(down0)
  50. down0 = Conv2D(32, (3, 3), padding='same', use_bias=use_bias)(down0)
  51. down0 = BatchNormalization()(down0)
  52. down0 = LeakyReLU(alpha=0.1)(down0)
  53. down0_pool = MaxPooling2D((2, 2), strides=(2, 2))(down0)
  54. # 128
  55. down1 = Conv2D(64, (3, 3), padding='same', use_bias=use_bias)(down0_pool)
  56. down1 = BatchNormalization()(down1)
  57. down1 = LeakyReLU(alpha=0.1)(down1)
  58. down1 = Conv2D(64, (3, 3), padding='same', use_bias=use_bias)(down1)
  59. down1 = BatchNormalization()(down1)
  60. down1 = LeakyReLU(alpha=0.1)(down1)
  61. down1_pool = MaxPooling2D((2, 2), strides=(2, 2))(down1)
  62. # 64
  63. down2 = Conv2D(128, (3, 3), padding='same', use_bias=use_bias)(down1_pool)
  64. down2 = BatchNormalization()(down2)
  65. down2 = LeakyReLU(alpha=0.1)(down2)
  66. down2 = Conv2D(128, (3, 3), padding='same', use_bias=use_bias)(down2)
  67. down2 = BatchNormalization()(down2)
  68. down2 = LeakyReLU(alpha=0.1)(down2)
  69. down2_pool = MaxPooling2D((2, 2), strides=(2, 2))(down2)
  70. # 32
  71. down3 = Conv2D(256, (3, 3), padding='same', use_bias=use_bias)(down2_pool)
  72. down3 = BatchNormalization()(down3)
  73. down3 = LeakyReLU(alpha=0.1)(down3)
  74. down3 = Conv2D(256, (3, 3), padding='same', use_bias=use_bias)(down3)
  75. down3 = BatchNormalization()(down3)
  76. down3 = LeakyReLU(alpha=0.1)(down3)
  77. down3_pool = MaxPooling2D((2, 2), strides=(2, 2))(down3)
  78. # 16
  79. down4 = Conv2D(512, (3, 3), padding='same', use_bias=use_bias)(down3_pool)
  80. down4 = BatchNormalization()(down4)
  81. down4 = LeakyReLU(alpha=0.1)(down4)
  82. down4 = Conv2D(512, (3, 3), padding='same', use_bias=use_bias)(down4)
  83. down4 = BatchNormalization()(down4)
  84. down4 = LeakyReLU(alpha=0.1)(down4)
  85. down4_pool = MaxPooling2D((2, 2), strides=(2, 2))(down4)
  86. # 8
  87. center = Conv2D(1024, (3, 3), padding='same', use_bias=use_bias)(down4_pool)
  88. center = BatchNormalization()(center)
  89. center = LeakyReLU(alpha=0.1)(center)
  90. center = Conv2D(1024, (3, 3), padding='same', use_bias=use_bias)(center)
  91. center = BatchNormalization()(center)
  92. center = LeakyReLU(alpha=0.1)(center)
  93. # center
  94. up4 = UpSampling2D((2, 2))(center)
  95. up4 = concatenate([down4, up4], axis=3)
  96. up4 = Conv2D(512, (3, 3), padding='same', use_bias=use_bias)(up4)
  97. up4 = BatchNormalization()(up4)
  98. up4 = LeakyReLU(alpha=0.1)(up4)
  99. up4 = Conv2D(512, (3, 3), padding='same', use_bias=use_bias)(up4)
  100. up4 = BatchNormalization()(up4)
  101. up4 = LeakyReLU(alpha=0.1)(up4)
  102. up4 = Conv2D(512, (3, 3), padding='same', use_bias=use_bias)(up4)
  103. up4 = BatchNormalization()(up4)
  104. up4 = LeakyReLU(alpha=0.1)(up4)
  105. # 16
  106. up3 = UpSampling2D((2, 2))(up4)
  107. up3 = concatenate([down3, up3], axis=3)
  108. up3 = Conv2D(256, (3, 3), padding='same', use_bias=use_bias)(up3)
  109. up3 = BatchNormalization()(up3)
  110. up3 = LeakyReLU(alpha=0.1)(up3)
  111. up3 = Conv2D(256, (3, 3), padding='same', use_bias=use_bias)(up3)
  112. up3 = BatchNormalization()(up3)
  113. up3 = LeakyReLU(alpha=0.1)(up3)
  114. up3 = Conv2D(256, (3, 3), padding='same', use_bias=use_bias)(up3)
  115. up3 = BatchNormalization()(up3)
  116. up3 = LeakyReLU(alpha=0.1)(up3)
  117. # 32
  118. up2 = UpSampling2D((2, 2))(up3)
  119. up2 = concatenate([down2, up2], axis=3)
  120. up2 = Conv2D(128, (3, 3), padding='same', use_bias=use_bias)(up2)
  121. up2 = BatchNormalization()(up2)
  122. up2 = LeakyReLU(alpha=0.1)(up2)
  123. up2 = Conv2D(128, (3, 3), padding='same', use_bias=use_bias)(up2)
  124. up2 = BatchNormalization()(up2)
  125. up2 = LeakyReLU(alpha=0.1)(up2)
  126. up2 = Conv2D(128, (3, 3), padding='same', use_bias=use_bias)(up2)
  127. up2 = BatchNormalization()(up2)
  128. up2 = LeakyReLU(alpha=0.1)(up2)
  129. # 64
  130. up1 = UpSampling2D((2, 2))(up2)
  131. up1 = concatenate([down1, up1], axis=3)
  132. up1 = Conv2D(64, (3, 3), padding='same', use_bias=use_bias)(up1)
  133. up1 = BatchNormalization()(up1)
  134. up1 = LeakyReLU(alpha=0.1)(up1)
  135. up1 = Conv2D(64, (3, 3), padding='same', use_bias=use_bias)(up1)
  136. up1 = BatchNormalization()(up1)
  137. up1 = LeakyReLU(alpha=0.1)(up1)
  138. up1 = Conv2D(64, (3, 3), padding='same', use_bias=use_bias)(up1)
  139. up1 = BatchNormalization()(up1)
  140. up1 = LeakyReLU(alpha=0.1)(up1)
  141. # 128
  142. up0 = UpSampling2D((2, 2))(up1)
  143. up0 = concatenate([down0, up0], axis=3)
  144. up0 = Conv2D(32, (3, 3), padding='same', use_bias=use_bias)(up0)
  145. up0 = BatchNormalization()(up0)
  146. up0 = LeakyReLU(alpha=0.1)(up0)
  147. up0 = Conv2D(32, (3, 3), padding='same', use_bias=use_bias)(up0)
  148. up0 = BatchNormalization()(up0)
  149. up0 = LeakyReLU(alpha=0.1)(up0)
  150. up0 = Conv2D(32, (3, 3), padding='same', use_bias=use_bias)(up0)
  151. up0 = BatchNormalization()(up0)
  152. up0 = LeakyReLU(alpha=0.1)(up0)
  153. # 256
  154. up0a = UpSampling2D((2, 2))(up0)
  155. up0a = concatenate([down0a, up0a], axis=3)
  156. up0a = Conv2D(16, (3, 3), padding='same', use_bias=use_bias)(up0a)
  157. up0a = BatchNormalization()(up0a)
  158. up0a = LeakyReLU(alpha=0.1)(up0a)
  159. up0a = Conv2D(16, (3, 3), padding='same', use_bias=use_bias)(up0a)
  160. up0a = BatchNormalization()(up0a)
  161. up0a = LeakyReLU(alpha=0.1)(up0a)
  162. up0a = Conv2D(16, (3, 3), padding='same', use_bias=use_bias)(up0a)
  163. up0a = BatchNormalization()(up0a)
  164. up0a = LeakyReLU(alpha=0.1)(up0a)
  165. # 512
  166. classify = Conv2D(num_classes, (1, 1), activation='sigmoid')(up0a)
  167. model = Model(inputs=inputs, outputs=classify)
  168. return model
  169. model = table_net((None, None, 3), 2)
  170. def table_line(img, model, size=(512, 1024), hprob=0.5, vprob=0.5, row=50, col=30, alph=15):
  171. sizew, sizeh = size
  172. # [..., ::-1] 最后一维内部反向输出
  173. # inputBlob, fx, fy = letterbox_image(img[..., ::-1], (sizew, sizeh))
  174. # pred = model.predict(np.array([np.array(inputBlob)]))
  175. # pred = model.predict(np.array([np.array(inputBlob)/255.0]))
  176. img_new = cv2.resize(img, (sizew, sizeh), interpolation=cv2.INTER_AREA)
  177. pred = model.predict(np.array([img_new]))
  178. pred = pred[0]
  179. vpred = pred[..., 1] > vprob # 横线
  180. hpred = pred[..., 0] > hprob # 竖线
  181. vpred = vpred.astype(int)
  182. hpred = hpred.astype(int)
  183. # print("vpred shape", vpred)
  184. # print("hpred shape", hpred)
  185. colboxes = get_table_line(vpred, axis=1, lineW=col)
  186. rowboxes = get_table_line(hpred, axis=0, lineW=row)
  187. # if len(rowboxes) > 0:
  188. # rowboxes = np.array(rowboxes)
  189. # rowboxes[:, [0, 2]] = rowboxes[:, [0, 2]]/fx
  190. # rowboxes[:, [1, 3]] = rowboxes[:, [1, 3]]/fy
  191. # rowboxes = rowboxes.tolist()
  192. # if len(colboxes) > 0:
  193. # colboxes = np.array(colboxes)
  194. # colboxes[:, [0, 2]] = colboxes[:, [0, 2]]/fx
  195. # colboxes[:, [1, 3]] = colboxes[:, [1, 3]]/fy
  196. # colboxes = colboxes.tolist()
  197. nrow = len(rowboxes)
  198. ncol = len(colboxes)
  199. for i in range(nrow):
  200. for j in range(ncol):
  201. rowboxes[i] = line_to_line(rowboxes[i], colboxes[j], 10)
  202. colboxes[j] = line_to_line(colboxes[j], rowboxes[i], 10)
  203. # 删掉贴着边框的line
  204. temp_list = []
  205. threshold = 5
  206. for line in rowboxes:
  207. if line[1]-0 <= threshold or size[1]-line[1] <= threshold:
  208. continue
  209. temp_list.append(line)
  210. rowboxes = temp_list
  211. temp_list = []
  212. for line in colboxes:
  213. if line[0]-0 <= threshold or size[0]-line[0] <= threshold:
  214. continue
  215. temp_list.append(line)
  216. colboxes = temp_list
  217. return rowboxes, colboxes, img_new
  218. def get_outline(points, image_np):
  219. # 取出x, y的最大值最小值
  220. x_min = points[0][0]
  221. x_max = points[-1][0]
  222. points.sort(key=lambda x: (x[1], x[0]))
  223. y_min = points[0][1]
  224. y_max = points[-1][1]
  225. # 创建空图
  226. # outline_img = np.zeros(image_size, np.uint8)
  227. outline_img = np.copy(image_np)
  228. cv2.rectangle(outline_img, (x_min-5, y_min-5), (x_max+5, y_max+5), (0, 0, 0), 2)
  229. # cv2.imshow("outline_img", outline_img)
  230. # cv2.waitKey(0)
  231. return outline_img
  232. def get_split_line(points, col_lines, image_np):
  233. # print("get_split_line", image_np.shape)
  234. points.sort(key=lambda x: (x[1], x[0]))
  235. # 遍历y坐标,并判断y坐标与上一个y坐标是否存在连接线
  236. i = 0
  237. split_line_y = []
  238. for point in points:
  239. # 从已分开的线下面开始判断
  240. if split_line_y:
  241. if point[1] <= split_line_y[-1] + 5:
  242. last_y = point[1]
  243. continue
  244. if last_y <= split_line_y[-1] + 5:
  245. last_y = point[1]
  246. continue
  247. if i == 0:
  248. last_y = point[1]
  249. i += 1
  250. continue
  251. current_line = (last_y, point[1])
  252. split_flag = 1
  253. for col in col_lines:
  254. # 只要找到一条col包含就不是分割线
  255. if current_line[0] >= col[1]-3 and current_line[1] <= col[3]+3:
  256. split_flag = 0
  257. # print("img", img.shape)
  258. # print("col", col)
  259. # print("current_line", current_line)
  260. break
  261. if split_flag:
  262. split_line_y.append(current_line[0]+5)
  263. split_line_y.append(current_line[1]-5)
  264. last_y = point[1]
  265. # 加上收尾分割线
  266. points.sort(key=lambda x: (x[1], x[0]))
  267. y_min = points[0][1]
  268. y_max = points[-1][1]
  269. # print("加上收尾分割线", y_min, y_max)
  270. if y_min-5 < 0:
  271. split_line_y.append(0)
  272. else:
  273. split_line_y.append(y_min-5)
  274. if y_max+5 > image_np.shape[0]:
  275. split_line_y.append(image_np.shape[0])
  276. else:
  277. split_line_y.append(y_max+5)
  278. split_line_y = list(set(split_line_y))
  279. # 剔除两条相隔太近分割线
  280. temp_split_line_y = []
  281. split_line_y.sort(key=lambda x: x)
  282. last_y = -20
  283. for y in split_line_y:
  284. # print(y)
  285. if y - last_y >= 20:
  286. # print(y, last_y)
  287. temp_split_line_y.append(y)
  288. last_y = y
  289. split_line_y = temp_split_line_y
  290. # print("split_line_y", split_line_y)
  291. # 生成分割线
  292. split_line = []
  293. last_y = 0
  294. for y in split_line_y:
  295. # if y - last_y <= 15:
  296. # continue
  297. split_line.append([(0, y), (image_np.shape[1], y)])
  298. last_y = y
  299. split_line.append([(0, 0), (image_np.shape[1], 0)])
  300. split_line.append([(0, image_np.shape[0]), (image_np.shape[1], image_np.shape[0])])
  301. split_line.sort(key=lambda x: x[0][1])
  302. # print("split_line", split_line)
  303. # 画图画线
  304. # split_line_img = np.copy(image_np)
  305. # for y in split_line_y:
  306. # cv2.line(split_line_img, (0, y), (image_np.shape[1], y), (0, 0, 0), 1)
  307. # cv2.imshow("split_line_img", split_line_img)
  308. # cv2.waitKey(0)
  309. return split_line, split_line_y
  310. def get_points(row_lines, col_lines, image_size):
  311. # 创建空图
  312. row_img = np.zeros(image_size, np.uint8)
  313. col_img = np.zeros(image_size, np.uint8)
  314. # 画线
  315. for row in row_lines:
  316. cv2.line(row_img, (int(row[0]), int(row[1])), (int(row[2]), int(row[3])), (255, 255, 255), 1)
  317. for col in col_lines:
  318. cv2.line(col_img, (int(col[0]), int(col[1])), (int(col[2]), int(col[3])), (255, 255, 255), 1)
  319. # 求出交点
  320. point_img = np.bitwise_and(row_img, col_img)
  321. # cv2.imshow("point_img", np.bitwise_not(point_img))
  322. # cv2.waitKey(0)
  323. # 识别黑白图中的白色交叉点,将横纵坐标取出
  324. ys, xs = np.where(point_img > 0)
  325. points = []
  326. for i in range(len(xs)):
  327. points.append((xs[i], ys[i]))
  328. points.sort(key=lambda x: (x[0], x[1]))
  329. return points
  330. def get_minAreaRect(image):
  331. gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
  332. gray = cv2.bitwise_not(gray)
  333. thresh = cv2.threshold(gray, 0, 255,
  334. cv2.THRESH_BINARY | cv2.THRESH_OTSU)[1]
  335. coords = np.column_stack(np.where(thresh > 0))
  336. return cv2.minAreaRect(coords)
  337. def get_contours(image):
  338. gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
  339. ret, binary = cv2.threshold(gray, 127, 255, cv2.THRESH_BINARY)
  340. contours, hierarchy = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
  341. cv2.drawContours(image, contours, -1, (0, 0, 255), 3)
  342. cv2.imshow("get contours", image)
  343. cv2.waitKey(0)
  344. def fix_outline(image, row_lines, col_lines, points, split_y):
  345. # 分割线纵坐标
  346. # print("len(split_y)", split_y)
  347. if len(split_y) < 2:
  348. # split_y = [2000., 2000., 2000., 2000.]
  349. return [], [], []
  350. elif len(split_y) == 2:
  351. split_y = [2000., 2000., 2000., 2000.]
  352. split_y.sort(key=lambda x: x)
  353. split_y = split_y[1:-1]
  354. new_split_y = []
  355. for i in range(1, len(split_y), 2):
  356. new_split_y.append(int((split_y[i]+split_y[i-1])/2))
  357. # print("len(new_split_y)", len(new_split_y))
  358. # 预测线根据分割线纵坐标分为多个分割区域
  359. row_lines.sort(key=lambda x: (x[3], x[2]))
  360. col_lines.sort(key=lambda x: (x[3], x[2]))
  361. row_count = 0
  362. col_count = 0
  363. split_row_list = []
  364. split_col_list = []
  365. for y in new_split_y:
  366. row_lines = row_lines[row_count:]
  367. col_lines = col_lines[col_count:]
  368. row_count = 0
  369. col_count = 0
  370. for row in row_lines:
  371. if row[3] <= y:
  372. row_count += 1
  373. else:
  374. split_row_list.append(row_lines[:row_count])
  375. break
  376. if row_count == len(row_lines):
  377. split_row_list.append(row_lines[:row_count])
  378. break
  379. for col in col_lines:
  380. if col[3] <= y:
  381. col_count += 1
  382. else:
  383. split_col_list.append(col_lines[:col_count])
  384. break
  385. if col_count == len(col_lines):
  386. split_col_list.append(col_lines[:col_count])
  387. break
  388. if row_count < len(row_lines) - 1:
  389. row_lines = row_lines[row_count:]
  390. split_row_list.append(row_lines)
  391. if col_count < len(col_lines) - 1:
  392. col_lines = col_lines[col_count:]
  393. split_col_list.append(col_lines)
  394. # print("len(split_row_list)", len(split_row_list))
  395. # print("len(split_col_list)", len(split_col_list))
  396. # 预测线取上下左右4个边(会有超出表格部分) [(), ()]
  397. area_row_line = []
  398. area_col_line = []
  399. for area in split_row_list:
  400. if not area:
  401. continue
  402. area.sort(key=lambda x: (x[1], x[0]))
  403. # print(area)
  404. up_line = area[0]
  405. bottom_line = area[-1]
  406. area_row_line.append([up_line, bottom_line])
  407. for area in split_col_list:
  408. if not area:
  409. continue
  410. area.sort(key=lambda x: x[0])
  411. left_line = area[0]
  412. right_line = area[-1]
  413. area_col_line.append([left_line, right_line])
  414. # 线交点根据分割线纵坐标分为多个分割区域
  415. points.sort(key=lambda x: (x[1], x[0]))
  416. point_count = 0
  417. split_point_list = []
  418. for y in new_split_y:
  419. points = points[point_count:len(points)]
  420. point_count = 0
  421. for point in points:
  422. if point[1] <= y:
  423. point_count += 1
  424. else:
  425. split_point_list.append(points[:point_count])
  426. break
  427. if point_count == len(points):
  428. split_point_list.append(points[:point_count])
  429. break
  430. if point_count < len(points) - 1:
  431. points = points[point_count:len(points)]
  432. split_point_list.append(points)
  433. # print("len(split_point_list)", len(split_point_list))
  434. # 取每个分割区域的4条线(无超出表格部分)
  435. area_row_line2 = []
  436. area_col_line2 = []
  437. for area in split_point_list:
  438. if not area:
  439. continue
  440. area.sort(key=lambda x: (x[0], x[1]))
  441. left_up = area[0]
  442. right_bottom = area[-1]
  443. up_line = [left_up[0], left_up[1], right_bottom[0], left_up[1]]
  444. bottom_line = [left_up[0], right_bottom[1], right_bottom[0], right_bottom[1]]
  445. left_line = [left_up[0], left_up[1], left_up[0], right_bottom[1]]
  446. right_line = [right_bottom[0], left_up[1], right_bottom[0], right_bottom[1]]
  447. area_row_line2.append([up_line, bottom_line])
  448. area_col_line2.append([left_line, right_line])
  449. # 判断超出部分的长度,超出一定长度就补线
  450. new_row_lines = []
  451. new_col_lines = []
  452. longer_row_lines = []
  453. longer_col_lines = []
  454. # print(len(area_col_line))
  455. # print(len(area_col_line2))
  456. for i in range(len(area_row_line)):
  457. up_line = area_row_line[i][0]
  458. up_line2 = area_row_line2[i][0]
  459. bottom_line = area_row_line[i][1]
  460. bottom_line2 = area_row_line2[i][1]
  461. left_line = area_col_line[i][0]
  462. left_line2 = area_col_line2[i][0]
  463. right_line = area_col_line[i][1]
  464. right_line2 = area_col_line2[i][1]
  465. # 补左右两条竖线超出来的线的row
  466. # print("left_line2[1] - left_line[1]", left_line2[1] - left_line[1])
  467. if left_line2[1] - left_line[1] >= 10 and right_line2[1] - right_line[1] >= 10:
  468. if left_line2[1] - left_line[1] >= right_line2[1] - right_line[1]:
  469. new_row_lines.append([left_line[0], left_line[1], right_line[0], left_line[1]])
  470. new_col_y = left_line[1]
  471. # 补了row,要将其他短的col连到row上
  472. for col in split_col_list[i]:
  473. longer_col_lines.append([col[0], new_col_y, col[2], col[3]])
  474. else:
  475. new_row_lines.append([left_line[0], right_line[1], right_line[0], right_line[1]])
  476. new_col_y = right_line[1]
  477. # 补了row,要将其他短的col连到row上
  478. for col in split_col_list[i]:
  479. longer_col_lines.append([col[0], new_col_y, col[2], col[3]])
  480. # print("left_line[3] - left_line2[3]", left_line[3] - left_line2[3])
  481. if left_line[3] - left_line2[3] >= 10 and right_line[3] - right_line2[3] >= 10:
  482. if left_line[3] - left_line2[3] >= right_line[3] - right_line2[3]:
  483. new_row_lines.append([left_line[2], left_line[3], right_line[2], left_line[3]])
  484. new_col_y = left_line[3]
  485. # 补了row,要将其他短的col连到row上
  486. for col in split_col_list[i]:
  487. longer_col_lines.append([col[0], col[1], col[2], new_col_y])
  488. else:
  489. new_row_lines.append([left_line[2], right_line[3], right_line[2], right_line[3]])
  490. new_col_y = right_line[3]
  491. # 补了row,要将其他短的col连到row上
  492. for col in split_col_list[i]:
  493. longer_col_lines.append([col[0], col[1], col[2], new_col_y])
  494. # 补上下两条横线超出来的线的col
  495. if up_line2[0] - up_line[0] >= 10 and bottom_line2[0] - bottom_line[0] >= 10:
  496. if up_line2[0] - up_line[0] >= bottom_line2[0] - bottom_line[0]:
  497. new_col_lines.append([up_line[0], up_line[1], up_line[0], bottom_line[1]])
  498. new_row_x = up_line[0]
  499. # 补了col,要将其他短的row连到col上
  500. for row in split_row_list[i]:
  501. longer_row_lines.append([new_row_x, row[1], row[2], row[3]])
  502. # new_row_lines.append([bottom_line[0], up_line[1], bottom_line[2], bottom_line[3]])
  503. else:
  504. new_col_lines.append([bottom_line[0], up_line[1], bottom_line[0], bottom_line[1]])
  505. new_row_x = bottom_line[0]
  506. # 补了col,要将其他短的row连到col上
  507. for row in split_row_list[i]:
  508. longer_row_lines.append([new_row_x, row[1], row[2], row[3]])
  509. # new_row_lines.append([bottom_line[0], up_line[1], up_line[2], up_line[3]])
  510. if up_line[2] - up_line2[2] >= 10 and bottom_line[2] - bottom_line2[2] >= 10:
  511. if up_line[2] - up_line2[2] >= bottom_line[2] - bottom_line2[2]:
  512. new_col_lines.append([up_line[2], up_line[3], up_line[2], bottom_line[3]])
  513. new_row_x = up_line[2]
  514. # 补了col,要将其他短的row连到col上
  515. for row in split_row_list[i]:
  516. longer_row_lines.append([row[0], row[1], new_row_x, row[3]])
  517. else:
  518. new_col_lines.append([bottom_line[2], up_line[3], bottom_line[2], bottom_line[3]])
  519. new_row_x = bottom_line[2]
  520. # 补了col,要将其他短的row连到col上
  521. for row in split_row_list[i]:
  522. longer_row_lines.append([row[0], row[1], new_row_x, row[3]])
  523. # 删除表格内部的补线
  524. # temp_list = []
  525. # for row in new_row_lines:
  526. # if up_line[1]-5 <= row[1] <= bottom_line[1]+5:
  527. # continue
  528. # temp_list.append(row)
  529. # print("fix_outline", new_row_lines)
  530. # new_row_lines = temp_list
  531. # print("fix_outline", new_row_lines)
  532. # temp_list = []
  533. # for col in new_col_lines:
  534. # if left_line[0]-5 <= col[0] <= right_line[0]+5:
  535. # continue
  536. # temp_list.append(col)
  537. #
  538. # new_col_lines = temp_list
  539. # print("fix_outline", new_col_lines)
  540. # print("fix_outline", new_row_lines)
  541. # 删除重复包含的补线
  542. # temp_list = []
  543. # for row in new_row_lines:
  544. # if up_line[1]-5 <= row[1] <= bottom_line[1]+5:
  545. # continue
  546. # temp_list.append(row)
  547. # new_row_lines = temp_list
  548. # 展示上下左右边框线
  549. # for i in range(len(area_row_line)):
  550. # print("row1", area_row_line[i])
  551. # print("row2", area_row_line2[i])
  552. # print("col1", area_col_line[i])
  553. # print("col2", area_col_line2[i])
  554. # cv2.line(image, (int(area_row_line[i][0][0]), int(area_row_line[i][0][1])),
  555. # (int(area_row_line[i][0][2]), int(area_row_line[i][0][3])), (0, 255, 0), 2)
  556. # cv2.line(image, (int(area_row_line2[i][1][0]), int(area_row_line2[i][1][1])),
  557. # (int(area_row_line2[i][1][2]), int(area_row_line2[i][1][3])), (0, 0, 255), 2)
  558. # cv2.imshow("fix_outline", image)
  559. # cv2.waitKey(0)
  560. return new_row_lines, new_col_lines, longer_row_lines, longer_col_lines
  561. def delete_close_points(point_list, row_point_list, col_point_list, threshold=10):
  562. new_point_list = []
  563. delete_point_list = []
  564. point_list.sort(key=lambda x: (x[1], x[0]))
  565. for i in range(len(point_list)):
  566. point1 = point_list[i]
  567. if point1 in delete_point_list:
  568. continue
  569. if i == len(point_list) - 1:
  570. new_point_list.append(point1)
  571. break
  572. point2 = point_list[i+1]
  573. # 判断坐标
  574. if abs(point1[0] - point2[0]) > threshold or abs(point1[1] - point2[1]) > threshold:
  575. new_point_list.append(point1)
  576. else:
  577. # 看两个点上的相同坐标点哪个多,就保留哪个
  578. count1 = 0
  579. count2 = 0
  580. for col in col_point_list:
  581. if point1[0] == col[0][0]:
  582. count1 += len(col)
  583. elif point2[0] == col[0][0]:
  584. count2 += len(col)
  585. if count1 >= count2:
  586. new_point_list.append(point1)
  587. delete_point_list.append(point2)
  588. else:
  589. new_point_list.append(point2)
  590. delete_point_list.append(point1)
  591. point_list = new_point_list
  592. new_point_list = []
  593. delete_point_list = []
  594. point_list.sort(key=lambda x: (x[0], x[1]))
  595. for i in range(len(point_list)):
  596. point1 = point_list[i]
  597. if point1 in delete_point_list:
  598. continue
  599. if i == len(point_list) - 1:
  600. new_point_list.append(point1)
  601. break
  602. point2 = point_list[i+1]
  603. # 判断坐标
  604. if abs(point1[0] - point2[0]) > threshold or abs(point1[1] - point2[1]) > threshold:
  605. new_point_list.append(point1)
  606. else:
  607. count1 = 0
  608. count2 = 0
  609. for row in row_point_list:
  610. if point1[0] == row[0][0]:
  611. count1 += len(row)
  612. elif point2[0] == row[0][0]:
  613. count2 += len(row)
  614. if count1 >= count2:
  615. new_point_list.append(point1)
  616. delete_point_list.append(point2)
  617. else:
  618. new_point_list.append(point2)
  619. delete_point_list.append(point1)
  620. return new_point_list
  621. def get_bbox2(image_np, points):
  622. # # 坐标点按行分
  623. # row_point_list = []
  624. # row_point = []
  625. # points.sort(key=lambda x: (x[0], x[1]))
  626. # for p in points:
  627. # if len(row_point) == 0:
  628. # x = p[0]
  629. # if x-5 <= p[0] <= x+5:
  630. # row_point.append(p)
  631. # else:
  632. # row_point_list.append(row_point)
  633. # row_point = []
  634. # # 坐标点按列分
  635. # col_point_list = []
  636. # col_point = []
  637. # points.sort(key=lambda x: (x[1], x[0]))
  638. # for p in points:
  639. # if len(col_point) == 0:
  640. # y = p[1]
  641. # if y-5 <= p[1] <= y+5:
  642. # col_point.append(p)
  643. # else:
  644. # col_point_list.append(col_point)
  645. # col_point = []
  646. row_point_list = get_points_row(points)
  647. col_point_list = get_points_col(points)
  648. print("len(points)", len(points))
  649. for point in points:
  650. cv2.circle(image_np, point, 1, (0, 255, 0), 1)
  651. cv2.imshow("points_deleted", image_np)
  652. points = delete_close_points(points, row_point_list, col_point_list)
  653. print("len(points)", len(points))
  654. for point in points:
  655. cv2.circle(image_np, point, 1, (255, 0, 0), 3)
  656. cv2.imshow("points_deleted", image_np)
  657. cv2.waitKey(0)
  658. row_point_list = get_points_row(points, 5)
  659. col_point_list = get_points_col(points, 5)
  660. print("len(row_point_list)", len(row_point_list))
  661. for row in row_point_list:
  662. print("row", len(row))
  663. print("col_point_list", len(col_point_list))
  664. for col in col_point_list:
  665. print("col", len(col))
  666. bbox = []
  667. for i in range(len(row_point_list)):
  668. if i == len(row_point_list) - 1:
  669. break
  670. # 遍历每个row的point,找到其所在列的下一个点和所在行的下一个点
  671. current_row = row_point_list[i]
  672. for j in range(len(current_row)):
  673. current_point = current_row[j]
  674. if j == len(current_row) - 1:
  675. break
  676. next_row_point = current_row[j+1]
  677. # 找出当前点所在的col,得到该列下一个point
  678. current_col = col_point_list[j]
  679. for k in range(len(current_col)):
  680. if current_col[k][1] > current_point[1] + 10:
  681. next_col_point = current_col[k]
  682. break
  683. next_row = row_point_list[k]
  684. for k in range(len(next_row)):
  685. if next_row[k][0] >= next_row_point[0] + 5:
  686. next_point = next_row[k]
  687. break
  688. # 得到bbox
  689. bbox.append([(current_point[0], current_point[1]), (next_point[0], next_point[1])])
  690. # bbox = []
  691. # for p in points:
  692. # # print("p", p)
  693. # p_row = []
  694. # p_col = []
  695. # for row in row_point_list:
  696. # if p[0] == row[0][0]:
  697. # for p1 in row:
  698. # if abs(p[1]-p1[1]) <= 5:
  699. # continue
  700. # p_row.append([p1, abs(p[1]-p1[1])])
  701. # p_row.sort(key=lambda x: x[1])
  702. # for col in col_point_list:
  703. # if p[1] == col[0][1]:
  704. # for p2 in col:
  705. # if abs(p[0]-p2[0]) <= 5:
  706. # continue
  707. # p_col.append([p2, abs(p[0]-p2[0])])
  708. # p_col.sort(key=lambda x: x[1])
  709. # if len(p_row) == 0 or len(p_col) == 0:
  710. # continue
  711. # break_flag = 0
  712. # for i in range(len(p_row)):
  713. # for j in range(len(p_col)):
  714. # # print(p_row[i][0])
  715. # # print(p_col[j][0])
  716. # another_point = (p_col[j][0][0], p_row[i][0][1])
  717. # # print("another_point", another_point)
  718. # if abs(p[0]-another_point[0]) <= 5 or abs(p[1]-another_point[1]) <= 5:
  719. # continue
  720. # if p[0] >= another_point[0] or p[1] >= another_point[1]:
  721. # continue
  722. # if another_point in points:
  723. # box = [p, another_point]
  724. # box.sort(key=lambda x: x[0])
  725. # if box not in bbox:
  726. # bbox.append(box)
  727. # break_flag = 1
  728. # break
  729. # if break_flag:
  730. # break
  731. #
  732. # # delete duplicate
  733. # delete_bbox = []
  734. # for i in range(len(bbox)):
  735. # for j in range(i+1, len(bbox)):
  736. # if bbox[i][0] == bbox[j][0]:
  737. # if bbox[i][1][0] - bbox[j][1][0] <= 3 \
  738. # and bbox[i][1][1] - bbox[j][1][1] <= 3:
  739. # delete_bbox.append(bbox[j])
  740. # if bbox[i][1] == bbox[j][1]:
  741. # if bbox[i][0][0] - bbox[j][0][0] <= 3 \
  742. # and bbox[i][0][1] - bbox[j][0][1] <= 3:
  743. # delete_bbox.append(bbox[j])
  744. # # delete too small area
  745. # # for box in bbox:
  746. # # if box[1][0] - box[0][0] <=
  747. # for d_box in delete_bbox:
  748. # if d_box in bbox:
  749. # bbox.remove(d_box)
  750. # print bbox
  751. bbox.sort(key=lambda x: (x[0][0], x[0][1], x[1][0], x[1][1]))
  752. # origin bbox
  753. # origin_bbox = []
  754. # for box in bbox:
  755. # origin_bbox.append([(box[0][0], box[0][1] - 40), (box[1][0], box[1][1] - 40)])
  756. # for box in origin_bbox:
  757. # cv2.rectangle(origin_image, box[0], box[1], (0, 0, 255), 2, 8)
  758. # cv2.imshow('AlanWang', origin_image)
  759. # cv2.waitKey(0)
  760. for box in bbox:
  761. cv2.rectangle(image_np, box[0], box[1], (0, 0, 255), 2, 8)
  762. cv2.imshow('bboxes', image_np)
  763. cv2.waitKey(0)
  764. # for point in points:
  765. # print(point)
  766. # cv2.circle(image_np, point, 1, (0, 0, 255), 3)
  767. # cv2.imshow('points', image_np)
  768. # cv2.waitKey(0)
  769. return bbox
  770. def get_bbox1(image_np, points, split_y):
  771. # 分割线纵坐标
  772. # print("split_y", split_y)
  773. if len(split_y) < 2:
  774. return []
  775. # 计算行列,剔除相近交点
  776. row_point_list = get_points_row(points)
  777. col_point_list = get_points_col(points)
  778. print("len(row_point_list)", row_point_list)
  779. print("len(col_point_list)", len(col_point_list))
  780. # for point in points:
  781. # cv2.circle(image_np, point, 1, (0, 255, 0), 1)
  782. # cv2.imshow("points", image_np)
  783. points = delete_close_points(points, row_point_list, col_point_list)
  784. # print("len(points)", len(points))
  785. # for point in points:
  786. # cv2.circle(image_np, point, 1, (255, 0, 0), 3)
  787. # cv2.imshow("points_deleted", image_np)
  788. # cv2.waitKey(0)
  789. # 获取bbox
  790. bbox = []
  791. # 每个点获取与其x最相近和y最相近的点
  792. for i in range(1, len(split_y)):
  793. for point1 in points:
  794. if point1[1] <= split_y[i-1] or point1[1] >= split_y[i]:
  795. continue
  796. distance_x = 10000
  797. distance_y = 10000
  798. x = 0
  799. y = 0
  800. threshold = 10
  801. for point2 in points:
  802. if point2[1] <= split_y[i-1] or point2[1] >= split_y[i]:
  803. continue
  804. # 最近 x y
  805. if 2 < point2[0] - point1[0] < distance_x and point2[1] - point1[1] <= threshold:
  806. distance_x = point2[0] - point1[0]
  807. x = point2[0]
  808. if 2 < point2[1] - point1[1] < distance_y and point2[0] - point1[0] <= threshold:
  809. distance_y = point2[1] - point1[1]
  810. y = point2[1]
  811. if not x or not y:
  812. continue
  813. bbox.append([(point1[0], point1[1]), (x, y)])
  814. # 删除包含关系bbox
  815. temp_list = []
  816. for i in range(len(bbox)):
  817. box1 = bbox[i]
  818. for j in range(len(bbox)):
  819. if i == j:
  820. continue
  821. box2 = bbox[j]
  822. contain_flag = 0
  823. if box2[0][0] <= box1[0][0] <= box1[1][0] <= box2[1][0] and \
  824. box2[0][1] <= box1[0][1] <= box1[1][1] <= box2[1][1]:
  825. contain_flag = 1
  826. break
  827. temp_list.append(box1)
  828. bbox = temp_list
  829. # 展示
  830. for box in bbox:
  831. # print(box[0], box[1])
  832. # if abs(box[0][1] - box[1][1]) > abs(box[0][0] - box[1][0]):
  833. # continue
  834. cv2.rectangle(image_np, box[0], box[1], (0, 0, 255), 2, 8)
  835. cv2.imshow('bboxes', image_np)
  836. cv2.waitKey(0)
  837. return bbox
  838. def get_bbox0(image_np, row_point_list, col_point_list, split_y, row_lines, col_lines):
  839. # 分割线纵坐标
  840. if len(split_y) < 2:
  841. return []
  842. # 计算行列,剔除相近交点
  843. # row_point_list = get_points_row(points)
  844. # col_point_list = get_points_col(points)
  845. # points = delete_close_points(points, row_point_list, col_point_list)
  846. # row_point_list = get_points_row(points)
  847. # col_point_list = get_points_col(points)
  848. # 获取bbox
  849. bbox = []
  850. # print("get_bbox split_y", split_y)
  851. # 每个点获取与其x最相近和y最相近的点
  852. for i in range(1, len(split_y)):
  853. # 循环每行
  854. for row in row_point_list:
  855. row.sort(key=lambda x: (x[0], x[1]))
  856. # 行不在该区域跳过
  857. if row[0][1] <= split_y[i-1] or row[0][1] >= split_y[i]:
  858. continue
  859. # 循环行中的点
  860. for j in range(len(row)):
  861. if j == len(row) - 1:
  862. break
  863. current_point = row[j]
  864. next_point_in_row = row[j+1]
  865. # 查询下个点所在列
  866. next_col = []
  867. for col in col_point_list:
  868. col.sort(key=lambda x: (x[1], x[0]))
  869. # 列不在该区域跳过
  870. if col[0][1] <= split_y[i-1] or col[-1][1] >= split_y[i]:
  871. continue
  872. if col[0][0]-3 <= next_point_in_row[0] <= col[0][0]+3:
  873. next_col = col
  874. break
  875. # 循环匹配当前点和下一列点
  876. for point1 in next_col:
  877. # 同一行的就跳过
  878. if current_point[1]-3 <= point1[1] <= current_point[1]+3:
  879. continue
  880. if point1[1] <= current_point[1]-3:
  881. continue
  882. # 候选bbox
  883. candidate_bbox = [current_point[0], current_point[1], point1[0], point1[1]]
  884. # 判断该bbox是否存在,线条包含关系
  885. contain_flag1 = 0
  886. contain_flag2 = 0
  887. for row1 in row_lines:
  888. # 行不在该区域跳过
  889. if row1[1] <= split_y[i-1] or row1[1] >= split_y[i]:
  890. continue
  891. # bbox上边框 y一样
  892. if not contain_flag1:
  893. if row1[1]-3 <= candidate_bbox[1] <= row1[1]+3:
  894. # candidate的x1,x2需被包含在row线中
  895. if row1[0]-3 <= candidate_bbox[0] <= candidate_bbox[2] <= row1[2]+3:
  896. contain_flag1 = 1
  897. # bbox下边框 y一样
  898. if not contain_flag2:
  899. if row1[1]-3 <= candidate_bbox[3] <= row1[1]+3:
  900. # candidate的x1,x2需被包含在row线中
  901. if row1[0]-3 <= candidate_bbox[0] <= candidate_bbox[2] <= row1[2]+3:
  902. contain_flag2 = 1
  903. # 找到了该bbox,并且是存在的
  904. if contain_flag1 and contain_flag2:
  905. bbox.append([(candidate_bbox[0], candidate_bbox[1]),
  906. (candidate_bbox[2], candidate_bbox[3])])
  907. break
  908. return bbox
  909. def get_bbox(image_np, row_point_list, col_point_list, split_y, row_lines, col_lines):
  910. # 分割线纵坐标
  911. if len(split_y) < 2:
  912. return []
  913. # 获取bbox
  914. bbox = []
  915. # 每个点获取与其x最相近和y最相近的点
  916. for i in range(1, len(split_y)):
  917. # 循环每行
  918. for row in row_point_list:
  919. row.sort(key=lambda x: (x[0], x[1]))
  920. # 行不在该区域跳过
  921. if row[0][1] <= split_y[i-1] or row[0][1] >= split_y[i]:
  922. continue
  923. # 循环行中的点
  924. for j in range(len(row)):
  925. if j == len(row) - 1:
  926. break
  927. current_point = row[j]
  928. next_point_in_row_list = row[j+1:]
  929. # 循环这一行的下一个点
  930. for next_point_in_row in next_point_in_row_list:
  931. # 是否在这一行点找到,找不到就这一行的下个点
  932. not_found = 1
  933. # 查询下个点所在列
  934. next_col = []
  935. for col in col_point_list:
  936. col.sort(key=lambda x: (x[1], x[0]))
  937. # 列不在该区域跳过
  938. if col[0][1] <= split_y[i-1] or col[-1][1] >= split_y[i]:
  939. continue
  940. if col[0][0]-3 <= next_point_in_row[0] <= col[0][0]+3:
  941. next_col = col
  942. break
  943. # 循环匹配当前点和下一列点
  944. for point1 in next_col:
  945. # 同一行的就跳过
  946. if current_point[1]-3 <= point1[1] <= current_point[1]+3:
  947. continue
  948. if point1[1] <= current_point[1]-3:
  949. continue
  950. # 候选bbox
  951. candidate_bbox = [current_point[0], current_point[1], point1[0], point1[1]]
  952. # 判断该bbox是否存在,判断bbox的上下两条边是否有包含在row中
  953. contain_flag1 = 0
  954. contain_flag2 = 0
  955. for row1 in row_lines:
  956. # 行不在该区域跳过
  957. if row1[1] <= split_y[i-1] or row1[1] >= split_y[i]:
  958. continue
  959. # bbox上边框 y一样
  960. if not contain_flag1:
  961. if row1[1]-3 <= candidate_bbox[1] <= row1[1]+3:
  962. # candidate的x1,x2需被包含在row线中
  963. if row1[0]-3 <= candidate_bbox[0] <= candidate_bbox[2] <= row1[2]+3:
  964. contain_flag1 = 1
  965. # bbox下边框 y一样
  966. if not contain_flag2:
  967. if row1[1]-3 <= candidate_bbox[3] <= row1[1]+3:
  968. # candidate的x1,x2需被包含在row线中
  969. if row1[0]-3 <= candidate_bbox[0] <= candidate_bbox[2] <= row1[2]+3:
  970. contain_flag2 = 1
  971. # 判断该bbox是否存在,判断bbox的左右两条边是否有包含在col中
  972. contain_flag3 = 0
  973. contain_flag4 = 0
  974. for col1 in col_lines:
  975. # 列不在该区域跳过
  976. if col1[1] <= split_y[i-1] or col1[3] >= split_y[i]:
  977. continue
  978. # bbox左边线 x一样
  979. if not contain_flag3:
  980. if col1[0]-3 <= candidate_bbox[0] <= col1[0]+3:
  981. # candidate的y1,y2需被包含在col线中
  982. if col1[1]-3 <= candidate_bbox[1] <= candidate_bbox[3] <= col1[3]+3:
  983. contain_flag3 = 1
  984. # bbox右边框 x一样
  985. if not contain_flag4:
  986. if col1[0]-3 <= candidate_bbox[2] <= col1[0]+3:
  987. # candidate的y1,y2需被包含在col线中
  988. if col1[1]-3 <= candidate_bbox[1] <= candidate_bbox[3] <= col1[3]+3:
  989. contain_flag4 = 1
  990. # 找到了该bbox,并且是存在的
  991. if contain_flag1 and contain_flag2 and contain_flag3 and contain_flag4:
  992. bbox.append([(candidate_bbox[0], candidate_bbox[1]),
  993. (candidate_bbox[2], candidate_bbox[3])])
  994. not_found = 0
  995. break
  996. if not not_found:
  997. break
  998. return bbox
  999. def get_bbox_by_contours(image_np):
  1000. img_gray = cv2.cvtColor(image_np, cv2.COLOR_BGR2GRAY)
  1001. ret, img_bin = cv2.threshold(img_gray, 127, 255, cv2.THRESH_BINARY)
  1002. # 3.连通域分析
  1003. img_bin, contours, hierarchy = cv2.findContours(img_bin,
  1004. cv2.RETR_LIST,
  1005. cv2.CHAIN_APPROX_SIMPLE)
  1006. # 4.获取最小外接圆 圆心 半径
  1007. center, radius = cv2.minEnclosingTriangle(contours[0])
  1008. center = np.int0(center)
  1009. # 5.绘制最小外接圆
  1010. img_result = image_np.copy()
  1011. cv2.circle(img_result, tuple(center), int(radius), (255, 255, 255), 2)
  1012. # # 读入图片
  1013. # img = image_np
  1014. # cv2.imshow("get_bbox_by_contours ", image_np)
  1015. # # 中值滤波,去噪
  1016. # img = cv2.medianBlur(img, 3)
  1017. # gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
  1018. # cv2.namedWindow('original', cv2.WINDOW_AUTOSIZE)
  1019. # cv2.imshow('original', gray)
  1020. #
  1021. # # 阈值分割得到二值化图片
  1022. # ret, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)
  1023. #
  1024. # # 膨胀操作
  1025. # kernel2 = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
  1026. # bin_clo = cv2.dilate(binary, kernel2, iterations=2)
  1027. #
  1028. # # 连通域分析
  1029. # num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(bin_clo, connectivity=8)
  1030. #
  1031. # # 查看各个返回值
  1032. # # 连通域数量
  1033. # print('num_labels = ',num_labels)
  1034. # # 连通域的信息:对应各个轮廓的x、y、width、height和面积
  1035. # print('stats = ',stats)
  1036. # # 连通域的中心点
  1037. # print('centroids = ',centroids)
  1038. # # 每一个像素的标签1、2、3.。。,同一个连通域的标签是一致的
  1039. # print('labels = ',labels)
  1040. #
  1041. # # 不同的连通域赋予不同的颜色
  1042. # output = np.zeros((img.shape[0], img.shape[1], 3), np.uint8)
  1043. # for i in range(1, num_labels):
  1044. #
  1045. # mask = labels == i
  1046. # output[:, :, 0][mask] = np.random.randint(0, 255)
  1047. # output[:, :, 1][mask] = np.random.randint(0, 255)
  1048. # output[:, :, 2][mask] = np.random.randint(0, 255)
  1049. # cv2.imshow('oginal', output)
  1050. # cv2.waitKey()
  1051. # cv2.destroyAllWindows()
  1052. def get_points_col(points, split_y, threshold=5):
  1053. # 坐标点按行分
  1054. row_point_list = []
  1055. row_point = []
  1056. points.sort(key=lambda x: (x[0], x[1]))
  1057. x = points[0][0]
  1058. for i in range(1, len(split_y)):
  1059. for p in points:
  1060. if p[1] <= split_y[i-1] or p[1] >= split_y[i]:
  1061. continue
  1062. if x-threshold <= p[0] <= x+threshold:
  1063. row_point.append(p)
  1064. else:
  1065. row_point.sort(key=lambda x: (x[1], x[0]))
  1066. if row_point:
  1067. row_point_list.append(row_point)
  1068. row_point = []
  1069. x = p[0]
  1070. row_point.append(p)
  1071. if row_point:
  1072. row_point_list.append(row_point)
  1073. return row_point_list
  1074. def get_points_row(points, split_y, threshold=5):
  1075. # 坐标点按列分
  1076. col_point_list = []
  1077. col_point = []
  1078. points.sort(key=lambda x: (x[1], x[0]))
  1079. y = points[0][1]
  1080. for i in range(len(split_y)):
  1081. for p in points:
  1082. if p[1] <= split_y[i-1] or p[1] >= split_y[i]:
  1083. continue
  1084. if y-threshold <= p[1] <= y+threshold:
  1085. col_point.append(p)
  1086. else:
  1087. col_point.sort(key=lambda x: (x[0], x[1]))
  1088. if col_point:
  1089. col_point_list.append(col_point)
  1090. col_point = []
  1091. y = p[1]
  1092. col_point.append(p)
  1093. if col_point:
  1094. col_point_list.append(col_point)
  1095. return col_point_list
  1096. def get_outline_point(points, split_y):
  1097. # 分割线纵坐标
  1098. # print("get_outline_point split_y", split_y)
  1099. if len(split_y) < 2:
  1100. return []
  1101. outline_2point = []
  1102. points.sort(key=lambda x: (x[1], x[0]))
  1103. for i in range(1, len(split_y)):
  1104. area_points = []
  1105. for point in points:
  1106. if point[1] <= split_y[i-1] or point[1] >= split_y[i]:
  1107. continue
  1108. area_points.append(point)
  1109. if area_points:
  1110. area_points.sort(key=lambda x: (x[1], x[0]))
  1111. outline_2point.append([area_points[0], area_points[-1]])
  1112. return outline_2point
  1113. # def merge_row(row_lines):
  1114. # for row in row_lines:
  1115. # for row1 in row_lines:
  1116. def get_best_predict_size(image_np):
  1117. sizes = [1280, 1152, 1024, 896, 768, 640, 512, 384, 256, 128]
  1118. min_len = 10000
  1119. best_height = sizes[0]
  1120. for height in sizes:
  1121. if abs(image_np.shape[0] - height) < min_len:
  1122. min_len = abs(image_np.shape[0] - height)
  1123. best_height = height
  1124. min_len = 10000
  1125. best_width = sizes[0]
  1126. for width in sizes:
  1127. if abs(image_np.shape[1] - width) < min_len:
  1128. min_len = abs(image_np.shape[1] - width)
  1129. best_width = width
  1130. return best_height, best_width
  1131. def choose_longer_row(lines):
  1132. new_row = []
  1133. jump_row = []
  1134. for i in range(len(lines)):
  1135. row1 = lines[i]
  1136. jump_flag = 0
  1137. if row1 in jump_row:
  1138. continue
  1139. for j in range(i+1, len(lines)):
  1140. row2 = lines[j]
  1141. if row2 in jump_row:
  1142. continue
  1143. if row2[1]-5 <= row1[1] <= row2[1]+5:
  1144. if row1[0] <= row2[0] and row1[2] >= row2[2]:
  1145. new_row.append(row1)
  1146. jump_row.append(row1)
  1147. jump_row.append(row2)
  1148. jump_flag = 1
  1149. break
  1150. elif row2[0] <= row1[0] and row2[2] >= row1[2]:
  1151. new_row.append(row2)
  1152. jump_row.append(row1)
  1153. jump_row.append(row2)
  1154. jump_flag = 1
  1155. break
  1156. if not jump_flag:
  1157. new_row.append(row1)
  1158. jump_row.append(row1)
  1159. return new_row
  1160. def choose_longer_col(lines):
  1161. new_col = []
  1162. jump_col = []
  1163. for i in range(len(lines)):
  1164. col1 = lines[i]
  1165. jump_flag = 0
  1166. if col1 in jump_col:
  1167. continue
  1168. for j in range(i+1, len(lines)):
  1169. col2 = lines[j]
  1170. if col2 in jump_col:
  1171. continue
  1172. if col2[0]-5 <= col1[0] <= col2[0]+5:
  1173. if col1[1] <= col2[1] and col1[3] >= col2[3]:
  1174. new_col.append(col1)
  1175. jump_col.append(col1)
  1176. jump_col.append(col2)
  1177. jump_flag = 1
  1178. break
  1179. elif col2[1] <= col1[1] and col2[3] >= col1[3]:
  1180. new_col.append(col2)
  1181. jump_col.append(col1)
  1182. jump_col.append(col2)
  1183. jump_flag = 1
  1184. break
  1185. if not jump_flag:
  1186. new_col.append(col1)
  1187. jump_col.append(col1)
  1188. return new_col
  1189. if __name__ == '__main__':
  1190. # p = "开标记录表3_page_0.png"
  1191. # p = "train_data/label_1.jpg"
  1192. # p = "train_463.jpg"
  1193. p = "8.png"
  1194. # p = "无边框3.jpg"
  1195. # p = "part1.png"
  1196. # p = "D:\\Project\\format_conversion\\appendix_test\\temp\\00e959a0bc9011ebaf5a00163e0ae709" + \
  1197. # "\\00e95f7cbc9011ebaf5a00163e0ae709_pdf_page0.png"
  1198. # p = "D:\\Project\\format_conversion\\appendix_test\\temp\\00fb3e52bc7e11eb836000163e0ae709" + \
  1199. # "\\00fb43acbc7e11eb836000163e0ae709.png"
  1200. # p = "table.jpg"
  1201. # p = "data_process/create_data/0.jpg"
  1202. img = cv2.imread(p)
  1203. t = time.time()
  1204. model.load_weights(tableModeLinePath)
  1205. best_h, best_w = get_best_predict_size(img)
  1206. print(img.shape)
  1207. print((best_h, best_w))
  1208. # row_boxes, col_boxes = table_line(img[..., ::-1], model, size=(512, 1024), hprob=0.5, vprob=0.5)
  1209. # row_boxes, col_boxes, img = table_line(img[..., ::-1], model, size=(best_w, best_h), hprob=0.5, vprob=0.5)
  1210. row_boxes, col_boxes, img = table_line(img, model, size=(best_w, best_h), hprob=0.5, vprob=0.5)
  1211. # 创建空图
  1212. test_img = np.zeros((img.shape), np.uint8)
  1213. test_img.fill(255)
  1214. for box in row_boxes+col_boxes:
  1215. cv2.line(test_img, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])), (0, 0, 0), 1)
  1216. cv2.imshow("test_image", test_img)
  1217. cv2.waitKey(0)
  1218. cv2.imwrite("temp.jpg", test_img)
  1219. # 计算交点、分割线
  1220. crossover_points = get_points(row_boxes, col_boxes, (img.shape[0], img.shape[1]))
  1221. print("len(col_boxes)", len(col_boxes))
  1222. split_lines, split_y = get_split_line(crossover_points, col_boxes, img)
  1223. print("split_y", split_y)
  1224. # for point in crossover_points:
  1225. # cv2.circle(test_img, point, 1, (0, 255, 0), 3)
  1226. # cv2.imshow("point image1", test_img)
  1227. # cv2.waitKey(0)
  1228. # 计算行列,剔除相近交点
  1229. row_point_list = get_points_row(crossover_points, split_y, 0)
  1230. col_point_list = get_points_col(crossover_points, split_y, 0)
  1231. print("1", len(crossover_points), len(row_point_list), len(col_point_list))
  1232. crossover_points = delete_close_points(crossover_points, row_point_list, col_point_list)
  1233. row_point_list = get_points_row(crossover_points, split_y)
  1234. col_point_list = get_points_col(crossover_points, split_y)
  1235. print("2", len(crossover_points), len(row_point_list), len(col_point_list))
  1236. for point in crossover_points:
  1237. cv2.circle(test_img, point, 1, (0, 0, 255), 3)
  1238. cv2.imshow("point image1", test_img)
  1239. cv2.waitKey(0)
  1240. print("len(row_boxes)", len(row_boxes))
  1241. print("len(col_boxes)", len(col_boxes))
  1242. # 修复边框
  1243. new_row_boxes, new_col_boxes, long_row_boxes, long_col_boxes = \
  1244. fix_outline(img, row_boxes, col_boxes, crossover_points, split_y)
  1245. if new_row_boxes or new_col_boxes:
  1246. if long_row_boxes:
  1247. print("long_row_boxes", long_row_boxes)
  1248. row_boxes = long_row_boxes
  1249. if long_col_boxes:
  1250. print("long_col_boxes", long_col_boxes)
  1251. col_boxes = long_col_boxes
  1252. if new_row_boxes:
  1253. row_boxes += new_row_boxes
  1254. if new_col_boxes:
  1255. col_boxes += new_col_boxes
  1256. # row_boxes += new_row_boxes
  1257. # col_boxes += new_col_boxes
  1258. # row_boxes = choose_longer_row(row_boxes)
  1259. # col_boxes = choose_longer_col(col_boxes)
  1260. # 创建空图
  1261. test_img = np.zeros((img.shape), np.uint8)
  1262. test_img.fill(255)
  1263. for box in row_boxes+col_boxes:
  1264. cv2.line(test_img, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])), (0, 0, 0), 1)
  1265. cv2.imshow("test_image2", test_img)
  1266. cv2.waitKey(0)
  1267. # 展示补线
  1268. for row in new_row_boxes:
  1269. cv2.line(test_img, (int(row[0]), int(row[1])),
  1270. (int(row[2]), int(row[3])), (0, 0, 255), 1)
  1271. for col in new_col_boxes:
  1272. cv2.line(test_img, (int(col[0]), int(col[1])),
  1273. (int(col[2]), int(col[3])), (0, 0, 255), 1)
  1274. cv2.imshow("fix_outline", test_img)
  1275. cv2.waitKey(0)
  1276. cv2.imwrite("temp.jpg", test_img)
  1277. # 修复边框后重新计算交点、分割线
  1278. print("crossover_points", len(crossover_points))
  1279. crossover_points = get_points(row_boxes, col_boxes, (img.shape[0], img.shape[1]))
  1280. print("crossover_points new", len(crossover_points))
  1281. split_lines, split_y = get_split_line(crossover_points, col_boxes, img)
  1282. # 计算行列,剔除相近交点
  1283. row_point_list = get_points_row(crossover_points, split_y, 0)
  1284. col_point_list = get_points_col(crossover_points, split_y, 0)
  1285. print(len(crossover_points), len(row_point_list), len(col_point_list))
  1286. crossover_points = delete_close_points(crossover_points, row_point_list, col_point_list)
  1287. print(len(crossover_points), len(row_point_list), len(col_point_list))
  1288. row_point_list = get_points_row(crossover_points, split_y)
  1289. col_point_list = get_points_col(crossover_points, split_y)
  1290. for point in crossover_points:
  1291. cv2.circle(test_img, point, 1, (0, 255, 0), 3)
  1292. cv2.imshow("point image2", test_img)
  1293. cv2.waitKey(0)
  1294. # 获取每个表格的左上右下两个点
  1295. outline_point = get_outline_point(crossover_points, split_y)
  1296. # print(outline_point)
  1297. for outline in outline_point:
  1298. cv2.circle(test_img, outline[0], 1, (255, 0, 0), 5)
  1299. cv2.circle(test_img, outline[1], 1, (255, 0, 0), 5)
  1300. cv2.imshow("outline point", test_img)
  1301. cv2.waitKey(0)
  1302. # 获取bbox
  1303. # get_bbox(img, crossover_points, split_y)
  1304. # for point in crossover_points:
  1305. # cv2.circle(test_img, point, 1, (0, 255, 0), 3)
  1306. # cv2.imshow("point image3", test_img)
  1307. # cv2.waitKey(0)
  1308. # split_y = []
  1309. # for outline in outline_point:
  1310. # split_y.extend([outline[0][1]-5, outline[1][1]+5])
  1311. print("len(row_boxes)", len(row_boxes))
  1312. print("len(col_boxes)", len(col_boxes))
  1313. bboxes = get_bbox(img, row_point_list, col_point_list, split_y, row_boxes, col_boxes)
  1314. # 展示
  1315. for box in bboxes:
  1316. # print(box[0], box[1])
  1317. # if abs(box[0][1] - box[1][1]) > abs(box[0][0] - box[1][0]):
  1318. # continue
  1319. cv2.rectangle(test_img, box[0], box[1], (0, 0, 255), 2, 8)
  1320. cv2.imshow('bboxes', test_img)
  1321. cv2.waitKey(0)
  1322. # img = draw_lines(img, row_boxes+col_boxes, color=(255, 0, 0), lineW=2)
  1323. # img = draw_boxes(img, rowboxes+colboxes, color=(0, 0, 255))
  1324. print(time.time()-t, len(row_boxes), len(col_boxes))
  1325. # cv2.imwrite('temp.jpg', test_img)
  1326. # cv2.imshow('main', img)
  1327. # cv2.waitKey(0)