BIDI_ML_INFO_EXTRACTION/BiddingKG/dl/table_head/pre_process.py

import os
import random
import sys
import psycopg2
import numpy as np
sys.path.append(os.path.dirname(__file__) + "/../")
from common.Utils import embedding_word, embedding_word_forward


def get_sentence_index_list(sentence, dict_path='utils/ppocr_keys_v1.txt'):
    with open(dict_path, 'r') as f:
        character_list = f.readlines()
    for i in range(len(character_list)):
        character_list[i] = character_list[i][:-1]

    index_list = []
    for character in sentence:
        if character == '':
            index_list.append(0)
        elif character in character_list:
            _index = character_list.index(character) + 1
            index_list.append(_index)
        else:
            index_list.append(0)
    return index_list


def postgresql_util(sql, limit):
    conn = psycopg2.connect(dbname="table_head_label", user="postgres", password="postgres",
                            host="192.168.2.103")
    cursor = conn.cursor()
    cursor.execute(sql)
    print(sql)

    rows = cursor.fetchmany(1000)
    cnt = 0
    all_rows = []
    while rows:
        if cnt >= limit:
            break
        all_rows += rows
        cnt += len(rows)
        rows = cursor.fetchmany(1000)
    return all_rows


def get_data_from_sql(dim=10, whole_table=False, padding=True):
    sql = """
    select table_text, pre_label, post_label, id
    from label_table_head_info
    where status = 0 and (update_user='test9' or update_user='test1' or update_user='test7' or update_user='test26')
    ;
    """
    # sql = """
    # select table_text, pre_label, post_label, id
    # from label_table_head_info
    # where status = 1 and update_time >= '2022-01-17' and update_time <= '2022-01-22'
    # ;
    # """

    result_list = postgresql_util(sql, limit=1000000)

    # 需排除的id
    with open(r"C:\Users\Administrator\Desktop\table_not_eval.txt", "r") as f:
        delete_id_list = eval(f.read())
    with open(r"C:\Users\Administrator\Desktop\table_delete.txt", "r") as f:
        delete_id_list += eval(f.read())

    all_data_list = []
    all_data_label_list = []
    i = 0
    # 一行就是一篇表格
    for table in result_list:
        i += 1
        if i % 100 == 0:
            print("Loop", i)

        pre_label = eval(table[1])
        post_label = eval(table[2])
        _id = table[3]

        if _id in delete_id_list:
            print("pass", _id)
            continue

        # table_text需要特殊处理
        try:
            table_text = table[0]
            if table_text[0] == '"':
                table_text = eval(table_text)
            else:
                table_text = table_text
            table_text = table_text.replace('\\', '/')
            table_text = eval(table_text)
        except:
            print("无法识别table_text", _id)
            continue

        if whole_table:
            if len(post_label) >= 2:
                data_list, data_label_list = table_pre_process_2(table_text, post_label,
                                                                 _id, padding=padding)
            elif len(pre_label) >= 2:
                data_list, data_label_list = table_pre_process_2(table_text, pre_label,
                                                                 _id, padding=padding)
            else:
                data_list, data_label_list = [], []
        else:
            # 只有一行的也不要
            if len(post_label) >= 2:
                data_list, data_label_list = table_pre_process(table_text, post_label, _id)
            elif len(pre_label) >= 2:
                data_list, data_label_list = table_pre_process(table_text, pre_label, _id)
            else:
                data_list, data_label_list = [], []

        all_data_list += data_list
        all_data_label_list += data_label_list

    # 按维度大小排序
    if whole_table:
        _list = []
        for data, label in zip(all_data_list, all_data_label_list):
            _list.append([data, label])
        _list.sort(key=lambda x: (len(x[0]), len(x[0][0])))
        all_data_list[:], all_data_label_list[:] = zip(*_list)

    print("len(all_data_list)", len(all_data_list))
    return all_data_list, all_data_label_list


def table_pre_process(text_list, label_list, _id, is_train=True):
    """
    表格处理，每个单元格生成2条数据，横竖各1条

    :param text_list:
    :param label_list:
    :param _id:
    :param is_train:
    :return:
    """
    if is_train:
        if len(text_list) != len(label_list):
            print("文字单元格与标注单元格数量不匹配！", _id)
            print("len(text_list)", len(text_list), "len(label_list)", len(label_list))
            return [], []

    data_list = []
    data_label_list = []
    for i in range(len(text_list)):
        row = text_list[i]
        if is_train:
            row_label = label_list[i]

        if i > 0:
            last_row = text_list[i-1]
            if is_train:
                last_row_label = label_list[i-1]
        else:
            last_row = []
            if is_train:
                last_row_label = []

        if i < len(text_list) - 1:
            next_row = text_list[i+1]
            if is_train:
                next_row_label = label_list[i+1]
        else:
            next_row = []
            if is_train:
                next_row_label = []

        for j in range(len(row)):
            col = row[j]
            if is_train:
                col_label = row_label[j]

            # 超出表格置为None, 0
            if j > 0:
                last_col = row[j-1]
                if is_train:
                    last_col_label = row_label[j-1]
            else:
                last_col = col
                if is_train:
                    last_col_label = col_label

            if j < len(row) - 1:
                next_col = row[j+1]
                if is_train:
                    next_col_label = row_label[j+1]
            else:
                next_col = col
                if is_train:
                    next_col_label = col_label

            if last_row:
                last_row_col = last_row[j]
                if is_train:
                    last_row_col_label = last_row_label[j]
            else:
                last_row_col = col
                if is_train:
                    last_row_col_label = col_label

            if next_row:
                next_row_col = next_row[j]
                if is_train:
                    next_row_col_label = next_row_label[j]
            else:
                next_row_col = col
                if is_train:
                    next_row_col_label = col_label

            # data_list.append([last_col, col, next_col])
            # if is_train:
            #     data_label_list.append([int(last_col_label), int(col_label),
            #                             int(next_col_label)])
            #
            # data_list.append([last_row_col, col, next_row_col])
            # if is_train:
            #     data_label_list.append([int(last_row_col_label), int(col_label),
            #                             int(next_row_col_label)])

            if is_train:
                dup_list = [str(x) for x in data_list]
                data = [last_col, col, next_col, last_row_col, col, next_row_col]
                if str(data) not in dup_list:
                    data_list.append([last_col, col, next_col, last_row_col, col, next_row_col])
                    data_label_list.append(int(col_label))
            else:
                data_list.append([last_col, col, next_col, last_row_col, col, next_row_col])
    if is_train:
        return data_list, data_label_list
    else:
        return data_list


def table_pre_process_2(text_list, label_list, _id, is_train=True, padding=True):
    """
    表格处理，整个表格为一个数组，且填充长宽维度

    :param text_list:
    :param label_list:
    :param _id:
    :param is_train:
    :return:
    """
    # 判断表格长宽是否合理
    row_len = len(text_list)
    best_row_len = get_best_padding_size(row_len, min_len=8)
    col_len = len(text_list[0])
    best_col_len = get_best_padding_size(col_len, min_len=8)
    if best_row_len is None:
        if is_train:
            return [], []
        else:
            return []
    if best_col_len is None:
        if is_train:
            return [], []
        else:
            return []

    if is_train:
        if len(text_list) != len(label_list):
            print("文字单元格与标注单元格数量不匹配！", _id)
            print("len(text_list)", len(text_list), "len(label_list)", len(label_list))
            return [], []

        if padding:
            for i in range(row_len):
                col_len = len(text_list[i])
                text_list[i] += [None]*(best_col_len-col_len)
                if is_train:
                    label_list[i] += ["0"]*(best_col_len-col_len)
            text_list += [[None]*best_col_len]*(best_row_len-row_len)
            if is_train:
                label_list += [["0"]*best_col_len]*(best_row_len-row_len)

    if is_train:
        for i in range(len(label_list)):
            for j in range(len(label_list[i])):
                label_list[i][j] = int(label_list[i][j])
        return [text_list], [label_list]
    else:
        return [text_list]


def get_best_padding_size(axis_len, min_len=3, max_len=300):
    # sizes = [8, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120,
    #          128, 136, 144, 152, 160, 168, 176, 184, 192, 200, 208, 216, 224,
    #          232, 240, 248, 256, 264, 272, 280, 288, 296]
    # sizes = [3, 6, 9, 12, 15, 18, 21, 24, 27, 30, 33, 36, 39, 42, 45, 48, 51, 54, 57,
    #          60, 63, 66, 69, 72, 75, 78, 81, 84, 87, 90, 93, 96, 99, 102, 105, 108, 111,
    #          114, 117, 120, 123, 126, 129, 132, 135, 138, 141, 144, 147, 150, 153, 156,
    #          159, 162, 165, 168, 171, 174, 177, 180, 183, 186, 189, 192, 195, 198, 201,
    #          204, 207, 210, 213, 216, 219, 222, 225, 228, 231, 234, 237, 240, 243, 246,
    #          249, 252, 255, 258, 261, 264, 267, 270, 273, 276, 279, 282, 285, 288, 291,
    #          294, 297]
    sizes = []
    for i in range(1, max_len):
        if i * min_len <= max_len:
            sizes.append(i * min_len)
    if axis_len > sizes[-1]:
        return axis_len
    best_len = sizes[-1]
    for height in sizes:
        if axis_len <= height:
            best_len = height
            break
    # print("get_best_padding_size", axis_len, best_len)
    return best_len


def get_data_from_file(file_type, model_id=1):
    if file_type == 'np':
        data_path = 'train_data/data_3.npy'
        data_label_path = 'train_data/data_label_3.npy'

        array1 = np.load(data_path)
        array2 = np.load(data_label_path)
        return array1, array2
    elif file_type == 'txt':
        if model_id == 1:
            data_path = 'train_data/data1.txt'
            data_label_path = 'train_data/data_label1.txt'
        elif model_id == 2:
            data_path = 'train_data/data2.txt'
            data_label_path = 'train_data/data_label2.txt'
        elif model_id == 3:
            data_path = 'train_data/data3.txt'
            data_label_path = 'train_data/data_label3.txt'
        with open(data_path, 'r') as f:
            data_list = f.readlines()
        with open(data_label_path, 'r') as f:
            data_label_list = f.readlines()

        return data_list, data_label_list
    else:
        print("file type error! only np and txt supported")
        raise Exception


def processed_save_to_np():
    array1, array2 = get_data_from_sql()
    np.save('train_data/data_3.npy', array1)
    np.save('train_data/data_label_3.npy', array2)

    # with open('train_data/data.txt', 'w') as f:
    #     for line in list1:
    #         f.write(str(line) + "\n")
    # with open('train_data/data_label.txt', 'w') as f:
    #     for line in list2:
    #         f.write(str(line) + "\n")


def processed_save_to_txt(whole_table=False, padding=True):
    list1, list2 = get_data_from_sql(whole_table=whole_table, padding=padding)

    # 打乱
    # if not whole_table or not padding:
    zip_list = list(zip(list1, list2))
    random.shuffle(zip_list)
    list1[:], list2[:] = zip(*zip_list)

    with open('train_data/data1.txt', 'w') as f:
        for line in list1:
            f.write(str(line) + "\n")
    with open('train_data/data_label1.txt', 'w') as f:
        for line in list2:
            f.write(str(line) + "\n")


def data_balance():
    data_list, data_label_list = get_data_from_file('txt')
    all_cnt = len(data_label_list)
    cnt_0 = 0
    cnt_1 = 0
    for data in data_label_list:
        if eval(data[:-1])[1] == 1:
            cnt_1 += 1
        else:
            cnt_0 += 1
    print("all_cnt", all_cnt)
    print("label has 1", cnt_1)
    print("label all 0", cnt_0)


def test_embedding():
    output_shape = (2, 1, 60)
    data = [[None], [None]]
    result = embedding_word(data, output_shape)
    print(result)


def my_data_loader(data_list, data_label_list, batch_size, is_train=True):
    data_num = len(data_list)

    # 定义Embedding输出
    output_shape = (6, 20, 60)

    # batch循环取数据
    i = 0
    if is_train:
        while True:
            new_data_list = []
            new_data_label_list = []
            for j in range(batch_size):
                if i >= data_num:
                    i = 0

                # 中文字符映射为Embedding
                data = eval(data_list[i][:-1])
                data_label = eval(data_label_list[i][:-1])
                data = embedding_word(data, output_shape)
                if data.shape == output_shape:
                    new_data_list.append(data)
                    new_data_label_list.append(data_label)
                i += 1

            new_data_list = np.array(new_data_list)
            new_data_label_list = np.array(new_data_label_list)
            X = new_data_list
            Y = new_data_label_list
            # (table_num, 3 sentences, dim characters, embedding) -> (3, table_num, dim, embedding)
            X = np.transpose(X, (1, 0, 2, 3))

            if (X[0] == X[1]).all():
                X[0] = np.zeros_like(X[1], dtype='float32')
            if (X[2] == X[1]).all():
                X[2] = np.zeros_like(X[1], dtype='float32')
            if (X[3] == X[1]).all():
                X[3] = np.zeros_like(X[1], dtype='float32')
            if (X[5] == X[1]).all():
                X[5] = np.zeros_like(X[1], dtype='float32')

            yield {'input_1': X[0], 'input_2': X[1], 'input_3': X[2],
                   'input_4': X[3], 'input_5': X[4], 'input_6': X[5]}, \
                  {'output': Y}

    else:
        while True:
            new_data_list = []
            for j in range(batch_size):
                if i >= data_num:
                    i = 0

                # 中文字符映射为Embedding
                data = data_list[i]
                data = embedding_word(data, output_shape)
                if data.shape == output_shape:
                    new_data_list.append(data)
                i += 1

            new_data_list = np.array(new_data_list)
            X = new_data_list
            X = np.transpose(X, (1, 0, 2, 3))
            yield {'input_1': X[0], 'input_2': X[1], 'input_3': X[2],
                   'input_4': X[3], 'input_5': X[4], 'input_6': X[5], }


def my_data_loader_2(table_list, table_label_list, batch_size, is_train=True):
    pad_len = 0

    table_num = len(table_list)
    if is_train and batch_size == 1:
        table_list, table_label_list = get_random(table_list, table_label_list)

    # Embedding shape
    output_shape = (20, 60)

    # batch循环取数据
    i = 0
    last_shape = None
    while True:
        new_table_list = []
        new_table_label_list = []
        for j in range(batch_size):
            if i >= table_num:
                i = 0
                if is_train:
                    table_list, table_label_list = get_random(table_list, table_label_list,
                                                              seed=random.randint(1, 40))

            if type(table_list[i]) != list:
                table = eval(table_list[i][:-1])
            else:
                table = table_list[i]

            if batch_size > 1:
                if last_shape is None:
                    last_shape = (len(table), len(table[0]))
                    continue
                if (len(table), len(table[0])) != last_shape:
                    last_shape = (len(table), len(table[0]))
                    break

            if is_train:
                table_label = eval(table_label_list[i][:-1])

            # 中文字符映射为Embedding
            for k in range(len(table)):
                table[k] = embedding_word_forward(table[k], (len(table[k]),
                                                     output_shape[0],
                                                     output_shape[1]))
            new_table_list.append(table)
            if is_train:
                new_table_label_list.append(table_label)
            i += 1
        new_table_list = np.array(new_table_list)
        X = new_table_list
        if X.shape[-2:] != output_shape:
            # print("Dimension not match!", X.shape)
            # print("\n")
            continue

        # 获取Padding大小
        pad_height = get_best_padding_size(X.shape[1], pad_len)
        pad_width = get_best_padding_size(X.shape[2], pad_len)
        input_2 = np.zeros([1, X.shape[1], X.shape[2], pad_height, pad_width])

        if is_train:
            new_table_label_list = np.array(new_table_label_list)
            Y = new_table_label_list
            # Y = Y.astype(np.float32)
            # yield {"input_1": X, "input_2": input_2}, \
            #       {"output_1": Y, "output_2": Y}
            yield {"input_1": X, "input_2": input_2}, \
                  {"output": Y}
        else:
            yield {"input_1": X, "input_2": input_2}


def check_train_data():
    data_list, label_list = get_data_from_file('txt', model_id=2)
    for data in data_list:
        data = eval(data)
        if len(data) % 8 != 0:
            print(len(data))
            print(len(data[0]))
        for row in data:
            if len(row) % 8 != 0:
                print(len(data))
                print(len(row))


def get_random(text_list, label_list, seed=42):
    random.seed(seed)
    zip_list = list(zip(text_list, label_list))
    random.shuffle(zip_list)
    text_list[:], label_list[:] = zip(*zip_list)
    return text_list, label_list


if __name__ == '__main__':
    processed_save_to_txt(whole_table=False, padding=False)
    # data_balance()

    # test_embedding()
    # check_train_data()

    # _list = []
    # for i in range(1, 100):
    #     _list.append(i*3)
    # print(_list)

    # print(get_best_padding_size(9, 5))