import json
import random
import numpy as np
import cv2
import pandas as pd
import imgkit
import time
from PIL import Image
from selenium.webdriver.chrome.options import Options
from selenium import webdriver
from create_labelme_data import create_lines_labelme, create_official_seal
DESKTOP_PATH = r"C:\Users\Administrator\Desktop"
def read_csv(path):
df = pd.read_csv(path, encoding='gbk')
df = df[:2000]
_list = []
for index, row in df.iterrows():
if type(row['json_table']) == float:
print(row['json_table'])
continue
_list.append(json.loads(row['json_table']))
# print(_list)
return _list
def create_html_table(_list, html_dir, count=0):
html_path_list = []
i = count
for table in _list:
table_text = '' + "\n"
for row in table:
table_text += "" + "\n"
for col in row:
table_text += "| " + col[0] + " | " + "\n"
table_text += "
" + "\n"
table_text += "
" + "\n"
css_text = ''
# 虚线框
# css_text = ''
# 粗虚线 上下或左右无边线
# css_text = ''
# css_text = ''
html_text = "" + "\n" \
+ "" + "\n" \
+ css_text + "\n" \
+ '' + "\n" \
+ '' + "\n" \
+ table_text + "\n" \
+ '
' + "\n" \
+ ""
html_path = html_dir + str(i) + ".html"
with open(html_path, 'w') as f:
f.write(html_text)
html_path_list.append(html_path)
i += 1
return html_path_list
def get_table_size(html_path):
print(html_path)
chrome_options = Options()
chrome_options.add_argument('--headless')
chrome_options.add_argument('--no-sandbox')
driver = webdriver.Chrome(chrome_options=chrome_options)
# driver = webdriver.Chrome()
driver.maximize_window()
driver.implicitly_wait(0)
driver.get(html_path)
# JS
js = """
_t = document.getElementById('table0')
function myFunc(){
var cells_size = new Array();
for (i = 0; i<_t.rows.length; i++){
var rows_size = new Array();
for (j = 0; j < _t.rows[i].cells.length; j++){
col = _t.rows[i].cells[j];
col_width = window.getComputedStyle(col).width;
col_height = window.getComputedStyle(col).height;
rows_size[j] = [col_width, col_height];
}
cells_size[i] = rows_size
}
return cells_size;
}
return myFunc();
"""
rows_list = driver.execute_script(js)
# print(rows_list)
# print(type(rows_list))
js = """
_t = document.getElementById('table0')
function myFunc(){
table_width = window.getComputedStyle(_t).width;
table_height = window.getComputedStyle(_t).height;
var table_size = [table_width, table_height];
return table_size
}
return myFunc()
"""
table_size = driver.execute_script(js)
table_size = [int(float(table_size[0][:-2])), int(float(table_size[1][:-2]))]
print("table_size", table_size)
# get_table_lines(rows_list, table_size)
# table = driver.find_element_by_id('table0')
# print("rows_list", rows_list)
return rows_list, table_size
def html2image1(html_path, table_size):
# 工具路径
path_wkimg = r'D:\Software\html_to_pdf\wkhtmltopdf\bin\wkhtmltoimage.exe'
cfg = imgkit.config(wkhtmltoimage=path_wkimg)
options = {
'width': table_size[0],
'height': table_size[1],
'encoding': 'UTF-8',
}
# 1、将html文件转为图片
image_path = html_path.split(".")[0]+".jpg"
print("html2image", image_path)
imgkit.from_file(html_path, image_path, config=cfg, options=options)
# 2、从url获取html，再转为图片
# imgkit.from_url('https://httpbin.org/ip', 'ip.jpg', config=cfg)
# 3、将字符串转为图片
# imgkit.from_string('Hello!','hello.jpg', config=cfg)
return image_path
def html2image(html_path, table_size):
# 将html文件转为图片
image_path = html_path.split(".")[0]+".jpg"
print("html2image", image_path)
chrome_options = Options()
chrome_options.add_argument('--headless')
chrome_options.add_argument('--no-sandbox')
broswer = webdriver.Chrome(chrome_options=chrome_options)
broswer.get(html_path)
# 截全图
width = broswer.execute_script("return document.documentElement.scrollWidth")
height = broswer.execute_script("return document.documentElement.scrollHeight")
broswer.set_window_size(width, height)
broswer.save_screenshot(image_path)
table = broswer.find_element_by_id('table0')
left = table.location['x']
top = table.location['y']
elementWidth = table.location['x'] + table.size['width']
elementHeight = table.location['y'] + table.size['height']
picture = Image.open(image_path)
picture = picture.crop((left, top, elementWidth, elementHeight))
picture = picture.convert("RGB")
picture.save(image_path)
return image_path
def get_table_lines(rows_list):
row_line_list = []
col_line_list = []
x = 0
y = 0
# 横线line
width = 0
height = 0
i = 0
for row in rows_list:
if i == 0:
for col in row:
width += float(col[0][:-2]) + 1
row_line_list.append([[x, y], [x+width, y]])
height += float(row[0][1][:-2]) + 1
else:
row_line_list.append([[x, y+height], [x+width, y+height]])
height += float(row[0][1][:-2]) + 1
i += 1
row_line_list.append([[x, y+height], [x+width, y+height]])
# 竖线line
width = 0
height = 0
i = 0
for col_num in range(len(rows_list)):
height += float(rows_list[col_num][0][1][:-2]) + 1
# print("height", height)
for row in rows_list:
if i == 0:
col_line_list.append([[x, y], [x, y+height]])
for col in row:
width += float(col[0][:-2]) + 1
col_line_list.append([[x+width, y], [x+width, y+height]])
break
# print("row_line_list", row_line_list)
# print("col_line_list", col_line_list)
# draw_lines(row_line_list+col_line_list, table_size, )
return row_line_list+col_line_list
def draw_lines(line_list, table_size, image_path, expand=False):
img = np.zeros((table_size[1], table_size[0]), np.uint8)
img.fill(255)
if expand:
image_origin = cv2.imread(image_path)
# print(image_origin.shape, img.shape)
# expand_height = int((image_origin.shape[1] - img.shape[1]) / 2)
# expand_width = int((image_origin.shape[0] - img.shape[0]) / 2)
# img = cv2.copyMakeBorder(img, expand_height, expand_height, expand_width,
# expand_width, cv2.BORDER_CONSTANT, value=(255, 255, 255))
img = np.zeros((image_origin.shape[0], image_origin.shape[1]), np.uint8)
print(image_origin.shape, img.shape)
img.fill(255)
for line in line_list:
cv2.line(img, (int(line[0][0]), int(line[0][1])),
(int(line[1][0]), int(line[1][1])), (0, 0, 255), 1)
cv2.imwrite(image_path.split(".")[0] + ".png", img)
# cv2.imshow("label", img)
# cv2.waitKey(0)
# image = cv2.imread(image_path)
# cv2.imshow("image", image)
# cv2.waitKey(0)
return line_list
def image_expand(image_path, line_list):
# image_np = cv2.imdecode(np.frombuffer(image_bytes, np.uint8), cv2.IMREAD_COLOR)
# 随机选择边缘扩充px
# expand_height = random.randint(0, 500)
# expand_width = int(expand_height / 1.3333)
image_origin = cv2.imread(image_path)
expand_height = int((1123 - image_origin.shape[0]) / 2)
expand_width = int((794 - image_origin.shape[1]) / 2)
if expand_width < 0:
expand_width = 0
if expand_height < 0:
expand_height = 0
# print(expand_height, expand_width)
# 图像边缘扩充
image_np = cv2.imread(image_path)
image_np = cv2.copyMakeBorder(image_np, expand_height, expand_height, expand_width,
expand_width, cv2.BORDER_CONSTANT, value=(255, 255, 255))
cv2.imwrite(image_path, image_np)
# 线条坐标全部加上增加的宽高
new_line_list = []
for line in line_list:
new_line_list.append([[line[0][0]+expand_width, line[0][1]+expand_height],
[line[1][0]+expand_width, line[1][1]+expand_height]])
return image_path, new_line_list
if __name__ == '__main__':
csv_path = "D:\\BIDI_DOC\\比地_文档\\websource_67000_table.csv"
table_list = read_csv(csv_path)
i = 500
stop_i = 700
table_list = table_list[i:]
save_dir = "D:\\Project\\table-detect-master\\data_process\\create_data\\"
html_path_list = create_html_table(table_list, save_dir, i)
for html in html_path_list:
print("Loop", i)
rows_list, table_size = get_table_size(html)
image_path = html2image(html, table_size)
line_list = get_table_lines(rows_list)
# 图片扩展
# image_path, line_list = image_expand(image_path, line_list)
# 添加公章
image_np = cv2.imread(image_path)
image_np = create_official_seal(image_np)
cv2.imwrite(image_path, image_np)
with open(image_path, 'rb') as f:
image_bytes = f.read()
image_np = cv2.imread(image_path)
labelme = create_lines_labelme(line_list, image_bytes, image_np.shape[1], image_np.shape[0])
with open('../train/dataset-line/6/train_' + str(i) + '.json', 'w') as f:
json.dump(labelme, f)
draw_lines(line_list, table_size, image_path, False)
i += 1
if i > stop_i:
break
# break
# img_path = DESKTOP_PATH + '/1.jpg'
# get_table_lines(img_path)