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

import json
from os.path import join

import pytest
from funcy import first

from cv_analysis.locations import TEST_DATA_DIR
from cv_analysis.table_parsing import parse_tables
from cv_analysis.utils.preprocessing import open_pdf
from cv_analysis.utils.test_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(score_threshold, image_with_tables, expected_table_annotation, test_file_index):

    result = [x.json_xywh() for x in parse_tables(image_with_tables)]
    formatted_result = {"pages": [{"page": str(test_file_index), "cells": result}]}

    score = compute_document_score(formatted_result, expected_table_annotation)

    assert round(score, 3) >= score_threshold


@pytest.fixture
def image_with_tables(test_file_index):
    img_path = join(TEST_DATA_DIR, f"test{test_file_index}.png")
    return first(open_pdf(img_path))


@pytest.fixture
def expected_table_annotation(test_file_index):
    json_path = join(TEST_DATA_DIR, f"test{test_file_index}.json")
    with open(json_path) as f:
        return json.load(f)