TABLE_DETECTION/data_process/create_labelme_data.py

import base64
import copy
import json
import random

import cv2
import numpy as np
from PIL import Image, ImageFont, ImageDraw


def create_table_image(if_add_seal=False):
    lines_of_table = []
    outline_of_table = []

    # 空图
    # weight = random.randint(512, 896)
    height = random.randint(512, 640)
    weight = random.randint(512, 768)
    weight_origion = weight
    height_origion = height
    img = np.zeros((weight, height), np.uint8)
    img.fill(255)

    # 画矩形表格轮廓
    num = random.choice([1, 1, 2])
    num = 1
    scale_x = random.randint(10, int(weight/5))
    # scale_y = random.randint(25, int(height/5))
    scale_y = random.randint(int(height/12), int(height/10))
    if num == 1:
        x1 = scale_x
        y1 = scale_y
        x2 = scale_x + (weight - 2 * scale_x)
        y2 = scale_y + (height - 2 * scale_y)
        img = cv2.rectangle(img, (y1, x1), (y2, x2), (0, 0, 0), 1)
        lines_of_table.extend([[(y1, x1), (y2, x1)], [(y2, x1), (y2, x2)],
                               [(y2, x2), (y1, x2)], [(y1, x2), (y1, x1)]])
    if num == 2:
        scale_num = random.choice([0.4, 0.5, 0.6, 0.7])
        x1 = scale_x
        y1 = scale_y
        x2 = int((x1 + (weight - 2 * scale_x)) * scale_num)
        y2 = (y1 + (height - 2 * scale_y))
        # print(y1, x1, y2, x2)
        img = cv2.rectangle(img, (y1, x1), (y2, x2), (0, 0, 0), 1)
        lines_of_table.extend([[(y1, x1), (y2, x1)], [(y2, x1), (y2, x2)],
                               [(y2, x2), (y1, x2)], [(y1, x2), (y1, x1)]])

        weight_origin = weight
        weight = weight-x2
        scale_x = random.randint(10, int(weight/5))
        scale_y = random.randint(25, int(height/5))
        x3 = scale_x + x2
        y3 = scale_y
        x4 = x3 + (weight - 2 * scale_x)
        y4 = y3 + (height - 2 * scale_y)
        # print(weight_origin)
        # print(y3, x3, y4, x4)
        if x2 + 10 <= x3 or y4 - y3 >= 20 or x4 - 10 <= weight_origin:
            img = cv2.rectangle(img, (y3, x3), (y4, x4), (0, 0, 0), 1)
            lines_of_table.extend([[(y3, x3), (y4, x3)], [(y4, x3), (y4, x4)],
                                   [(y4, x4), (y3, x4)], [(y3, x4), (y3, x3)]])

    # 画表格内单元格线
    # 第一个表格
    row_num = random.randint(23, 25)
    col_num = random.randint(3, 4)
    margin_x = int((x2 - x1)/row_num)
    margin_y = int((y2 - y1)/col_num)
    col_points = []
    row_points = []
    # for rn in range(0, row_num):
    #     if y1+margin_y*rn < y2 - margin_y:
    #         col_points.append(y1+margin_y*rn)
    # for cn in range(0, col_num):
    #     if x1+margin_x*cn < x2 - margin_x:
    #         row_points.append(x1+margin_x*cn)

    # 紧凑表格制作
    print("margin_y, margin_x", margin_y, margin_x)
    last_x = copy.copy(x1)
    row_points.append(last_x)
    for rn in range(0, row_num):
        last_x += margin_x
        if last_x < x2 - margin_x:
            row_points.append(last_x)
    print("row_points", row_points)
    last_y = copy.copy(y1)
    col_points.append(last_y)
    for cn in range(0, col_num):
        last_y += margin_y
        if last_y < y2 - margin_y:
            col_points.append(last_y)
    print("col_points", col_points)

    # 画线
    for col in col_points:
        # if random.choice([0, 1, 1, 1]):
        # 不跨行
        if random.choice([1, 1, 1, 1]):
            img = cv2.line(img, (col, x1), (col, x2), (0, 0, 0), 1)
            lines_of_table.append([(col, x1), (col, x2)])
        else:
            img = cv2.line(img, (col, x1+margin_x), (col, x2), (0, 0, 0), 1)
            lines_of_table.append([(col, x1+margin_x), (col, x2)])
    for row in row_points:
        # if random.choice([0, 1, 1, 1]):
        if random.choice([1, 1, 1, 1]):
            img = cv2.line(img, (y1, row), (y2, row), (0, 0, 0), 1)
            lines_of_table.append([(y1, row), (y2, row)])
        else:
            img = cv2.line(img, (y1+margin_y, row), (y2, row), (0, 0, 0), 1)
            lines_of_table.append([(y1+margin_y, row), (y2, row)])

    # 第二个表格
    if num == 2:
        row_num2 = random.randint(6, 7)
        col_num2 = random.randint(6, 7)
        margin_x2 = int((x4 - x3)/row_num2)
        margin_y2 = int((y4 - y3)/col_num2)
        col_points2 = []
        row_points2 = []
        for rn in range(0, row_num2):
            if y3+margin_y2*rn < y4 - margin_y2:
                col_points2.append(y3+margin_y2*rn)
        for cn in range(0, col_num2):
            if x3+margin_x2*cn < x4 - margin_x2:
                row_points2.append(x3+margin_x2*cn)

        # 画线
        for col in col_points2:
            if random.choice([0, 1, 1, 1]):
                img = cv2.line(img, (col, x3), (col, x4), (0, 0, 0), 1)
                lines_of_table.append([(col, x3), (col, x4)])
            else:
                img = cv2.line(img, (col, x3+margin_x2), (col, x4), (0, 0, 0), 1)
                lines_of_table.append([(col, x3+margin_x2), (col, x4)])
        for row in row_points2:
            if random.choice([0, 1, 1, 1]):
                img = cv2.line(img, (y3, row), (y4, row), (0, 0, 0), 1)
                lines_of_table.append([(y3, row), (y4, row)])
            else:
                img = cv2.line(img, (y3+margin_y2, row), (y4, row), (0, 0, 0), 1)
                lines_of_table.append([(y3+margin_y2, row), (y4, row)])

    # 画轮廓
    point1 = (y1 - 5, x1 - 5)
    point2 = (y2 + 5, x2 + 5)
    # img = cv2.rectangle(img, point1, point2, (0, 0, 255), 2)
    outline_of_table.append([point1, point2])

    if num == 2:
        point3 = (y3 - 5, x3 - 5)
        point4 = (y4 + 5, x4 + 5)
        # img = cv2.rectangle(img, point3, point4, (0, 0, 255), 2)
        outline_of_table.append([point3, point4])

    # 表格中填字
    font_list = ["卡", "啊", "的", "我", "你", "吧", "为", "看", "他", "个", "来", "哦", "啊"]
    font_num = random.randint(7, 10)
    text = "".join(random.sample(font_list, font_num))
    # print(text)
    # text_size = int(margin_y / 1.6) if margin_y < margin_x else int(margin_x) / 1.6
    # font_size = text_size / font_num
    font_size = 10
    text_point = []

    image_pil = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
    for i in range(100):
        # text_x = random.choice(col_points) + int(margin_y/5)
        # text_y = random.choice(row_points) + int(margin_x/5)
        text_x = random.choice(col_points)
        text_y = random.choice(row_points)
        if (text_x, text_y) not in text_point:
            text_point.append((text_x, text_y))
            font = ImageFont.truetype("msyh.ttc", int(font_size))
            dr = ImageDraw.Draw(image_pil)

            bias_x = random.choice([0, 0, 0, 15, 30])
            dr.text((text_x+bias_x, text_y), text, font=font, fill="#000000")

    img = cv2.cvtColor(np.asarray(image_pil), cv2.COLOR_RGB2BGR)

    if num == 2:
        # text_size = int(margin_y2 / 1.6) if margin_y2 < margin_x2 else int(margin_x2) / 1.6
        # font_size = text_size / font_num
        font_size = 15
        text_point = []

        image_pil = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
        for i in range(100):
            text_x = random.choice(col_points2) + int(margin_y2/5)
            text_y = random.choice(row_points2) + int(margin_x2/5)
            if (text_x, text_y) not in text_point:
                text_point.append((text_x, text_y))
                font = ImageFont.truetype("msyh.ttc", int(font_size))
                dr = ImageDraw.Draw(image_pil)
                dr.text((text_x, text_y), text, font=font, fill="#000000")

    img = cv2.cvtColor(np.asarray(image_pil), cv2.COLOR_RGB2BGR)

    # 加印章
    if if_add_seal:
        img = create_official_seal(img)
        # cv2.imshow("add_seal", img)
        # cv2.waitKey(0)

    # 获取图片的二进制格式数据
    image_bytes = np.array(cv2.imencode('.jpg', img)[1]).tobytes()
    # with open('temp.jpg', 'wb') as f:
    #     f.write(image_bytes)

    # 显示
    # cv2.imshow("image", img)
    # cv2.waitKey(0)
    return lines_of_table, outline_of_table, image_bytes, weight_origion, height_origion


def create_outline_labelme(outline_list, image_bytes):
    labelme_data = {}
    shapes_list = []
    for outline in outline_list:
        shape_dict = {"label": "table",
                      "points": [[outline[0][0], outline[0][1]], [outline[1][0], outline[1][1]]],
                      "shape_type": "rectangle"
                      }
        shapes_list.append(shape_dict)
    labelme_data["shapes"] = shapes_list
    image_base64 = base64.b64encode(image_bytes)
    labelme_data["imageData"] = str(image_base64)
    return labelme_data


def create_lines_labelme(line_list, image_bytes, weight, height):
    labelme_data = {}
    shapes_list = []
    for line in line_list:
        if abs(line[1][0] - line[0][0]) >= abs(line[1][1] - line[0][1]):
            shape_dict = {"label": "0",
                          "points": [[line[0][0], line[0][1]], [line[1][0], line[1][1]]],
                          "shape_type": "line"
                          }
        else:
            shape_dict = {"label": "1",
                          "points": [[line[0][0], line[0][1]], [line[1][0], line[1][1]]],
                          "shape_type": "line"
                          }
        shapes_list.append(shape_dict)
    labelme_data["shapes"] = shapes_list
    # image_str = str(image_bytes, encoding='utf-8')
    image_base64_bytes = base64.b64encode(image_bytes)
    image_base64_string = image_base64_bytes.decode('utf-8')
    labelme_data["imageData"] = image_base64_string
    labelme_data["imageHeight"] = height
    labelme_data["imageWidth"] = weight
    return labelme_data


def create_official_seal(main_image_np):
    img1 = main_image_np

    # 随机选择公章
    seal_no = random.randint(0, 8)
    img2 = cv2.imread("../seal" + str(seal_no) + ".jpg")

    rows1, cols1, channels1 = img1.shape

    # 随机公章大小
    width = random.randint(int(cols1/6), int(cols1-cols1/6))
    # width = random.randint(150, 250)
    # width = 150
    height = width
    if width > rows1:
        width = rows1
        height = width
    img2 = cv2.resize(img2, (width, height))
    rows2, cols2, channels2 = img2.shape

    # 随机选择公章放置位置
    row_bias = random.randint(int(rows1/6), int(rows1-rows1/6))
    col_bias = random.randint(int(cols1/6), int(cols1-cols1/6))
    # I want to put logo on top-left corner, So I create a ROI

    # roi = img1[rows1-rows2-row_bias:rows1-row_bias, cols1-cols2-col_bias:cols1-col_bias]

    if rows1 >= rows2+row_bias and cols1 >= cols2+col_bias:
        roi_height = [row_bias, rows2+row_bias]
        roi_width = [col_bias, cols2+col_bias]
    elif rows1 < rows2+row_bias and cols1 < cols2+col_bias:
        roi_height = [rows1-width, rows1]
        roi_width = [cols1-width, cols1]
    elif cols1 < cols2+col_bias:
        roi_height = [row_bias, rows2+row_bias]
        roi_width = [cols1-width, cols1]
    else:
        roi_height = [rows1-width, rows1]
        roi_width = [col_bias, cols2+col_bias]
    roi = img1[roi_height[0]:roi_height[1], roi_width[0]:roi_width[1]]

    # Now create a mask of logo and create its inverse mask also
    img2gray = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)
    ret, mask = cv2.threshold(img2gray, 200, 255, cv2.THRESH_BINARY)
    mask_inv = cv2.bitwise_not(mask)

    # Now black-out the area of logo in ROI
    print("row_bias, col_bias", row_bias, col_bias)
    print("img1.shape, img2.shape", img1.shape, img2.shape)
    print("roi.shape, mask.shape", roi.shape, mask.shape)
    img1_bg = cv2.bitwise_and(roi, roi, mask=mask)

    # Take only region of logo from logo image.
    img2_fg = cv2.bitwise_and(img2, img2, mask=mask_inv)

    # Put logo in ROI and modify the main image
    dst = cv2.add(img1_bg, img2_fg)

    img1[roi_height[0]:roi_height[1], roi_width[0]:roi_width[1]] = dst

    # cv2.imshow('res', img1)
    # cv2.waitKey(0)
    return img1


if __name__ == '__main__':
    # 生成 单元格线 数据
    for i in range(2601, 2801):
        if i % 100 == 0:
            print(i)
        lines, outlines, image_data, weight, height = create_table_image(if_add_seal=True)
        labelme = create_lines_labelme(lines, image_data, weight, height)
        with open('../train/dataset-line/6/train_' + str(i) + '.json', 'w') as f:
            json.dump(labelme, f)

    # main_image = cv2.imread("../13.png")
    # create_official_seal(main_image)