import os import sys import numpy as np import cv2 sys.path.append(os.path.dirname(os.path.abspath(__file__)) + "/../../") from format_convert.utils import log def shrink_bbox(img, bbox_list): def return_first_black_index(image_np): lower = np.array([0, 0, 0]) upper = np.array([150, 150, 150]) mask = cv2.inRange(image_np, lower, upper) black_index_list = np.where(mask != 0) return black_index_list new_bbox_list = [] for bbox in bbox_list: img_bbox = img[int(bbox[0][1]):int(bbox[2][1]), int(bbox[0][0]):int(bbox[2][0]), :] if 0 in img_bbox.shape: new_bbox_list.append(bbox) continue # 左右上下开始扫描,碰到黑像素即停 index_list = return_first_black_index(img_bbox[:, :, :]) if index_list[0].size == 0 or index_list[1].size == 0: new_bbox_list.append(bbox) continue min_h = index_list[0][0] max_h = index_list[0][-1] img_bbox1 = np.swapaxes(img_bbox, 0, 1) index_list = return_first_black_index(img_bbox1[:, :, :]) if index_list[0].size == 0 or index_list[1].size == 0: new_bbox_list.append(bbox) continue min_w = index_list[0][0] max_w = index_list[0][-1] real_min_w = bbox[0][0] + min_w real_max_w = bbox[0][0] + max_w real_min_h = bbox[0][1] + min_h real_max_h = bbox[0][1] + max_h new_bbox = [[real_min_w, real_min_h], [real_min_w, real_max_h], [real_max_w, real_max_h], [real_max_w, real_min_h]] new_bbox_list.append(new_bbox) # cv2.imshow('img', img_bbox) # cv2.imshow('shrink', img[int(new_bbox[0][1]):int(new_bbox[2][1]), int(new_bbox[0][0]):int(new_bbox[2][0]), :]) # cv2.waitKey(0) return new_bbox_list def split_bbox(img, bbox, bbox_text_dict): text = bbox_text_dict.get(str(bbox)) sub_img = img[int(bbox[0][1]):int(bbox[2][1]), int(bbox[0][0]):int(bbox[2][0]), :] split_line_list = [] last_i_status = 1 # 从左到右遍历img for i in range(1, sub_img.shape[1]): # 若这一列黑色像素超过一定值 if np.where(sub_img[:, i, :] < 200)[0].size > sub_img.shape[0]/5: i_status = 0 else: i_status = 1 # 异或,上个像素列为黑且这个像素列为白,或上个像素列为白且这个像素列为黑 if last_i_status ^ i_status: split_line_list.append(int(i)) last_i_status = i_status # 两条分割线太近的去重 min_len = 5 last_l = split_line_list[0] temp_list = [split_line_list[0]] for l in split_line_list[1:]: if l - last_l > min_len: temp_list.append(l) last_l = l split_line_list = temp_list # 若两个分割线间无黑像素,则是应该分割的 split_pair_list = [] last_line = split_line_list[0] for line in split_line_list[1:]: print('last_line, line', last_line, line, np.where(sub_img[:, last_line:line, :] < 100)[0].size) if line - last_line >= 10 and np.where(sub_img[:, last_line:line, :] < 100)[0].size < 10: split_pair_list.append([last_line, line]) last_line = line print('split_pair_list', split_pair_list) for l in split_line_list: l = int(l + bbox[0][0]) cv2.line(img, (l, int(bbox[0][1])), (l, int(bbox[2][1])), (0, 255, 0), 2) cv2.rectangle(img, (int(bbox[0][0]), int(bbox[0][1])), (int(bbox[2][0]), int(bbox[2][1])), (0, 0, 255), 1) cv2.imshow('img', img) cv2.waitKey(0) # 分割得到新bbox split_bbox_list = [] if split_pair_list: start_line = 0 for line1, line2 in split_pair_list: w1 = start_line + bbox[0][0] w2 = line1 + bbox[0][0] start_line = line2 split_bbox_list.append([[w1, bbox[0][1]], [], [w2, bbox[2][1]], []]) w1 = start_line + bbox[0][0] w2 = bbox[2][0] split_bbox_list.append([[w1, bbox[0][1]], [], [w2, bbox[2][1]], []]) print('split_bbox_list', split_bbox_list) # 计算每个字长度 all_len = 0 bbox_len_list = [] for bbox in split_bbox_list: _len = abs(bbox[2][0] - bbox[0][0]) all_len += _len bbox_len_list.append(_len) single_char_len = all_len / len(text) # 根据bbox长度和单字长度比例计算得到截取后的text split_text_list = [] text_start = 0 for _len in bbox_len_list: text_num = int(_len / single_char_len + 0.5) text_end = text_start+text_num if text_end >= len(text): text_end = len(text) split_text_list.append(text[text_start:text_end]) text_start = text_end print('split_text_list', split_text_list) # 更新bbox_text_dict for i, bbox in enumerate(split_bbox_list): bbox_text_dict[str(bbox)] = split_text_list[i] return split_bbox_list, bbox_text_dict def count_black(image_np, threshold=150): lower = np.array([0, 0, 0]) upper = np.array([threshold, threshold, threshold]) mask = cv2.inRange(image_np, lower, upper) cnt = np.sum(mask != 0) # print("count color ", cnt) return cnt def get_points_by_line(img, row_lines, col_lines): row_img = np.zeros_like(img[:, :, 0], dtype=np.uint8) col_img = np.zeros_like(img[:, :, 0], dtype=np.uint8) for r in row_lines: cv2.line(row_img, (r[0], r[1]), (r[2], r[3]), (255, 255, 255), 1) for c in col_lines: cv2.line(col_img, (c[0], c[1]), (c[2], c[3]), (255, 255, 255), 1) point_img = np.bitwise_and(row_img, col_img) # 识别黑白图中的白色交叉点,将横纵坐标取出 ys, xs = np.where(point_img > 0) points = [] for i in range(len(xs)): points.append((xs[i], ys[i])) points.sort(key=lambda x: (x[0], x[1])) return points def get_table_bbox_list(img, area_row_lines, area_col_lines, table_location_list, bbox_list): area_table_bbox_list = [] area_table_cell_list = [] for i in range(len(table_location_list)): row_lines = area_row_lines[i] col_lines = area_col_lines[i] # 求线交点 cross_points = get_points_by_line(img, row_lines, col_lines) # for p in cross_points: # cv2.circle(img, p, 2, (0, 0, 255), 2) # cv2.imshow('cross_points', img) # 交点分行 cross_points.sort(key=lambda x: (x[1], x[0])) row_point_list = [] if not cross_points: area_table_bbox_list.append([]) area_table_cell_list.append([]) continue current_row = [cross_points[0]] for p in cross_points[1:]: if current_row[0][1] == p[1]: current_row.append(p) else: row_point_list.append(current_row) current_row = [p] if current_row: row_point_list.append(current_row) # bbox以表格格式排列 used_bbox_list = [] row_list = [] row_cell_list = [] for j in range(1, len(row_point_list)): last_row = row_point_list[j-1] row = row_point_list[j] col_list = [] col_cell_list = [] for k in range(1, len(row)): last_p = last_row[k-1] p = row[k] cell = [] for bbox in bbox_list: if bbox in used_bbox_list: continue bbox_h_center = (bbox[0][1]+bbox[2][1]) / 2 bbox_w_center = (bbox[0][0]+bbox[2][0]) / 2 if last_p[1] <= bbox_h_center <= p[1] and last_p[0] <= bbox_w_center <= p[0]: cell.append(bbox) used_bbox_list.append(bbox) col_list.append(cell) col_cell_list.append([last_p, p]) row_list.append(col_list) row_cell_list.append(col_cell_list) area_table_bbox_list.append(row_list) area_table_cell_list.append(row_cell_list) return area_table_bbox_list[0], area_table_cell_list[0]