75 lines
2.8 KiB
Python
75 lines
2.8 KiB
Python
from itertools import starmap
|
|
|
|
import cv2
|
|
import pytest
|
|
from funcy import lmap, compose, zipdict
|
|
|
|
from cv_analysis.table_parsing import parse_tables
|
|
from cv_analysis.utils import lift
|
|
from cv_analysis.utils.rectangle import Rectangle
|
|
from cv_analysis.utils.metrics import compute_document_score
|
|
|
|
|
|
@pytest.mark.parametrize("score_threshold", [0.95])
|
|
@pytest.mark.parametrize("test_file_index", range(1, 11))
|
|
def test_table_parsing_on_client_pages(
|
|
score_threshold, client_page_with_table, expected_table_annotation, test_file_index
|
|
):
|
|
|
|
results = compose(lift(rectangle_to_dict), parse_tables)(client_page_with_table)
|
|
formatted_result = {"pages": [{"cells": results}]}
|
|
|
|
score = compute_document_score(formatted_result, expected_table_annotation)
|
|
|
|
assert round(score, 3) >= score_threshold
|
|
|
|
|
|
@pytest.fixture
|
|
def error_tolerance(line_thickness):
|
|
return line_thickness * 7
|
|
|
|
|
|
def rectangle_to_dict(rectangle: Rectangle):
|
|
return zipdict(["x", "y", "width", "height"], rectangle_to_xywh(rectangle))
|
|
|
|
|
|
def rectangle_to_xywh(rectangle: Rectangle):
|
|
return rectangle.x1, rectangle.y1, abs(rectangle.x1 - rectangle.x2), abs(rectangle.y1 - rectangle.y2)
|
|
|
|
|
|
@pytest.mark.parametrize("line_thickness", [1, 2, 3])
|
|
@pytest.mark.parametrize("line_type", [cv2.LINE_4, cv2.LINE_AA, cv2.LINE_8])
|
|
@pytest.mark.parametrize("table_style", ["closed horizontal vertical", "open horizontal vertical"])
|
|
@pytest.mark.parametrize("n_tables", [1, 2])
|
|
@pytest.mark.parametrize("background_color", [255, 220])
|
|
@pytest.mark.parametrize("table_shape", [(5, 8)])
|
|
def test_table_parsing_on_generic_pages(page_with_table, expected_gold_page_with_table, error_tolerance):
|
|
result = lmap(rectangle_to_xywh, parse_tables(page_with_table))
|
|
assert (
|
|
result == expected_gold_page_with_table
|
|
or average_error(result, expected_gold_page_with_table) <= error_tolerance
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("line_thickness", [1, 2, 3])
|
|
@pytest.mark.parametrize("line_type", [cv2.LINE_4, cv2.LINE_AA, cv2.LINE_8])
|
|
@pytest.mark.parametrize("table_style", ["closed horizontal vertical", "open horizontal vertical"])
|
|
@pytest.mark.parametrize("n_tables", [1, 2])
|
|
@pytest.mark.parametrize("background_color", [255, 220])
|
|
@pytest.mark.parametrize("table_shape", [(5, 8)])
|
|
@pytest.mark.xfail
|
|
def test_low_quality_table(page_with_patchy_table, expected_gold_page_with_table, error_tolerance):
|
|
result = lmap(rectangle_to_xywh, parse_tables(page_with_patchy_table))
|
|
assert (
|
|
result == expected_gold_page_with_table
|
|
or average_error(result, expected_gold_page_with_table) <= error_tolerance
|
|
)
|
|
|
|
|
|
def average_error(result, expected):
|
|
return sum(starmap(calc_rect_diff, zip(result, expected))) / len(expected)
|
|
|
|
|
|
def calc_rect_diff(rect1, rect2):
|
|
return sum(abs(c1 - c2) for c1, c2 in zip(rect1, rect2))
|