PaddleOCR/ppocr/data/text2Image.py

import random
import re
import numpy as np
import cv2
# import psycopg2
from PIL import Image, ImageFont, ImageDraw
import os
from PIL import Image, ImageFont, ImageDraw
Image.MAX_IMAGE_PIXELS = None
import pandas as pd


# project_path = "D:\\Project\\PaddleOCR-release-2.0\\"
from bs4 import BeautifulSoup

current_path = os.path.dirname(__file__)
project_path = current_path + "/../../"
image_output_path = project_path + "train_data/bidi_data/mix_data5/"
train_data_path = image_output_path + "rec_gt_train.txt"
test_data_path = image_output_path + "rec_gt_test.txt"


def create_image(data_dir, file_name, text):
    # list1 = re.findall('[a-zA-Z\d]', text)
    # list2 = re.findall('[\u4e00-\u9fa5。，！？￥《》【】’“：；·、（）]', text)
    # list3 = re.findall('[,.!?&@*+=~%()#<>-''|/:{}$;]', text)
    # english_len = len(list1)
    # chinese_len = len(list2)
    # character_len = len(list3)
    #
    # if english_len + chinese_len + character_len == 0:
    #     character_len = len(text)

    # 根据各字体大小生成图片
    # font 10 : a1-6 字-10 image-len*, 16
    # font 20 : a1-12 字-20 image-len*, 32
    font_list = [7, 8, 9, 10, 11, 12, 15, 20, 25, 30, 35, 40]
    # 随机选字体大小
    font_index = random.randint(0, len(font_list)-1)
    font = font_list[font_index]

    # 根据字体大小计算各字符长度
    # chinese_charc_len = font * 1
    # english_charc_len = int(font * 0.7)
    # number_charc_len = int(font * 0.3)
    # image_width = int(font * 1.6)
    # text_len = english_len * english_charc_len + chinese_len * chinese_charc_len \
    #            + character_len * number_charc_len

    # 获取字体及其实际像素大小
    font_list = []
    for root, dirs, files in os.walk(project_path + "tools/fonts"):
        for file in files:
            font_list.append(os.path.join(root, file))
    font_type = random.choice(font_list)
    font = ImageFont.truetype(font_type, font)
    font_width = font.getsize(text)[0]
    font_height = font.getsize(text)[1]

    # 增加上下左右边距
    margin_h = 0
    margin_w = 0

    if font_height > 9:
        if random.choice([0, 1]):
            margin_h = random.randint(3, 6)
            margin_w = random.randint(3, 6)

    im = Image.new("RGB", (font_width+margin_w*2, font_height+margin_h*2), (255, 255, 255))
    dr = ImageDraw.Draw(im)

    dr.text((0+margin_w, 0+margin_h), text, font=font, fill="#000000")

    # if random.choice([0, 1, 1]):
    #     # PIL -> CV2
    #     img = cv2.cvtColor(np.asarray(im), cv2.COLOR_RGB2BGR)
    #     # 文字加下划线
    #     img = cv2.line(img, (0+margin_w, font_height+margin_h),
    #                    (font_width+margin_w, font_height+margin_h),
    #                    (0, 0, 0), 1)
    #     # CV2 -> PIL
    #     im = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))

    #
    # # 随机缩放
    # resize_y = random.randint(1, 2)
    # resize_x = random.randint(1, 2)
    # img = cv2.resize(img, (img.shape[1]*resize_y, img.shape[0]*resize_x))
    #
    # # 模糊
    # # 高斯模糊
    # sigmaX = random.randint(1, 3)
    # sigmaY = random.randint(1, 3)
    # img = cv2.GaussianBlur(img, (5, 5), sigmaX, sigmaY)
    #

    # # resize_y = random.uniform(1, 3)
    # # resize_x = random.uniform(1, 3)
    # # img = im.resize((int(im.size[0]*resize_y), int(im.size[1]*resize_x)), Image.ANTIALIAS)
    #
    # # 保存
    # # cv2.imwrite(data_dir + file_name, img)

    # 图像增强
    im = my_image_aug(im)
    # print(im.size)
    # im.show("img")
    im.save(data_dir + file_name)


def my_image_aug(image_pil):
    # 图像增强
    # PIL -> CV2
    img = cv2.cvtColor(np.asarray(image_pil), cv2.COLOR_RGB2BGR)

    # 分辨率低的跳过
    if img.shape[0] > 20:
        # 随机缩放
        if random.choice([0, 1, 1]):
            resize_y = random.randint(1, 3)
            resize_x = random.randint(1, 3)
            img = cv2.resize(img, (img.shape[1]*resize_y, img.shape[0]*resize_x))

        if_dilate = random.choice([0, 1, 1, 2, 2])
        # if_dilate = 1
        if if_dilate == 2:
            gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (2, 2))
            iters = 1
            # if img.shape[0] > 40:
            #     iters = 2
            dilate = cv2.dilate(gray, kernel, iterations=iters)
            img = cv2.cvtColor(dilate, cv2.COLOR_GRAY2BGR)

        elif if_dilate == 1:
            # 高斯模糊
            sigmaX = random.randint(1, 2)
            sigmaY = random.randint(1, 2)
            img = cv2.GaussianBlur(img, (9, 9), sigmaX, sigmaY)

    # CV2 -> PIL
    im = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))

    return im


def create_orgs_image(df):
    df = df[:1000]
    label_file_train = project_path + "train_data\\bidi_data\\orgs_data\\rec_gt_train.txt"
    label_file_test = project_path + "train_data\\bidi_data\\orgs_data\\rec_gt_test.txt"
    image_output_path = project_path + "train_data\\bidi_data\\orgs_data\\"
    f1 = open(label_file_train, "w")
    f2 = open(label_file_test, "w")
    print(df.shape)
    for index, row in df.iterrows():
        text = row["name"]
        # text = "晋江滨江国家体育训练基地有限公司"
        im = Image.new("RGB", (len(text)*10, 16), (255, 255, 255))
        dr = ImageDraw.Draw(im)
        font = ImageFont.truetype(os.path.join(os.getcwd(), "fonts", "msyh.ttc"), 10)
        dr.text((0, 0), text, font=font, fill="#000000")
        # im.show()
        if index / df.shape[0] <= 0.8:
            mode = "train"
            f = f1
        else:
            mode = "test"
            f = f2
        im.save(image_output_path + mode + "\\" + "text_" + str(index) + ".jpg")
        f.write(mode + "/text_" + str(index) + ".jpg" + "\t" + text + "\n")

    f1.close()
    f2.close()


def create_longSentence_image(df):
    # df = df[:3000]
    label_file_train = project_path + "train_data\\bidi_data\\longSentence_data\\rec_gt_train.txt"
    label_file_test = project_path + "train_data\\bidi_data\\longSentence_data\\rec_gt_test.txt"
    image_output_path = project_path + "train_data\\bidi_data\\longSentence_data\\"
    f1 = open(label_file_train, "w")
    f2 = open(label_file_test, "w")
    print(df.shape)

    for index, row in df.iterrows():
        text = row["text"]
        # text = "晋江滨江国家体育训练基地有限公司"
        im = Image.new("RGB", (len(text)*10, 16), (255, 255, 255))
        dr = ImageDraw.Draw(im)
        font = ImageFont.truetype(os.path.join(os.getcwd(), "fonts", "msyh.ttc"), 10)
        dr.text((0, 0), text, font=font, fill="#000000")
        # im.show()
        if index <= int((df.shape[0]-1)*0.8):
            mode = "train"
            f = f1
        else:
            mode = "test"
            f = f2
        im.save(image_output_path + mode + "\\" + "text_" + str(index) + ".jpg")
        f.write(mode + "/text_" + str(index) + ".jpg" + "\t" + text + "\n")

    f1.close()
    f2.close()


# def readPostgreSQL():
#     conn_string = "host=192.168.2.101 port=5432 dbname=iepy " \
#                   "user=iepy_read password=iepy_read"
#     conn = psycopg2.connect(conn_string)
#
#     # 执行SQL语句
#     sql = "select text from corpus_iedocument " \
#           "where jump_signal=0"
#     df = pd.read_sql(sql, conn)
#     return df


# 生成多个场景混合数据


def create_mix_txt():
    # 最长字符串长度
    max_length = 100

    list1 = create_price(60000, True)
    list1 += create_price(40000)
    print("finish get list1", len(list1))

    list2 = get_long_sentence_from_file(100000)
    print("finish get list2", len(list2))

    with open("appendix_text.txt", "r") as f:
        list3 = f.readlines()
    list3 = list3[2000000:2300000]
    print("finish get list3", len(list3))

    list4 = create_org_list()
    list4 = list4[:100000]
    print("finish get list4", len(list4))

    list5 = create_no(800000)
    print("finish get list5", len(list5))

    list6 = get_data_from_paddle()
    list6 = list6[:200000]
    print("finish get list6", len(list6))

    train_data = list1[0:int(len(list1)*0.95)] + list2[0:int(len(list2)*0.95)] + \
                 list3[0:int(len(list3)*0.95)] + list4[0:int(len(list4)*0.95)] + \
                 list5[0:int(len(list5)*0.95)]
    test_data = list1[int(len(list1)*0.95):] + list2[int(len(list2)*0.95):] + \
                list3[int(len(list3)*0.95):] + list4[int(len(list4)*0.95):] + \
                list5[int(len(list5)*0.95):]

    print("len(train_data)", len(train_data))
    print("len(test_data)", len(test_data))

    data_index = 0
    for i in range(len(train_data)):
        prefix = "train/text_" + str(data_index) + ".png" + "\t"
        train_data[i] = prefix + train_data[i]
        data_index += 1
    for i in range(len(test_data)):
        prefix = "test/text_" + str(data_index) + ".png" + "\t"
        test_data[i] = prefix + test_data[i]
        data_index += 1
    train_data += list6[0:int(len(list6)*0.95)]
    test_data += list6[int(len(list6)*0.95):]

    with open(train_data_path, "w") as f:
        f.writelines(train_data)
    with open(test_data_path, "w") as f:
        f.writelines(test_data)
    print("finish write train data and test data!")

    # with open(train_data_path, "w") as f:
    #     for data in train_data:
    #         prefix = "train/text_" + str(data_index) + ".jpg" + "\t"
    #         data = prefix + data
    #         f.write(data)
    #         data_index += 1
    # print("finish write train data")
    # with open(test_data_path, "w") as f:
    #     for data in test_data:
    #         prefix = "test/text_" + str(data_index) + ".jpg" + "\t"
    #         data = prefix + data
    #         f.write(data)
    #         data_index += 1
    # print("finish write test data")
    return


def create_no_price():
    # 最长字符串长度
    max_length = 100
    # list1 = create_text_list(max_length)

    list1 = create_price(50000, True)
    print("finish get list1", len(list1))
    # list2 = create_org_list()

    list2 = get_long_sentence_from_file(10000)
    print("finish get list2", len(list2))
    # list2 = list2[0:100]

    with open("appendix_text.txt", "r") as f:
        list3 = f.readlines()
    list3 = list3[:10000]
    print("finish get list3", len(list3))

    list4 = create_org_list()
    list4 = list4[:10000]
    print("finish get list4", len(list4))

    list5 = create_no(50000)
    print("finish get list5", len(list5))

    train_data = list1[0:int(len(list1)*0.95)] + list2[0:int(len(list2)*0.95)] + \
                 list3[0:int(len(list3)*0.95)] + list4[0:int(len(list4)*0.95)] + \
                 list5[0:int(len(list5)*0.95)]
    test_data = list1[int(len(list1)*0.95):] + list2[int(len(list2)*0.95):] + \
                list3[int(len(list3)*0.95):] + list4[int(len(list4)*0.95):] + \
                list5[int(len(list5)*0.95):]
    print("len(train_data)", len(train_data))
    print("len(test_data)", len(test_data))

    data_index = 0
    with open(train_data_path, "w") as f:
        for data in train_data:
            prefix = "train/text_" + str(data_index) + ".jpg" + "\t"
            data = prefix + data
            f.write(data)
            data_index += 1
    print("finish write train data")
    with open(test_data_path, "w") as f:
        for data in test_data:
            prefix = "test/text_" + str(data_index) + ".jpg" + "\t"
            data = prefix + data
            f.write(data)
            data_index += 1
    print("finish write test data")
    return


# def create_text_list(max_length):
#     # 招投标文章语句
#     df1 = readPostgreSQL()
#     list1 = []
#     for index, row in df1.iterrows():
#         text = row["text"].split("，")
#         # print(len(text))
#
#         # 每篇文章最多取10个句子
#         max_sentence = 15
#         sentence_count = 0
#         while sentence_count < max_sentence:
#             if len(text) <= max_sentence:
#                 if sentence_count < len(text):
#                     sentence = text[sentence_count]
#                 else:
#                     break
#             else:
#                 r1 = random.randint(0, len(text) - 1)
#                 sentence = text[r1]
#             if len(sentence) > max_length:
#                 # 限制字数，随机截取前或后
#                 r2 = random.randint(0, 1)
#                 if r2:
#                     sentence = sentence[:max_length]
#                 else:
#                     sentence = sentence[-max_length:]
#
#             # sentence = re.sub("\n", "", sentence)
#             if sentence != "":
#                 list1.append(sentence+"\n")
#             sentence_count += 1
#     print("len(list1)", len(list1))
#     return list1


def delete_image(data_dir, file_name):
    if os.path.exists(data_dir + file_name):
        os.remove(data_dir + file_name)


def create_org_list():
    # 1kw公司名
    with open("C:\\Users\\Administrator\\Desktop\\LEGAL_ENTERPRISE.txt", "r") as f:
        list2 = f.readlines()
    # list2 = list2[:100]
    # print("len(list2)", len(list2))
    return list2


def create_number_list(number):
    no_list = []
    for i in range(number):
        # 随机选择生成几位小数
        decimal_place = random.choices([0, 1, 2, 3, 4, 5, 6])[0]

        if decimal_place == 0:
            no = random.randint(0, 10000000)
        else:
            no = random.uniform(0, 10000)
            no = round(no, decimal_place)
        no_list.append(str(no)+"\n")
    # print(no_list)
    return no_list


def get_mix_data_from_file(number):
    with open("../../train_data/bidi_data/mix_data/rec_gt_train.txt") as f:
        _list = f.readlines()
    _list = _list[:number]

    new_list = []
    for line in _list:
        s = line.split("\t")[1]
        new_list.append(s)
    # print(new_list)
    return new_list


def get_long_sentence_from_file(number):
    with open("../../train_data/bidi_data/longSentence_data/rec_gt_train.txt") as f:
        list1 = f.readlines()

    with open("../../train_data/bidi_data/longSentence_data/rec_gt_test.txt") as f:
        list2 = f.readlines()

    _list = list1 + list2
    _list = _list[:number]

    new_list = []
    for line in _list:
        s = line.split("\t")[1]
        new_list.append(s)
    # print(new_list)
    return new_list


def get_data_from_appendix():
    df = pd.read_excel("dochtmlcon.xlsx")
    text_list = []
    for index, row in df.iterrows():
        html_text = row["dochtmlcon"]

        # 创建一个BeautifulSoup解析对象
        soup = BeautifulSoup(html_text, "html.parser", from_encoding="utf-8")

        # 获取所有的链接
        appendix_text = soup.find_all('div', class_='richTextFetch')
        # print(str(appendix_text[0])[49:-6])

        appendix_text = str(appendix_text[0])[49:-6]
        ss = appendix_text.split("\n")
        for s in ss:
            text = re.sub(" ", "", s)
            text = re.sub("\t", "", text)
            if s == "":
                continue
            text_list.append(text + "\n")

    with open("appendix_text.txt", "w") as f:
        f.writelines(text_list)
    return


def get_data_from_paddle():
    path = "D:\\DataSet\\"
    with open(path + "char.txt", "r") as f:
        dictionary = f.readlines()
    with open(path + "data_train.txt") as f:
        train_list = f.readlines()
    with open(path + "data_test.txt") as f:
        test_list = f.readlines()

    data_list = train_list + test_list
    # data_list = data_list[-100:]

    text_list = []
    for data in data_list:
        ss = data[:-1].split(" ")
        image_path = "image/" + ss[0]
        text = ""
        for num in ss[1:]:
            char = dictionary[int(num)][:-1]
            text += char
        if text == "":
            print("no text!")
            continue
        text_list.append(image_path + "\t" + text + "\n")

    # with open("paddle_data.txt", "w") as f:
    #     f.writelines(text_list)

    return text_list


def create_number_list2(number):
    no_list = []
    for i in range(number):
        c1 = random.choice([0, 1, 1])
        if c1:
            no = random.randint(0, 10)
        else:
            no = random.randint(10, 100)
        no_list.append(str(no) + "\n")
    # print(no_list)
    return no_list


def create_price(number, add_rmb=False):
    no_list = []
    if not add_rmb:
        num_range = (10, 1000000000)
    else:
        num_range = (10, 100000)
    for i in range(number):
        # 选择小数整数
        c1 = random.choice([0, 1, 1])
        if c1:
            no = random.randint(num_range[0], num_range[1])
            no = str(no)
            # 加3位分割逗号
            # 选择中英文逗号
            c2 = random.choice([',', ',', ','])
            for i in range(len(no)-3, 0, -3):
                no = no[:i] + c2 + no[i:]
        else:
            no = random.uniform(num_range[0], num_range[1])
            n_digits = random.choice([1, 2, 3, 4])
            no = round(no, n_digits)
            no = str(no)
            nos = no.split(".")
            no_1 = nos[0]
            no_2 = nos[1]
            # 加3位分割逗号
            # 选择中英文逗号
            c2 = random.choice([',', ',', ','])
            for i in range(len(no_1)-3, 0, -3):
                no_1 = no_1[:i] + c2 + no_1[i:]
            no = no_1 + "." + no_2

        # 选择是否加￥符号
        if add_rmb:
            c3 = random.choice(["￥", "¥"])
        else:
            c3 = random.choice(['', "￥", "¥"])
        no_list.append(c3 + str(no) + "\n")
    # print(no_list[:100])
    return no_list


def create_no(number):
    no_list = []
    chn_no = ['一', '二', '三', '四', '五', '六', '七', '八', '九', '十',
              '壹', '贰', '叁', '肆', '伍', '陆', '柒', '捌', '玖', '拾']
    for j in range(number):
        if random.choice([0, 1, 1, 1, 1]):
            # 数字序号
            no = random.randint(0, 4)
        else:
            # 汉字序号
            no = random.choice(chn_no)
        no_list.append(str(no)+"\n")

    return no_list


if __name__ == '__main__':
    create_mix_txt()

    # with open("../../train_data/bidi_data/mix_data4/rec_gt_test.txt", "r") as f:
    #     _list = f.readlines()
    # print(len(_list))
    # str1 = _list[0].split("\t")[-1][:-1]
    # str2 = _list[200000].split("\t")[-1][:-1]
    # str3 = _list[-1].split("\t")[-1][:-1]
    #
    # for i in range(2):
    #     create_image("../../train_data/bidi_data/mix_data2/", "", str1)
    #     create_image("../../train_data/bidi_data/mix_data2/", "", str2)
    #     create_image("../../train_data/bidi_data/mix_data2/", "", str3)