datalab-to--surya
wehub-resource-sync
58 行
1.6 KiB
Python
58 行
1.6 KiB
Python
from PIL import Image, ImageDraw
|
||
|
||
|
||
def test_table_rec(table_rec_predictor):
|
||
data = [
|
||
["Name", "Age", "City"],
|
||
["Alice", 25, "New York"],
|
||
["Bob", 30, "Los Angeles"],
|
||
["Charlie", 35, "Chicago"],
|
||
]
|
||
test_image = draw_table(data)
|
||
|
||
results = table_rec_predictor([test_image])
|
||
assert len(results) == 1
|
||
assert results[0].image_bbox == [0, 0, test_image.size[0], test_image.size[1]]
|
||
if results[0].error:
|
||
return
|
||
|
||
rows = results[0].rows
|
||
cols = results[0].cols
|
||
cells = results[0].cells
|
||
assert len(rows) >= 1
|
||
assert len(cols) >= 1
|
||
# Geometric cells = rows × cols
|
||
assert len(cells) == len(rows) * len(cols) or len(cells) <= len(rows) * len(cols)
|
||
|
||
|
||
def draw_table(data, cell_width=100, cell_height=40):
|
||
rows = len(data)
|
||
cols = len(data[0])
|
||
width = cols * cell_width
|
||
height = rows * cell_height
|
||
|
||
image = Image.new("RGB", (width, height), "white")
|
||
draw = ImageDraw.Draw(image)
|
||
|
||
for i in range(rows + 1):
|
||
y = i * cell_height
|
||
draw.line([(0, y), (width, y)], fill="black", width=1)
|
||
|
||
for i in range(cols + 1):
|
||
x = i * cell_width
|
||
draw.line([(x, 0), (x, height)], fill="black", width=1)
|
||
|
||
for i in range(rows):
|
||
for j in range(cols):
|
||
text = str(data[i][j])
|
||
text_bbox = draw.textbbox((0, 0), text)
|
||
text_width = text_bbox[2] - text_bbox[0]
|
||
text_height = text_bbox[3] - text_bbox[1]
|
||
|
||
x = j * cell_width + (cell_width - text_width) // 2
|
||
y = i * cell_height + (cell_height - text_height) // 2
|
||
|
||
draw.text((x, y), text, fill="black")
|
||
|
||
return image
|