table_line_rec_utils.py 22 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676
  1. # -*- encoding: utf-8 -*-
  2. # @Author: SWHL
  3. # @Contact: liekkaskono@163.com
  4. import copy
  5. import math
  6. import cv2
  7. import numpy as np
  8. from scipy.spatial import distance as dist
  9. from skimage import measure
  10. def bbox_decode(heat, wh, reg=None, K=100):
  11. """bbox组成:[V1, V2, V3, V4]
  12. V1~V4: bbox的4个坐标点
  13. """
  14. batch = heat.shape[0]
  15. heat, keep = _nms(heat)
  16. scores, inds, clses, ys, xs = _topk(heat, K=K)
  17. if reg is not None:
  18. reg = _tranpose_and_gather_feat(reg, inds)
  19. reg = reg.reshape(batch, K, 2)
  20. xs = xs.reshape(batch, K, 1) + reg[:, :, 0:1]
  21. ys = ys.reshape(batch, K, 1) + reg[:, :, 1:2]
  22. else:
  23. xs = xs.reshape(batch, K, 1) + 0.5
  24. ys = ys.reshape(batch, K, 1) + 0.5
  25. wh = _tranpose_and_gather_feat(wh, inds)
  26. wh = wh.reshape(batch, K, 8)
  27. clses = clses.reshape(batch, K, 1).astype(np.float32)
  28. scores = scores.reshape(batch, K, 1)
  29. bboxes = np.concatenate(
  30. [
  31. xs - wh[..., 0:1],
  32. ys - wh[..., 1:2],
  33. xs - wh[..., 2:3],
  34. ys - wh[..., 3:4],
  35. xs - wh[..., 4:5],
  36. ys - wh[..., 5:6],
  37. xs - wh[..., 6:7],
  38. ys - wh[..., 7:8],
  39. ],
  40. axis=2,
  41. )
  42. detections = np.concatenate([bboxes, scores, clses], axis=2)
  43. return detections, inds
  44. def _nms(heat, kernel=3):
  45. pad = (kernel - 1) // 2
  46. hmax = max_pool(heat, kernel_size=kernel, stride=1, padding=pad)
  47. keep = hmax == heat
  48. return heat * keep, keep
  49. def max_pool(img, kernel_size, stride, padding):
  50. h, w = img.shape[2:]
  51. img = np.pad(
  52. img,
  53. ((0, 0), (0, 0), (padding, padding), (padding, padding)),
  54. "constant",
  55. constant_values=0,
  56. )
  57. res_h = ((h + 2 - kernel_size) // stride) + 1
  58. res_w = ((w + 2 - kernel_size) // stride) + 1
  59. res = np.zeros((img.shape[0], img.shape[1], res_h, res_w))
  60. for i in range(res_h):
  61. for j in range(res_w):
  62. temp = img[
  63. :,
  64. :,
  65. i * stride : i * stride + kernel_size,
  66. j * stride : j * stride + kernel_size,
  67. ]
  68. res[:, :, i, j] = temp.max()
  69. return res
  70. def _topk(scores, K=40):
  71. batch, cat, height, width = scores.shape
  72. topk_scores, topk_inds = find_topk(scores.reshape(batch, cat, -1), K)
  73. topk_inds = topk_inds % (height * width)
  74. topk_ys = topk_inds / width
  75. topk_xs = np.float32(np.int32(topk_inds % width))
  76. topk_score, topk_ind = find_topk(topk_scores.reshape(batch, -1), K)
  77. topk_clses = np.int32(topk_ind / K)
  78. topk_inds = _gather_feat(topk_inds.reshape(batch, -1, 1), topk_ind).reshape(
  79. batch, K
  80. )
  81. topk_ys = _gather_feat(topk_ys.reshape(batch, -1, 1), topk_ind).reshape(batch, K)
  82. topk_xs = _gather_feat(topk_xs.reshape(batch, -1, 1), topk_ind).reshape(batch, K)
  83. return topk_score, topk_inds, topk_clses, topk_ys, topk_xs
  84. def find_topk(a, k, axis=-1, largest=True, sorted=True):
  85. if axis is None:
  86. axis_size = a.size
  87. else:
  88. axis_size = a.shape[axis]
  89. assert 1 <= k <= axis_size
  90. a = np.asanyarray(a)
  91. if largest:
  92. index_array = np.argpartition(a, axis_size - k, axis=axis)
  93. topk_indices = np.take(index_array, -np.arange(k) - 1, axis=axis)
  94. else:
  95. index_array = np.argpartition(a, k - 1, axis=axis)
  96. topk_indices = np.take(index_array, np.arange(k), axis=axis)
  97. topk_values = np.take_along_axis(a, topk_indices, axis=axis)
  98. if sorted:
  99. sorted_indices_in_topk = np.argsort(topk_values, axis=axis)
  100. if largest:
  101. sorted_indices_in_topk = np.flip(sorted_indices_in_topk, axis=axis)
  102. sorted_topk_values = np.take_along_axis(
  103. topk_values, sorted_indices_in_topk, axis=axis
  104. )
  105. sorted_topk_indices = np.take_along_axis(
  106. topk_indices, sorted_indices_in_topk, axis=axis
  107. )
  108. return sorted_topk_values, sorted_topk_indices
  109. return topk_values, topk_indices
  110. def _gather_feat(feat, ind):
  111. dim = feat.shape[2]
  112. ind = np.broadcast_to(ind[:, :, None], (ind.shape[0], ind.shape[1], dim))
  113. feat = _gather_np(feat, 1, ind)
  114. return feat
  115. def _gather_np(data, dim, index):
  116. """
  117. Gathers values along an axis specified by dim.
  118. For a 3-D tensor the output is specified by:
  119. out[i][j][k] = input[index[i][j][k]][j][k] # if dim == 0
  120. out[i][j][k] = input[i][index[i][j][k]][k] # if dim == 1
  121. out[i][j][k] = input[i][j][index[i][j][k]] # if dim == 2
  122. :param dim: The axis along which to index
  123. :param index: A tensor of indices of elements to gather
  124. :return: tensor of gathered values
  125. """
  126. idx_xsection_shape = index.shape[:dim] + index.shape[dim + 1 :]
  127. data_xsection_shape = data.shape[:dim] + data.shape[dim + 1 :]
  128. if idx_xsection_shape != data_xsection_shape:
  129. raise ValueError(
  130. "Except for dimension "
  131. + str(dim)
  132. + ", all dimensions of index and data should be the same size"
  133. )
  134. if index.dtype != np.int64:
  135. raise TypeError("The values of index must be integers")
  136. data_swaped = np.swapaxes(data, 0, dim)
  137. index_swaped = np.swapaxes(index, 0, dim)
  138. gathered = np.take_along_axis(data_swaped, index_swaped, axis=0)
  139. return np.swapaxes(gathered, 0, dim)
  140. def _tranpose_and_gather_feat(feat, ind):
  141. feat = np.ascontiguousarray(np.transpose(feat, [0, 2, 3, 1]))
  142. feat = feat.reshape(feat.shape[0], -1, feat.shape[3])
  143. feat = _gather_feat(feat, ind)
  144. return feat
  145. def gbox_decode(mk, st_reg, reg=None, K=400):
  146. """gbox的组成:[V1, P1, P2, P3, P4]
  147. P1~P4: 四个框的中心点
  148. V1: 四个框的交点
  149. """
  150. batch = mk.shape[0]
  151. mk, keep = _nms(mk)
  152. scores, inds, clses, ys, xs = _topk(mk, K=K)
  153. if reg is not None:
  154. reg = _tranpose_and_gather_feat(reg, inds)
  155. reg = reg.reshape(batch, K, 2)
  156. xs = xs.reshape(batch, K, 1) + reg[:, :, 0:1]
  157. ys = ys.reshape(batch, K, 1) + reg[:, :, 1:2]
  158. else:
  159. xs = xs.reshape(batch, K, 1) + 0.5
  160. ys = ys.reshape(batch, K, 1) + 0.5
  161. scores = scores.reshape(batch, K, 1)
  162. clses = clses.reshape(batch, K, 1).astype(np.float32)
  163. st_Reg = _tranpose_and_gather_feat(st_reg, inds)
  164. bboxes = np.concatenate(
  165. [
  166. xs - st_Reg[..., 0:1],
  167. ys - st_Reg[..., 1:2],
  168. xs - st_Reg[..., 2:3],
  169. ys - st_Reg[..., 3:4],
  170. xs - st_Reg[..., 4:5],
  171. ys - st_Reg[..., 5:6],
  172. xs - st_Reg[..., 6:7],
  173. ys - st_Reg[..., 7:8],
  174. ],
  175. axis=2,
  176. )
  177. return np.concatenate([xs, ys, bboxes, scores, clses], axis=2), keep
  178. def transform_preds(coords, center, scale, output_size, rot=0):
  179. target_coords = np.zeros(coords.shape)
  180. trans = get_affine_transform(center, scale, rot, output_size, inv=1)
  181. for p in range(coords.shape[0]):
  182. target_coords[p, 0:2] = affine_transform(coords[p, 0:2], trans)
  183. return target_coords
  184. def get_affine_transform(
  185. center, scale, rot, output_size, shift=np.array([0, 0], dtype=np.float32), inv=0
  186. ):
  187. if not isinstance(scale, np.ndarray) and not isinstance(scale, list):
  188. scale = np.array([scale, scale], dtype=np.float32)
  189. scale_tmp = scale
  190. src_w = scale_tmp[0]
  191. dst_w = output_size[0]
  192. dst_h = output_size[1]
  193. rot_rad = np.pi * rot / 180
  194. src_dir = get_dir([0, src_w * -0.5], rot_rad)
  195. dst_dir = np.array([0, dst_w * -0.5], np.float32)
  196. src = np.zeros((3, 2), dtype=np.float32)
  197. dst = np.zeros((3, 2), dtype=np.float32)
  198. src[0, :] = center + scale_tmp * shift
  199. src[1, :] = center + src_dir + scale_tmp * shift
  200. dst[0, :] = [dst_w * 0.5, dst_h * 0.5]
  201. dst[1, :] = np.array([dst_w * 0.5, dst_h * 0.5], np.float32) + dst_dir
  202. src[2:, :] = get_3rd_point(src[0, :], src[1, :])
  203. dst[2:, :] = get_3rd_point(dst[0, :], dst[1, :])
  204. if inv:
  205. trans = cv2.getAffineTransform(np.float32(dst), np.float32(src))
  206. else:
  207. trans = cv2.getAffineTransform(np.float32(src), np.float32(dst))
  208. return trans
  209. def affine_transform(pt, t):
  210. new_pt = np.array([pt[0], pt[1], 1.0], dtype=np.float32).T
  211. new_pt = np.dot(t, new_pt)
  212. return new_pt[:2]
  213. def get_dir(src_point, rot_rad):
  214. sn, cs = np.sin(rot_rad), np.cos(rot_rad)
  215. src_result = [0, 0]
  216. src_result[0] = src_point[0] * cs - src_point[1] * sn
  217. src_result[1] = src_point[0] * sn + src_point[1] * cs
  218. return src_result
  219. def get_3rd_point(a, b):
  220. direct = a - b
  221. return b + np.array([-direct[1], direct[0]], dtype=np.float32)
  222. def bbox_post_process(bbox, c, s, h, w):
  223. for i in range(bbox.shape[0]):
  224. bbox[i, :, 0:2] = transform_preds(bbox[i, :, 0:2], c[i], s[i], (w, h))
  225. bbox[i, :, 2:4] = transform_preds(bbox[i, :, 2:4], c[i], s[i], (w, h))
  226. bbox[i, :, 4:6] = transform_preds(bbox[i, :, 4:6], c[i], s[i], (w, h))
  227. bbox[i, :, 6:8] = transform_preds(bbox[i, :, 6:8], c[i], s[i], (w, h))
  228. return bbox
  229. def gbox_post_process(gbox, c, s, h, w):
  230. for i in range(gbox.shape[0]):
  231. gbox[i, :, 0:2] = transform_preds(gbox[i, :, 0:2], c[i], s[i], (w, h))
  232. gbox[i, :, 2:4] = transform_preds(gbox[i, :, 2:4], c[i], s[i], (w, h))
  233. gbox[i, :, 4:6] = transform_preds(gbox[i, :, 4:6], c[i], s[i], (w, h))
  234. gbox[i, :, 6:8] = transform_preds(gbox[i, :, 6:8], c[i], s[i], (w, h))
  235. gbox[i, :, 8:10] = transform_preds(gbox[i, :, 8:10], c[i], s[i], (w, h))
  236. return gbox
  237. def nms(dets, thresh):
  238. if len(dets) < 2:
  239. return dets
  240. index_keep, keep = [], []
  241. for i in range(len(dets)):
  242. box = dets[i]
  243. if box[-1] < thresh:
  244. break
  245. max_score_index = -1
  246. ctx = (dets[i][0] + dets[i][2] + dets[i][4] + dets[i][6]) / 4
  247. cty = (dets[i][1] + dets[i][3] + dets[i][5] + dets[i][7]) / 4
  248. for j in range(len(dets)):
  249. if i == j or dets[j][-1] < thresh:
  250. break
  251. x1, y1 = dets[j][0], dets[j][1]
  252. x2, y2 = dets[j][2], dets[j][3]
  253. x3, y3 = dets[j][4], dets[j][5]
  254. x4, y4 = dets[j][6], dets[j][7]
  255. a = (x2 - x1) * (cty - y1) - (y2 - y1) * (ctx - x1)
  256. b = (x3 - x2) * (cty - y2) - (y3 - y2) * (ctx - x2)
  257. c = (x4 - x3) * (cty - y3) - (y4 - y3) * (ctx - x3)
  258. d = (x1 - x4) * (cty - y4) - (y1 - y4) * (ctx - x4)
  259. if all(x > 0 for x in (a, b, c, d)) or all(x < 0 for x in (a, b, c, d)):
  260. if dets[i][8] > dets[j][8] and max_score_index < 0:
  261. max_score_index = i
  262. elif dets[i][8] < dets[j][8]:
  263. max_score_index = -2
  264. break
  265. if max_score_index > -1:
  266. index_keep.append(max_score_index)
  267. elif max_score_index == -1:
  268. index_keep.append(i)
  269. keep = [dets[index_keep[i]] for i in range(len(index_keep))]
  270. return np.array(keep)
  271. def group_bbox_by_gbox(
  272. bboxes, gboxes, score_thred=0.3, v2c_dist_thred=2, c2v_dist_thred=0.5
  273. ):
  274. def point_in_box(box, point):
  275. x1, y1, x2, y2 = box[0], box[1], box[2], box[3]
  276. x3, y3, x4, y4 = box[4], box[5], box[6], box[7]
  277. ctx, cty = point[0], point[1]
  278. a = (x2 - x1) * (cty - y1) - (y2 - y1) * (ctx - x1)
  279. b = (x3 - x2) * (cty - y2) - (y3 - y2) * (ctx - x2)
  280. c = (x4 - x3) * (cty - y3) - (y4 - y3) * (ctx - x3)
  281. d = (x1 - x4) * (cty - y4) - (y1 - y4) * (ctx - x4)
  282. if all(x > 0 for x in (a, b, c, d)) or all(x < 0 for x in (a, b, c, d)):
  283. return True
  284. return False
  285. def get_distance(pt1, pt2):
  286. return math.sqrt(
  287. (pt1[0] - pt2[0]) * (pt1[0] - pt2[0])
  288. + (pt1[1] - pt2[1]) * (pt1[1] - pt2[1])
  289. )
  290. dets = copy.deepcopy(bboxes)
  291. sign = np.zeros((len(dets), 4))
  292. for gbox in gboxes:
  293. if gbox[10] < score_thred:
  294. break
  295. vertex = [gbox[0], gbox[1]]
  296. for i in range(4):
  297. center = [gbox[2 * i + 2], gbox[2 * i + 3]]
  298. if get_distance(vertex, center) < v2c_dist_thred:
  299. continue
  300. for k, bbox in enumerate(dets):
  301. if bbox[8] < score_thred:
  302. break
  303. if sum(sign[k]) == 4:
  304. continue
  305. w = (abs(bbox[6] - bbox[0]) + abs(bbox[4] - bbox[2])) / 2
  306. h = (abs(bbox[3] - bbox[1]) + abs(bbox[5] - bbox[7])) / 2
  307. m = max(w, h)
  308. if point_in_box(bbox, center):
  309. min_dist, min_id = 1e4, -1
  310. for j in range(4):
  311. dist = get_distance(vertex, [bbox[2 * j], bbox[2 * j + 1]])
  312. if dist < min_dist:
  313. min_dist = dist
  314. min_id = j
  315. if (
  316. min_id > -1
  317. and min_dist < c2v_dist_thred * m
  318. and sign[k][min_id] == 0
  319. ):
  320. bboxes[k][2 * min_id] = vertex[0]
  321. bboxes[k][2 * min_id + 1] = vertex[1]
  322. sign[k][min_id] = 1
  323. return bboxes
  324. def get_table_line(binimg, axis=0, lineW=10):
  325. ##获取表格线
  326. ##axis=0 横线
  327. ##axis=1 竖线
  328. labels = measure.label(binimg > 0, connectivity=2) # 8连通区域标记
  329. regions = measure.regionprops(labels)
  330. if axis == 1:
  331. lineboxes = [
  332. min_area_rect(line.coords)
  333. for line in regions
  334. if line.bbox[2] - line.bbox[0] > lineW
  335. ]
  336. else:
  337. lineboxes = [
  338. min_area_rect(line.coords)
  339. for line in regions
  340. if line.bbox[3] - line.bbox[1] > lineW
  341. ]
  342. return lineboxes
  343. def min_area_rect(coords):
  344. """
  345. 多边形外接矩形
  346. """
  347. rect = cv2.minAreaRect(coords[:, ::-1])
  348. box = cv2.boxPoints(rect)
  349. box = box.reshape((8,)).tolist()
  350. box = image_location_sort_box(box)
  351. x1, y1, x2, y2, x3, y3, x4, y4 = box
  352. degree, w, h, cx, cy = calculate_center_rotate_angle(box)
  353. if w < h:
  354. xmin = (x1 + x2) / 2
  355. xmax = (x3 + x4) / 2
  356. ymin = (y1 + y2) / 2
  357. ymax = (y3 + y4) / 2
  358. else:
  359. xmin = (x1 + x4) / 2
  360. xmax = (x2 + x3) / 2
  361. ymin = (y1 + y4) / 2
  362. ymax = (y2 + y3) / 2
  363. # degree,w,h,cx,cy = solve(box)
  364. # x1,y1,x2,y2,x3,y3,x4,y4 = box
  365. # return {'degree':degree,'w':w,'h':h,'cx':cx,'cy':cy}
  366. return [xmin, ymin, xmax, ymax]
  367. def image_location_sort_box(box):
  368. x1, y1, x2, y2, x3, y3, x4, y4 = box[:8]
  369. pts = (x1, y1), (x2, y2), (x3, y3), (x4, y4)
  370. pts = np.array(pts, dtype="float32")
  371. (x1, y1), (x2, y2), (x3, y3), (x4, y4) = _order_points(pts)
  372. return [x1, y1, x2, y2, x3, y3, x4, y4]
  373. def calculate_center_rotate_angle(box):
  374. """
  375. 绕 cx,cy点 w,h 旋转 angle 的坐标,能一定程度缓解图片的内部倾斜,但是还是依赖模型稳妥
  376. x = cx-w/2
  377. y = cy-h/2
  378. x1-cx = -w/2*cos(angle) +h/2*sin(angle)
  379. y1 -cy= -w/2*sin(angle) -h/2*cos(angle)
  380. h(x1-cx) = -wh/2*cos(angle) +hh/2*sin(angle)
  381. w(y1 -cy)= -ww/2*sin(angle) -hw/2*cos(angle)
  382. (hh+ww)/2sin(angle) = h(x1-cx)-w(y1 -cy)
  383. """
  384. x1, y1, x2, y2, x3, y3, x4, y4 = box[:8]
  385. cx = (x1 + x3 + x2 + x4) / 4.0
  386. cy = (y1 + y3 + y4 + y2) / 4.0
  387. w = (
  388. np.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)
  389. + np.sqrt((x3 - x4) ** 2 + (y3 - y4) ** 2)
  390. ) / 2
  391. h = (
  392. np.sqrt((x2 - x3) ** 2 + (y2 - y3) ** 2)
  393. + np.sqrt((x1 - x4) ** 2 + (y1 - y4) ** 2)
  394. ) / 2
  395. # x = cx-w/2
  396. # y = cy-h/2
  397. sinA = (h * (x1 - cx) - w * (y1 - cy)) * 1.0 / (h * h + w * w) * 2
  398. angle = np.arcsin(sinA)
  399. return angle, w, h, cx, cy
  400. def _order_points(pts):
  401. # 根据x坐标对点进行排序
  402. """
  403. ---------------------
  404. 本项目中是为了排序后得到[(xmin,ymin),(xmax,ymin),(xmax,ymax),(xmin,ymax)]
  405. 作者:Tong_T
  406. 来源:CSDN
  407. 原文:https://blog.csdn.net/Tong_T/article/details/81907132
  408. 版权声明:本文为博主原创文章,转载请附上博文链接!
  409. """
  410. x_sorted = pts[np.argsort(pts[:, 0]), :]
  411. left_most = x_sorted[:2, :]
  412. right_most = x_sorted[2:, :]
  413. left_most = left_most[np.argsort(left_most[:, 1]), :]
  414. (tl, bl) = left_most
  415. distance = dist.cdist(tl[np.newaxis], right_most, "euclidean")[0]
  416. (br, tr) = right_most[np.argsort(distance)[::-1], :]
  417. return np.array([tl, tr, br, bl], dtype="float32")
  418. def sqrt(p1, p2):
  419. return np.sqrt((p1[0] - p2[0]) ** 2 + (p1[1] - p2[1]) ** 2)
  420. def adjust_lines(lines, alph=50, angle=50):
  421. lines_n = len(lines)
  422. new_lines = []
  423. for i in range(lines_n):
  424. x1, y1, x2, y2 = lines[i]
  425. cx1, cy1 = (x1 + x2) / 2, (y1 + y2) / 2
  426. for j in range(lines_n):
  427. if i != j:
  428. x3, y3, x4, y4 = lines[j]
  429. cx2, cy2 = (x3 + x4) / 2, (y3 + y4) / 2
  430. if (x3 < cx1 < x4 or y3 < cy1 < y4) or (
  431. x1 < cx2 < x2 or y1 < cy2 < y2
  432. ): # 判断两个横线在y方向的投影重不重合
  433. continue
  434. else:
  435. r = sqrt((x1, y1), (x3, y3))
  436. k = abs((y3 - y1) / (x3 - x1 + 1e-10))
  437. a = math.atan(k) * 180 / math.pi
  438. if r < alph and a < angle:
  439. new_lines.append((x1, y1, x3, y3))
  440. r = sqrt((x1, y1), (x4, y4))
  441. k = abs((y4 - y1) / (x4 - x1 + 1e-10))
  442. a = math.atan(k) * 180 / math.pi
  443. if r < alph and a < angle:
  444. new_lines.append((x1, y1, x4, y4))
  445. r = sqrt((x2, y2), (x3, y3))
  446. k = abs((y3 - y2) / (x3 - x2 + 1e-10))
  447. a = math.atan(k) * 180 / math.pi
  448. if r < alph and a < angle:
  449. new_lines.append((x2, y2, x3, y3))
  450. r = sqrt((x2, y2), (x4, y4))
  451. k = abs((y4 - y2) / (x4 - x2 + 1e-10))
  452. a = math.atan(k) * 180 / math.pi
  453. if r < alph and a < angle:
  454. new_lines.append((x2, y2, x4, y4))
  455. return new_lines
  456. def final_adjust_lines(rowboxes, colboxes):
  457. nrow = len(rowboxes)
  458. ncol = len(colboxes)
  459. for i in range(nrow):
  460. for j in range(ncol):
  461. rowboxes[i] = line_to_line(rowboxes[i], colboxes[j], alpha=20, angle=30)
  462. colboxes[j] = line_to_line(colboxes[j], rowboxes[i], alpha=20, angle=30)
  463. return rowboxes, colboxes
  464. def draw_lines(im, bboxes, color=(0, 0, 0), lineW=3):
  465. """
  466. boxes: bounding boxes
  467. """
  468. tmp = np.copy(im)
  469. c = color
  470. h, w = im.shape[:2]
  471. for box in bboxes:
  472. x1, y1, x2, y2 = box[:4]
  473. cv2.line(
  474. tmp, (int(x1), int(y1)), (int(x2), int(y2)), c, lineW, lineType=cv2.LINE_AA
  475. )
  476. return tmp
  477. def line_to_line(points1, points2, alpha=10, angle=30):
  478. """
  479. 线段之间的距离
  480. """
  481. x1, y1, x2, y2 = points1
  482. ox1, oy1, ox2, oy2 = points2
  483. xy = np.array([(x1, y1), (x2, y2)], dtype="float32")
  484. A1, B1, C1 = fit_line(xy)
  485. oxy = np.array([(ox1, oy1), (ox2, oy2)], dtype="float32")
  486. A2, B2, C2 = fit_line(oxy)
  487. flag1 = point_line_cor(np.array([x1, y1], dtype="float32"), A2, B2, C2)
  488. flag2 = point_line_cor(np.array([x2, y2], dtype="float32"), A2, B2, C2)
  489. if (flag1 > 0 and flag2 > 0) or (flag1 < 0 and flag2 < 0): # 横线或者竖线在竖线或者横线的同一侧
  490. if (A1 * B2 - A2 * B1) != 0:
  491. x = (B1 * C2 - B2 * C1) / (A1 * B2 - A2 * B1)
  492. y = (A2 * C1 - A1 * C2) / (A1 * B2 - A2 * B1)
  493. # x, y = round(x, 2), round(y, 2)
  494. p = (x, y) # 横线与竖线的交点
  495. r0 = sqrt(p, (x1, y1))
  496. r1 = sqrt(p, (x2, y2))
  497. if min(r0, r1) < alpha: # 若交点与线起点或者终点的距离小于alpha,则延长线到交点
  498. if r0 < r1:
  499. k = abs((y2 - p[1]) / (x2 - p[0] + 1e-10))
  500. a = math.atan(k) * 180 / math.pi
  501. if a < angle or abs(90 - a) < angle:
  502. points1 = np.array([p[0], p[1], x2, y2], dtype="float32")
  503. else:
  504. k = abs((y1 - p[1]) / (x1 - p[0] + 1e-10))
  505. a = math.atan(k) * 180 / math.pi
  506. if a < angle or abs(90 - a) < angle:
  507. points1 = np.array([x1, y1, p[0], p[1]], dtype="float32")
  508. return points1
  509. def min_area_rect_box(
  510. regions, flag=True, W=0, H=0, filtersmall=False, adjust_box=False
  511. ):
  512. """
  513. 多边形外接矩形
  514. """
  515. boxes = []
  516. for region in regions:
  517. if region.bbox_area > H * W * 3 / 4: # 过滤大的单元格
  518. continue
  519. rect = cv2.minAreaRect(region.coords[:, ::-1])
  520. box = cv2.boxPoints(rect)
  521. box = box.reshape((8,)).tolist()
  522. box = image_location_sort_box(box)
  523. x1, y1, x2, y2, x3, y3, x4, y4 = box
  524. angle, w, h, cx, cy = calculate_center_rotate_angle(box)
  525. # if adjustBox:
  526. # x1, y1, x2, y2, x3, y3, x4, y4 = xy_rotate_box(cx, cy, w + 5, h + 5, angle=0, degree=None)
  527. # x1, x4 = max(x1, 0), max(x4, 0)
  528. # y1, y2 = max(y1, 0), max(y2, 0)
  529. # if w > 32 and h > 32 and flag:
  530. # if abs(angle / np.pi * 180) < 20:
  531. # if filtersmall and (w < 10 or h < 10):
  532. # continue
  533. # boxes.append([x1, y1, x2, y2, x3, y3, x4, y4])
  534. # else:
  535. if w * h < 0.5 * W * H:
  536. if filtersmall and (
  537. w < 15 or h < 15
  538. ): # or w / h > 30 or h / w > 30): # 过滤小的单元格
  539. continue
  540. boxes.append([x1, y1, x2, y2, x3, y3, x4, y4])
  541. return boxes
  542. def point_line_cor(p, A, B, C):
  543. ##判断点与线之间的位置关系
  544. # 一般式直线方程(Ax+By+c)=0
  545. x, y = p
  546. r = A * x + B * y + C
  547. return r
  548. def fit_line(p):
  549. """A = Y2 - Y1
  550. B = X1 - X2
  551. C = X2*Y1 - X1*Y2
  552. AX+BY+C=0
  553. 直线一般方程
  554. """
  555. x1, y1 = p[0]
  556. x2, y2 = p[1]
  557. A = y2 - y1
  558. B = x1 - x2
  559. C = x2 * y1 - x1 * y2
  560. return A, B, C