FORMAT_CONVERSION_MAXCOMPUTE/botr/rules/table_utils.py

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]