Skip to content

Commit

Permalink
update tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@felixdittrich92
felixdittrich92 committed Oct 4, 2023
1 parent 1edd231 commit 28f4981
Show file tree
Hide file tree
Showing 2 changed files with 58 additions and 47 deletions.
53 changes: 29 additions & 24 deletions tests/pytorch/test_models_zoo_pt.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -73,10 +73,17 @@ def test_ocrpredictor(mock_pdf, mock_vocab, assume_straight_pages, straighten_pa
assert out.pages[0].orientation["value"] == orientation


def test_trained_ocr_predictor(mock_tilted_payslip):
doc = DocumentFile.from_images(mock_tilted_payslip)
def test_trained_ocr_predictor(mock_payslip):
doc = DocumentFile.from_images(mock_payslip)

det_predictor = detection_predictor("db_resnet50", pretrained=True, batch_size=2, assume_straight_pages=True)
det_predictor = detection_predictor(
"db_resnet50",
pretrained=True,
batch_size=2,
assume_straight_pages=True,
symmetric_pad=True,
preserve_aspect_ratio=False,
)
reco_predictor = recognition_predictor("crnn_vgg16_bn", pretrained=True, batch_size=128)

predictor = OCRPredictor(
Expand All @@ -90,16 +97,12 @@ def test_trained_ocr_predictor(mock_tilted_payslip):
out = predictor(doc)

assert out.pages[0].blocks[0].lines[0].words[0].value == "Mr."
geometry_mr = np.array(
[[0.08563021, 0.35584526], [0.11464554, 0.34078913], [0.1274898, 0.36012764], [0.09847447, 0.37518377]]
)
assert np.allclose(np.array(out.pages[0].blocks[0].lines[0].words[0].geometry), geometry_mr)
geometry_mr = np.array([[0.1083984375, 0.0634765625], [0.1494140625, 0.0859375]])
assert np.allclose(np.array(out.pages[0].blocks[0].lines[0].words[0].geometry), geometry_mr, rtol=0.05)

assert out.pages[0].blocks[1].lines[0].words[-1].value == "revised"
geometry_revised = np.array(
[[0.50422498, 0.19551784], [0.55741975, 0.16791493], [0.56705294, 0.18241881], [0.51385817, 0.21002172]]
)
assert np.allclose(np.array(out.pages[0].blocks[1].lines[0].words[-1].geometry), geometry_revised)
geometry_revised = np.array([[0.7548828125, 0.126953125], [0.8388671875, 0.1484375]])
assert np.allclose(np.array(out.pages[0].blocks[1].lines[0].words[-1].geometry), geometry_revised, rtol=0.05)

det_predictor = detection_predictor(
"db_resnet50",
Expand Down Expand Up @@ -181,10 +184,17 @@ def test_kiepredictor(mock_pdf, mock_vocab, assume_straight_pages, straighten_pa
assert out.pages[0].orientation["value"] == orientation


def test_trained_kie_predictor(mock_tilted_payslip):
doc = DocumentFile.from_images(mock_tilted_payslip)
def test_trained_kie_predictor(mock_payslip):
doc = DocumentFile.from_images(mock_payslip)

det_predictor = detection_predictor("db_resnet50", pretrained=True, batch_size=2, assume_straight_pages=True)
det_predictor = detection_predictor(
"db_resnet50",
pretrained=True,
batch_size=2,
assume_straight_pages=True,
symmetric_pad=True,
preserve_aspect_ratio=False,
)
reco_predictor = recognition_predictor("crnn_vgg16_bn", pretrained=True, batch_size=128)

predictor = KIEPredictor(
Expand All @@ -199,17 +209,12 @@ def test_trained_kie_predictor(mock_tilted_payslip):

assert isinstance(out, KIEDocument)
assert out.pages[0].predictions[CLASS_NAME][0].value == "Mr."
geometry_mr = np.array(
[[0.08563021, 0.35584526], [0.11464554, 0.34078913], [0.1274898, 0.36012764], [0.09847447, 0.37518377]]
)
assert np.allclose(np.array(out.pages[0].predictions[CLASS_NAME][0].geometry), geometry_mr)
geometry_mr = np.array([[0.1083984375, 0.0634765625], [0.1494140625, 0.0859375]])
assert np.allclose(np.array(out.pages[0].predictions[CLASS_NAME][0].geometry), geometry_mr, rtol=0.05)

print(out.pages[0].predictions[CLASS_NAME])
assert out.pages[0].predictions[CLASS_NAME][7].value == "revised"
geometry_revised = np.array(
[[0.50422498, 0.19551784], [0.55741975, 0.16791493], [0.56705294, 0.18241881], [0.51385817, 0.21002172]]
)
assert np.allclose(np.array(out.pages[0].predictions[CLASS_NAME][7].geometry), geometry_revised)
assert out.pages[0].predictions[CLASS_NAME][6].value == "revised"
geometry_revised = np.array([[0.7548828125, 0.126953125], [0.8388671875, 0.1484375]])
assert np.allclose(np.array(out.pages[0].predictions[CLASS_NAME][6].geometry), geometry_revised, rtol=0.05)

det_predictor = detection_predictor(
"db_resnet50",
Expand Down
52 changes: 29 additions & 23 deletions tests/tensorflow/test_models_zoo_tf.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -72,10 +72,17 @@ def test_ocrpredictor(mock_pdf, mock_vocab, assume_straight_pages, straighten_pa
assert out.pages[0].language["value"] == language


def test_trained_ocr_predictor(mock_tilted_payslip):
doc = DocumentFile.from_images(mock_tilted_payslip)
def test_trained_ocr_predictor(mock_payslip):
doc = DocumentFile.from_images(mock_payslip)

det_predictor = detection_predictor("db_resnet50", pretrained=True, batch_size=2, assume_straight_pages=True)
det_predictor = detection_predictor(
"db_resnet50",
pretrained=True,
batch_size=2,
assume_straight_pages=True,
symmetric_pad=True,
preserve_aspect_ratio=False,
)
reco_predictor = recognition_predictor("crnn_vgg16_bn", pretrained=True, batch_size=128)

predictor = OCRPredictor(
Expand All @@ -89,16 +96,12 @@ def test_trained_ocr_predictor(mock_tilted_payslip):
out = predictor(doc)

assert out.pages[0].blocks[0].lines[0].words[0].value == "Mr."
geometry_mr = np.array(
[[0.08844472, 0.35763523], [0.11625107, 0.34320644], [0.12588427, 0.35771032], [0.09807791, 0.37213911]]
)
assert np.allclose(np.array(out.pages[0].blocks[0].lines[0].words[0].geometry), geometry_mr)
geometry_mr = np.array([[0.1083984375, 0.0634765625], [0.1494140625, 0.0859375]])
assert np.allclose(np.array(out.pages[0].blocks[0].lines[0].words[0].geometry), geometry_mr, rtol=0.05)

assert out.pages[0].blocks[1].lines[0].words[-1].value == "revised"
geometry_revised = np.array(
[[0.50422498, 0.19551784], [0.55741975, 0.16791493], [0.56705294, 0.18241881], [0.51385817, 0.21002172]]
)
assert np.allclose(np.array(out.pages[0].blocks[1].lines[0].words[-1].geometry), geometry_revised)
geometry_revised = np.array([[0.7548828125, 0.126953125], [0.8388671875, 0.1484375]])
assert np.allclose(np.array(out.pages[0].blocks[1].lines[0].words[-1].geometry), geometry_revised, rtol=0.05)

det_predictor = detection_predictor(
"db_resnet50",
Expand Down Expand Up @@ -179,10 +182,17 @@ def test_kiepredictor(mock_pdf, mock_vocab, assume_straight_pages, straighten_pa
assert out.pages[0].language["value"] == language


def test_trained_kie_predictor(mock_tilted_payslip):
doc = DocumentFile.from_images(mock_tilted_payslip)
def test_trained_kie_predictor(mock_payslip):
doc = DocumentFile.from_images(mock_payslip)

det_predictor = detection_predictor("db_resnet50", pretrained=True, batch_size=2, assume_straight_pages=True)
det_predictor = detection_predictor(
"db_resnet50",
pretrained=True,
batch_size=2,
assume_straight_pages=True,
symmetric_pad=True,
preserve_aspect_ratio=False,
)
reco_predictor = recognition_predictor("crnn_vgg16_bn", pretrained=True, batch_size=128)

predictor = KIEPredictor(
Expand All @@ -197,16 +207,12 @@ def test_trained_kie_predictor(mock_tilted_payslip):

assert isinstance(out, KIEDocument)
assert out.pages[0].predictions[CLASS_NAME][0].value == "Mr."
geometry_mr = np.array(
[[0.08844472, 0.35763523], [0.11625107, 0.34320644], [0.12588427, 0.35771032], [0.09807791, 0.37213911]]
)
assert np.allclose(np.array(out.pages[0].predictions[CLASS_NAME][0].geometry), geometry_mr)
geometry_mr = np.array([[0.1083984375, 0.0634765625], [0.1494140625, 0.0859375]])
assert np.allclose(np.array(out.pages[0].predictions[CLASS_NAME][0].geometry), geometry_mr, rtol=0.05)

assert out.pages[0].predictions[CLASS_NAME][-1].value == "Kabir)"
geometry_revised = np.array(
[[0.43725992, 0.67232439], [0.49045468, 0.64472149], [0.50570724, 0.66768597], [0.452512473, 0.69528887]]
)
assert np.allclose(np.array(out.pages[0].predictions[CLASS_NAME][-1].geometry), geometry_revised)
assert out.pages[0].predictions[CLASS_NAME][3].value == "revised"
geometry_revised = np.array([[0.7548828125, 0.126953125], [0.8388671875, 0.1484375]])
assert np.allclose(np.array(out.pages[0].predictions[CLASS_NAME][3].geometry), geometry_revised, rtol=0.05)

det_predictor = detection_predictor(
"db_resnet50",
Expand Down

0 comments on commit 28f4981

Please sign in to comment.