import base64 import copy import json import random import cv2 import numpy as np from PIL import Image, ImageFont, ImageDraw def create_table_image2(if_add_seal=False): lines_of_table = [] outline_of_table = [] # 空图 # weight = random.randint(512, 896) height = random.randint(512, 2500) weight = random.randint(512, 1500) weight_origion = weight height_origion = height img = np.zeros((weight, height), np.uint8) img.fill(255) # 计算交点用,黑图 row_img = np.zeros((weight, height), np.uint8) col_img = np.zeros((weight, height), np.uint8) # 画矩形表格轮廓 num = random.choice([1, 1, 2]) num = 1 if random.choice([0]): 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)) else: scale_x = weight-2 scale_y = height-2 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) row_img = cv2.line(row_img, (y1, x1), (y2, x1), (255, 255, 255), 1) row_img = cv2.line(row_img, (y2, x2), (y1, x2), (255, 255, 255), 1) col_img = cv2.line(col_img, (y2, x1), (y2, x2), (255, 255, 255), 1) col_img = cv2.line(col_img, (y1, x2), (y1, x1), (255, 255, 255), 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) row_img = cv2.line(row_img, (y1, x1), (y2, x1), (255, 255, 255), 1) row_img = cv2.line(row_img, (y2, x2), (y1, x2), (255, 255, 255), 1) col_img = cv2.line(col_img, (y2, x1), (y2, x2), (255, 255, 255), 1) col_img = cv2.line(col_img, (y1, x2), (y1, x1), (255, 255, 255), 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) row_img = cv2.line(row_img, (y3, x3), (y4, x3), (255, 255, 255), 1) row_img = cv2.line(row_img, (y4, x4), (y3, x4), (255, 255, 255), 1) col_img = cv2.line(col_img, (y4, x3), (y4, x4), (255, 255, 255), 1) col_img = cv2.line(col_img, (y3, x4), (y3, x3), (255, 255, 255), 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, 6) margin_x = int((x2 - x1)/row_num) margin_y = int((y2 - y1)/col_num) print("margin_x", margin_x) print("margin_y", margin_y) if margin_x < 20: row_num = random.randint(11, 12) margin_x = int((x2 - x1)/row_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) col_img = cv2.line(col_img, (col, x1), (col, x2), (255, 255, 255), 1) lines_of_table.append([(col, x1), (col, x2)]) else: img = cv2.line(img, (col, x1+margin_x), (col, x2), (0, 0, 0), 1) col_img = cv2.line(col_img, (col, x1+margin_x), (col, x2), (255, 255, 255), 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) row_img = cv2.line(row_img, (y1, row), (y2, row), (255, 255, 255), 1) lines_of_table.append([(y1, row), (y2, row)]) else: img = cv2.line(img, (y1+margin_y, row), (y2, row), (0, 0, 0), 1) row_img = cv2.line(row_img, (y1+margin_y, row), (y2, row), (255, 255, 255), 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) col_img = cv2.line(col_img, (col, x3), (col, x4), (255, 255, 255), 1) lines_of_table.append([(col, x3), (col, x4)]) else: img = cv2.line(img, (col, x3+margin_x2), (col, x4), (0, 0, 0), 1) col_img = cv2.line(col_img, (col, x3+margin_x2), (col, x4), (255, 255, 255), 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) row_img = cv2.line(row_img, (y3, row), (y4, row), (255, 255, 255), 1) lines_of_table.append([(y3, row), (y4, row)]) else: img = cv2.line(img, (y3+margin_y2, row), (y4, row), (0, 0, 0), 1) row_img = cv2.line(row_img, (y3+margin_y2, row), (y4, row), (255, 255, 255), 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]) # 计算交点 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])) print("len(points)", len(points)) # 查看交点 # for p in points: # cv2.circle(img, p, 2, (0, 0, 255)) # cv2.imshow("points", img) # cv2.waitKey(0) # 表格中填字 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, 50, 70]) bias_y = random.choice([0, 10, 10, 15, 20, 30]) dr.text((text_x+bias_x, text_y+bias_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 random.choice([0, 1]): sigmaX = random.randint(1, 10) sigmaY = random.randint(1, 10) img = cv2.GaussianBlur(img, (5, 5), sigmaX, sigmaY) # 加印章 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, points def create_table_image(): def get_margin(h, w): top = random.randint(10, int(h / 3)) bottom = random.randint(10, int(h / 3)) left = random.randint(10, int(w / 3)) right = random.randint(10, int(w / 3)) line_list = [] outline_list = [] # 空图 height = random.randint(512, 2500) width = random.randint(512, 1500) img = np.zeros((width, height), np.uint8) img.fill(255) # 计算交点用,黑图 row_img = np.zeros((width, height), np.uint8) col_img = np.zeros((width, height), np.uint8) # 设置表格数量 table_num = random.choice([1, 2]) # 设置表格outline margin_top = random.randint(10, int(height / 3)) margin_bottom = random.randint(10, int(height / 3)) margin_left = random.randint(10, int(width / 3)) margin_right = random.randint(10, int(width / 3)) up_line = [(margin_left, margin_top), (width - margin_right, margin_top)] bottom_line = [(margin_left, height-margin_bottom), (width - margin_right, height - margin_bottom)] left_line = [(margin_left, margin_top), (margin_left, height-margin_bottom)] right_line = [(width - margin_right, margin_top), (width - margin_right, height - margin_bottom)] outline_list.extend([up_line, bottom_line, left_line, right_line]) # 设置 return 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, cross_points): 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 labelme_data["cross_points"] = str(cross_points) return labelme_data def create_official_seal(main_image_np): img1 = main_image_np # 随机选择公章 seal_no = random.randint(0, 8) seal_no = 0 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(2500, 5000): if i % 100 == 0: print("Loop", i) lines, outlines, image_data, weight, height, cross_points = create_table_image(if_add_seal=False) labelme = create_lines_labelme(lines, image_data, weight, height, cross_points) with open('../train/dataset-line/7/train_' + str(i) + '.json', 'w') as f: json.dump(labelme, f) # main_image = cv2.imread("../13.png") # create_official_seal(main_image)