FORMAT_CONVERSION_MAXCOMPUTE/botr/extract_table.py

import re
import time
import traceback
import cv2
from pdfminer.layout import LTLine
# from botr.nsp.predict import nsp_predict
from botr.rules.get_table_by_rules import get_table_by_rule
from botr.utils import line_iou, get_table_iou
from format_convert.convert_need_interface import from_yolo_interface
from format_convert.utils import log, np2bytes


def b_table_process(list_line, list_text_boxes, list_cell, table_location):
    def merge_textbox(textbox_list, in_objs):
        delete_obj = []
        threshold = 5
        textbox_list.sort(key=lambda x:x.bbox[0])
        for k in range(len(textbox_list)):
            tb1 = textbox_list[k]
            if tb1 not in in_objs and tb1 not in delete_obj:
                for m in range(k+1, len(textbox_list)):
                    tb2 = textbox_list[m]
                    if tb2 in in_objs:
                        continue
                    if abs(tb1.bbox[1]-tb2.bbox[1]) <= threshold \
                            and abs(tb1.bbox[3]-tb2.bbox[3]) <= threshold:
                        if tb1.bbox[0] <= tb2.bbox[0]:
                            tb1.text = tb1.text + tb2.text
                        else:
                            tb1.text = tb2.text + tb1.text
                        tb1.bbox[0] = min(tb1.bbox[0], tb2.bbox[0])
                        tb1.bbox[2] = max(tb1.bbox[2], tb2.bbox[2])
                        delete_obj.append(tb2)
        for _obj in delete_obj:
            if _obj in textbox_list:
                textbox_list.remove(_obj)
        return textbox_list
    try:
        if list_line:
            from format_convert.convert_tree import TableLine
            list_lines = []
            for line in list_line:
                list_lines.append(LTLine(1, (line[0], line[1]), (line[2], line[3])))

            # 先拿出在表格区域里的TextBox
            area_list_text_boxes = []
            threshold = 7
            for t_b in list_text_boxes:
                bbox = t_b.bbox
                if table_location[1] - threshold <= bbox[1] <= bbox[3] <= table_location[3] + threshold:
                    area_list_text_boxes.append(t_b)

            # 对TextBox进行分行，否则同样一行有些框偏上有些偏下，影响文本顺序
            area_list_text_boxes.sort(key=lambda x: (x.bbox[1], x.bbox[0], x.bbox[3], x.bbox[2]))
            current_y = area_list_text_boxes[0].bbox[1]
            threshold = 2.
            for t_b in area_list_text_boxes:
                bbox = t_b.bbox
                if current_y - threshold <= bbox[1] <= current_y + threshold:
                    t_b.bbox[1] = current_y
                else:
                    current_y = bbox[1]
            area_list_text_boxes.sort(key=lambda x: (x.bbox[1], x.bbox[0], x.bbox[3], x.bbox[2]))

            # list_cell 转化为 LineTable形式
            tables = []
            obj_in_table = []
            table_dict = {'bbox': table_location}
            row_list = []
            for row in list_cell:
                col_list = []
                for col in row:
                    col_dict = {'bbox': (col[0][0], col[0][1], col[1][0], col[1][1]),
                                'rowspan': 1, 'columnspan': 1, 'text': ''}
                    for t_b in area_list_text_boxes:
                        if t_b in obj_in_table:
                            continue
                        text = re.sub('\s', '', t_b.text)
                        bbox = t_b.bbox
                        iou = get_table_iou(col[0][0], col[0][1], col[1][0], col[1][1],
                                            bbox[0], bbox[1], bbox[2], bbox[3])
                        if iou >= 0.3:
                            col_dict['text'] += text
                            obj_in_table.append(t_b)
                    col_list.append(col_dict)
                row_list.append(col_list)
            table_dict['table'] = row_list
            tables.append(table_dict)

            # print('b_table_process tables', tables)

            # 合并同一行textbox
            # list_text_boxes = merge_textbox(list_text_boxes, obj_in_table)
            return list_text_boxes, tables, obj_in_table
        else:
            return list_text_boxes, [], set()
    except:
        traceback.print_exc()
        return [-8], [-8], [-8]


def get_text_box_obj(_text_list, _bbox_list):
    from format_convert.convert_tree import TextBox
    _text_box_list = []
    for i in range(len(_bbox_list)):
        bbox = _bbox_list[i]
        b_text = _text_list[i]
        _text_box_list.append(TextBox([bbox[0][0], bbox[0][1],
                                       bbox[2][0], bbox[2][1]], b_text))
    return _text_box_list


def get_table(img, table_list, text_list, bbox_list, text_box_list, show=0):
    log('start')
    # 检测无边框表格
    start_time_all = time.time()
    start_time = time.time()
    img_bytes = np2bytes(img)
    b_table_list = from_yolo_interface(img_bytes)
    log('yolo detect cost: ' + str(time.time()-start_time))
    b_table_list = b_table_list[0]
    if not b_table_list:
        log('detect not b_table_list')
        return [], [], []

    # if show:
    #     for b_table in b_table_list:
    #         # for line in b_table:
    #         cv2.rectangle(img, (int(b_table[0]), int(b_table[1])), (int(b_table[2]), int(b_table[3])),
    #                  (0, 0, 255), 2)
    #     cv2.namedWindow('b_table', cv2.WINDOW_NORMAL)
    #     cv2.imshow('b_table', img)
    #     cv2.waitKey(0)

    if show:
        print('b_table_list', b_table_list)
        print('table_list', table_list)

    # 排除otr结果
    b_table_location_list = []
    for b_table in b_table_list:
        # print('b_table', b_table)
        min_x, min_y = 1000000, 1000000
        max_x, max_y = 0, 0
        # for line in b_table:
        if b_table[1] < min_y:
            min_y = b_table[1]
        if b_table[3] > max_y:
            max_y = b_table[3]
        if b_table[0] < min_x:
            min_x = b_table[0]
        if b_table[2] > max_x:
            max_x = b_table[2]
        b_loc = [min_x, min_y, max_x, max_y, b_table[4]]
        inter_flag = False
        for table in table_list:
            loc = table.get('bbox')
            rows = table.get('table')
            iou = line_iou([[0, loc[1]], [0, loc[3]]], [[0, b_loc[1]], [0, b_loc[3]]], axis=1)
            if iou > 0.3:
                # if len(rows) <= 1:
                #     if loc[1] < b_loc[1] < loc[3] < b_loc[3]:
                #         b_loc[1] = loc[3]
                #     if b_loc[1] < loc[1] < b_loc[3] < loc[3]:
                #         b_loc[3] = loc[1]
                #     continue
                inter_flag = True
                # cv2.rectangle(img, [int(loc[0]), int(loc[1])], [int(loc[2]), int(loc[3])], (0, 0, 255))
                # cv2.rectangle(img, [int(b_loc[0]), int(b_loc[1])], [int(b_loc[2]), int(b_loc[3])], (0, 0, 255))
                # cv2.imshow('inter', img)
                # cv2.waitKey(0)
                break
        if not inter_flag:
            b_table_location_list.append(b_loc)
    if not b_table_location_list:
        log('except otr, not b_table_location_list')
        return [], [], []

    if show:
        print('len(b_table_location_list)', len(b_table_location_list))

    # 排除有重合的，取概率大的
    if len(b_table_location_list) > 1:
        temp_list = []
        used_b_loc = []
        for i in range(len(b_table_location_list)):
            b_loc1 = b_table_location_list[i]
            if b_loc1 in used_b_loc:
                continue
            inter_flag = False
            for j in range(i+1, len(b_table_location_list)):
                b_loc2 = b_table_location_list[j]
                iou = line_iou([[0, b_loc1[1]], [0, b_loc1[3]]], [[0, b_loc2[1]], [0, b_loc2[3]]], axis=1)
                if show:
                    print('iou2', iou)
                if iou > 0.3:
                    inter_flag = True
                    break
            if inter_flag:
                used_b_loc.append(b_loc2)
                if b_loc1[4] >= b_loc2[4]:
                    temp_list.append(b_loc1[:4])
                else:
                    temp_list.append(b_loc2[:4])
            else:
                temp_list.append(b_loc1[:4])
        b_table_location_list = temp_list

    if show:
        for b_loc in b_table_location_list:
            cv2.rectangle(img, (int(b_loc[0]), int(b_loc[1])), (int(b_loc[2]), int(b_loc[3])),
                          (0, 0, 255), 2)
        cv2.namedWindow('b_table_no_otr', cv2.WINDOW_NORMAL)
        cv2.imshow('b_table_no_otr', img)
        cv2.waitKey(0)

    table_list = []
    obj_in_table_list = []
    # print('len(b_table_location_list)', len(b_table_location_list))
    for b_loc in b_table_location_list:
        area_text_list = []
        area_bbox_list = []
        threshold = 5
        for i, bbox in enumerate(bbox_list):
            if b_loc[1] - threshold <= bbox[0][1] <= bbox[2][1] <= b_loc[3] + threshold:
                area_bbox_list.append(bbox)
                area_text_list.append(text_list[i])

        # 根据ocr bbox，规则生成表格线
        start_time = time.time()
        line_list, cell_list, table_location, bbox_text_dict = get_table_by_rule(img, area_text_list, area_bbox_list, b_loc, show=show)
        if not table_location:
            log('get_table_by_rule not table_location')
            continue
        # 获取最新的text_list, bbox_list
        area_text_list, area_bbox_list = [], []
        for key in bbox_text_dict.keys():
            area_bbox_list.append(eval(key))
            area_text_list.append(bbox_text_dict.get(key))
        b_text_box_list = get_text_box_obj(area_text_list, area_bbox_list)
        log('get_table_by_rule cost: ' + str(time.time()-start_time))

        # 根据表格线生成单元格
        start_time = time.time()
        b_text_box_list, _table_list, _obj_in_table_list = b_table_process(line_list, b_text_box_list, cell_list, table_location)
        table_list += _table_list
        obj_in_table_list += _obj_in_table_list
        log('b_table_process cost: ' + str(time.time()-start_time))

        # if not table_list:
        #     log('table_process not table_list')
        #     return [], [], []

        if not _table_list:
            log('table_process not table_list')
            continue

        # 单元格合并，nsp模型
        # 使用hanlp分词，判断上下句是否该合并 顺便拉数据统计
        # 1. 上下句ab，ab相连得到c
        # 2.1 c分词，若ab相连处合为一个词语，则ab相连
        # 2.2 ab相连处不为一个词语，a, b分别分词
        # 2.2.1 若b的第一个词，从其中分第一个字给a，然后
        # near_col_list = []
        # table = _table_list[0].get('table')
        # col_cnt = len(table[0])
        # for c_cnt in range(col_cnt):
        #     for i in range(len(table)-1):
        #         t = table[i][c_cnt].get('text')
        #         next_t = table[i+1][c_cnt].get('text')
        #         if t and next_t:
        #             near_col_list.append([t, next_t])
        #         elif t and next_t == '':
        #             if i+2 <= len(table)-1:
        #                 next_2_t = table[i+2][c_cnt].get('text')
        #                 near_col_list.append([t, next_2_t])
        #
        # is_next_list = nsp_predict(near_col_list, has_label=False)
        #
        # next_index = 0
        # for c_cnt in range(col_cnt):
        #     # 先把一列里的需合并的打上标签
        #     for i in range(len(table)-1):
        #         t = table[i][c_cnt].get('text')
        #         next_t = table[i+1][c_cnt].get('text')
        #         if t and next_t:
        #             table[i+1][c_cnt]['is_next'] = is_next_list[next_index]
        #             next_index += 1
        #         elif t and next_t == '':
        #             if i+2 <= len(table)-1:
        #                 table[i+1][c_cnt]['is_next'] = is_next_list[next_index]
        #                 table[i+2][c_cnt]['is_next'] = is_next_list[next_index]
        #                 next_index += 1
        #
        #     first_col = None
        #     for i in range(len(table)):
        #         if table[i][c_cnt].get('is_next'):
        #             if first_col is None:
        #                 first_col = table[i-1][c_cnt]
        #             first_col['text'] += table[i][c_cnt].get('text')
        #             first_col['rowspan'] += 1
        #         else:
        #             first_col = None
        #
        # # 删除标签为True的
        # new_table = []
        # for row in table:
        #     new_row = []
        #     for col in row:
        #         if col.get('is_next'):
        #             continue
        #         new_row.append(col)
        #     new_table.append(new_row)
        #
        # _table_list[0]['table'] = new_table

    log('get_table finish ' + str(time.time() - start_time_all))
    return text_box_list, table_list, obj_in_table_list