diff --git a/TableGeneration/GenerateTable.py b/TableGeneration/GenerateTable.py
index 4ed75f8..0c3482f 100644
--- a/TableGeneration/GenerateTable.py
+++ b/TableGeneration/GenerateTable.py
@@ -1,40 +1,43 @@
import json
import os
-import sys
+import platform
import random
import string
-from PIL import Image
+import sys
from io import BytesIO
-from tqdm import tqdm
+
import numpy as np
+from PIL import Image
from selenium.webdriver.common.by import By
-from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
+from selenium.webdriver.support.ui import WebDriverWait
+from tqdm import tqdm
from TableGeneration.Table import Table
class GenerateTable:
- def __init__(self,
- output,
- ch_dict_path,
- en_dict_path,
- cell_box_type='cell',
- min_row=3,
- max_row=20,
- min_col=3,
- max_col=10,
- max_span_row_count=3,
- max_span_col_count=3,
- max_span_value=20,
- min_txt_len=2,
- max_txt_len=7,
- color_prob=0,
- cell_max_width=0,
- cell_max_height=0,
- brower='chrome',
- brower_width=1920,
- brower_height=1920):
+ def __init__(
+ self,
+ output,
+ ch_dict_path,
+ en_dict_path,
+ cell_box_type="cell",
+ min_row=3,
+ max_row=20,
+ min_col=3,
+ max_col=10,
+ max_span_row_count=3,
+ max_span_col_count=3,
+ max_span_value=20,
+ min_txt_len=2,
+ max_txt_len=7,
+ color_prob=0,
+ cell_max_width=0,
+ cell_max_height=0,
+ brower="chrome",
+ brower_width=1920,
+ brower_height=1920, ):
self.output = output # wheter to store images separately or not
self.ch_dict_path = ch_dict_path
self.en_dict_path = en_dict_path
@@ -55,21 +58,24 @@ def __init__(self,
self.brower_height = brower_height # brower height
self.brower_width = brower_width # brower width
- if self.brower == 'chrome':
+ if self.brower == "chrome":
from selenium.webdriver import Chrome as Brower
from selenium.webdriver import ChromeOptions as Options
else:
from selenium.webdriver import Firefox as Brower
from selenium.webdriver import FirefoxOptions as Options
opts = Options()
- opts.add_argument('--headless')
- opts.add_argument('--no-sandbox')
+ opts.add_argument("--headless")
+ opts.add_argument("--no-sandbox")
self.driver = Brower(options=opts)
+ self.is_macos = platform.system() == "Darwin"
+ self.ratio = 2
+
def gen_table_img(self, img_count):
os.makedirs(self.output, exist_ok=True)
f_gt = open(
- os.path.join(self.output, 'gt.txt'), encoding='utf-8', mode='w')
+ os.path.join(self.output, "gt.txt"), encoding="utf-8", mode="w")
for i in tqdm(range(img_count)):
# data_arr contains the images of generated tables and all_table_categories contains the table category of each of the table
out = self.generate_table()
@@ -80,30 +86,30 @@ def gen_table_img(self, img_count):
im, contens = self.clip_white(im, contens)
# randomly select a name of length=20 for file.
- output_file_name = ''.join(
+ output_file_name = "".join(
random.choices(
string.ascii_uppercase + string.digits, k=20))
- output_file_name = '{}_{}_{}'.format(border, i, output_file_name)
+ output_file_name = "{}_{}_{}".format(border, i, output_file_name)
# print('{}/{}, {}'.format(i, img_count, output_file_name))
# if the image and equivalent html is need to be stored
- os.makedirs(os.path.join(self.output, 'html'), exist_ok=True)
- os.makedirs(os.path.join(self.output, 'img'), exist_ok=True)
-
- html_save_path = os.path.join(self.output, 'html',
- output_file_name + '.html')
- img_save_path = os.path.join(self.output, 'img',
- output_file_name + '.jpg')
- with open(html_save_path, encoding='utf-8', mode='w') as f:
+ os.makedirs(os.path.join(self.output, "html"), exist_ok=True)
+ os.makedirs(os.path.join(self.output, "img"), exist_ok=True)
+
+ html_save_path = os.path.join(self.output, "html",
+ output_file_name + ".html")
+ img_save_path = os.path.join(self.output, "img",
+ output_file_name + ".jpg")
+ with open(html_save_path, encoding="utf-8", mode="w") as f:
f.write(html_content)
im.save(img_save_path, dpi=(600, 600))
# 构造标注信息
- img_file_name = os.path.join('img', output_file_name + '.jpg')
+ img_file_name = os.path.join("img", output_file_name + ".jpg")
label_info = self.make_ppstructure_label(structure, contens,
img_file_name)
- f_gt.write('{}\n'.format(
+ f_gt.write("{}\n".format(
json.dumps(
label_info, ensure_ascii=False)))
# convert to PP-Structure label format
@@ -116,12 +122,20 @@ def generate_table(self):
rows = random.randint(self.min_row, self.max_row)
try:
# initialize table class
- table = Table(self.ch_dict_path, self.en_dict_path,
- self.cell_box_type, rows, cols, self.min_txt_len,
- self.max_txt_len, self.max_span_row_count,
- self.max_span_col_count, self.max_span_value,
- self.color_prob, self.cell_max_width,
- self.cell_max_height)
+ table = Table(
+ self.ch_dict_path,
+ self.en_dict_path,
+ self.cell_box_type,
+ rows,
+ cols,
+ self.min_txt_len,
+ self.max_txt_len,
+ self.max_span_row_count,
+ self.max_span_col_count,
+ self.max_span_value,
+ self.color_prob,
+ self.cell_max_width,
+ self.cell_max_height, )
# get table of rows and cols based on unlv distribution and get features of this table
# (same row, col and cell matrices, total unique ids, html conversion of table and its category)
id_count, html_content, structure, border = table.create()
@@ -132,46 +146,49 @@ def generate_table(self):
return im, html_content, structure, contens, border
except KeyboardInterrupt:
import sys
+
sys.exit()
except:
import traceback
+
traceback.print_exc()
return None
return None
def make_ppstructure_label(self, structure, bboxes, img_path):
d = {
- 'filename': img_path,
- 'html': {
- 'structure': {
- 'tokens': structure
+ "filename": img_path,
+ "html": {
+ "structure": {
+ "tokens": structure
}
}
}
cells = []
for bbox in bboxes:
text = bbox[1]
- cells.append({'tokens': list(text), 'bbox': bbox[2:]})
- d['html']['cells'] = cells
- d['gt'] = self.rebuild_html_from_ppstructure_label(d)
+ cells.append({"tokens": list(text), "bbox": bbox[2:]})
+ d["html"]["cells"] = cells
+ d["gt"] = self.rebuild_html_from_ppstructure_label(d)
return d
def rebuild_html_from_ppstructure_label(self, label_info):
from html import escape
- html_code = label_info['html']['structure']['tokens'].copy()
+
+ html_code = label_info["html"]["structure"]["tokens"].copy()
to_insert = [
- i for i, tag in enumerate(html_code) if tag in ('| ', '>')
+ i for i, tag in enumerate(html_code) if tag in (" | ", ">")
]
- for i, cell in zip(to_insert[::-1], label_info['html']['cells'][::-1]):
- if cell['tokens']:
+ for i, cell in zip(to_insert[::-1], label_info["html"]["cells"][::-1]):
+ if cell["tokens"]:
cell = [
escape(token) if len(token) == 1 else token
- for token in cell['tokens']
+ for token in cell["tokens"]
]
- cell = ''.join(cell)
+ cell = "".join(cell)
html_code.insert(i + 1, cell)
- html_code = ''.join(html_code)
- html_code = ''.format(
+ html_code = "".join(html_code)
+ html_code = "".format(
html_code)
return html_code
@@ -196,7 +213,7 @@ def clip_white(self, im, bboxes):
return im, bboxes
def html_to_img(self, html_content, id_count):
- '''converts html to image'''
+ """converts html to image"""
self.driver.get("data:text/html;charset=utf-8," + html_content)
self.driver.maximize_window()
self.driver.set_window_size(
@@ -204,7 +221,11 @@ def html_to_img(self, html_content, id_count):
height=self.brower_height,
windowHandle="current")
window_size = self.driver.get_window_size()
- max_height, max_width = window_size['height'], window_size['width']
+ max_height, max_width = window_size["height"], window_size["width"]
+ if self.is_macos:
+ max_height *= self.ratio
+ max_width *= self.ratio
+
# element = WebDriverWait(self.driver, 3).until(EC.presence_of_element_located((By.ID, '0')))
contens = []
for id in range(id_count):
@@ -216,10 +237,16 @@ def html_to_img(self, html_content, id_count):
lentext = len(txt)
loc = e.location
size_ = e.size
- xmin = loc['x']
- ymin = loc['y']
- xmax = int(size_['width'] + xmin)
- ymax = int(size_['height'] + ymin)
+ xmin = loc["x"]
+ ymin = loc["y"]
+ xmax = int(size_["width"] + xmin)
+ ymax = int(size_["height"] + ymin)
+
+ if self.is_macos:
+ xmin *= self.ratio
+ ymin *= self.ratio
+ xmax *= self.ratio
+ ymax *= self.ratio
contens.append([
lentext, txt, [[xmin, ymin], [xmax, ymin], [xmax, ymax],
@@ -228,7 +255,7 @@ def html_to_img(self, html_content, id_count):
png = self.driver.get_screenshot_as_png()
- im = Image.open(BytesIO(png)).convert('RGB')
+ im = Image.open(BytesIO(png)).convert("RGB")
im = im.crop((0, 0, max_width, max_height))
return im, contens
|