cv-analysis-service/test/unit_tests/table_inference_test.py

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))