12 lines
601 B
Python
12 lines
601 B
Python
from cv_analysis.server.pipeline import make_image_analysis_pipeline
|
|
from cv_analysis.table_inference import infer_lines
|
|
|
|
def test_table_inference():
|
|
pl = make_image_analysis_pipeline(infer_lines)
|
|
with open("test/test_data/article.pdf", "rb") as f:
|
|
pdf_bytes = f.read()
|
|
vlp_mock = {"data": [{"page_idx": 1, "image_boxes": [{"label": "table", "x1": 0.1, "y1": 0.3, "x2": 0.4, "y2": 0.6}]}]}
|
|
output = list(pl(pdf_bytes, vlp_mock))
|
|
lines = output[0]["tableLines"]
|
|
assert len(lines) > 1
|
|
assert all(map(lambda item: sorted(item.keys())==['x1', 'x2', 'y1', 'y2'], lines)) |