add functional test for complete prediction with trained predictor

doctr/utils/geometry.py

@@ -150,13 +150,6 @@ def remap_boxes(
     orig_height, orig_width = orig_shape
     dest_height, dest_width = dest_shape
     mboxes = loc_preds.copy()
-    # # remaps position of the box center for the destination image shape
-    # mboxes[:, 0] = ((loc_preds[:, 0] * orig_width) + (dest_width - orig_width) / 2) / dest_width
-    # mboxes[:, 1] = ((loc_preds[:, 1] * orig_height) + (dest_height - orig_height) / 2) / dest_height
-    # # remaps box dimension for the destination image shape
-    # mboxes[:, 2] = loc_preds[:, 2] * orig_width / dest_width
-    # mboxes[:, 3] = loc_preds[:, 3] * orig_height / dest_height
-
     mboxes[:, :, 0] = ((loc_preds[:, :, 0] * orig_width) + (dest_width - orig_width) / 2) / dest_width
     mboxes[:, :, 1] = ((loc_preds[:, :, 1] * orig_height) + (dest_height - orig_height) / 2) / dest_height
 

tests/common/test_models.py

@@ -83,19 +83,6 @@ def mock_image(tmpdir_factory):
     return image
 
 
-@pytest.fixture(scope="session")
-def mock_tilted_payslip(tmpdir_factory):
-    url = 'https://3.bp.blogspot.com/-Es0oHTCrVEk/UnYA-iW9rYI/AAAAAAAAAFI/hWExrXFbo9U/s1600/003.jpg'
-    file = BytesIO(requests.get(url).content)
-    folder = tmpdir_factory.mktemp("images")
-    fn = str(folder.join("mock_payslip.jpeg"))
-    with open(fn, 'wb') as f:
-        f.write(file.getbuffer())
-    image = reader.read_img_as_numpy(fn)
-    image = geometry.rotate_image(image, 30, expand=True)
-    return image
-
-
 @pytest.fixture(scope="function")
 def mock_bitmap(mock_image):
     bitmap = np.squeeze(cv2.cvtColor(mock_image, cv2.COLOR_BGR2GRAY) / 255.)
@@ -107,7 +94,7 @@ def test_get_bitmap_angle(mock_bitmap):
     assert abs(angle - 30.) < 1.
 
 
-def test_estimate_orientation(mock_image):
+def test_estimate_orientation(mock_image, mock_tilted_payslip):
     assert estimate_orientation(mock_image * 0) == 0
 
     angle = estimate_orientation(mock_image)
@@ -117,10 +104,9 @@ def test_estimate_orientation(mock_image):
     angle_rotated = estimate_orientation(rotated)
     assert abs(angle_rotated) < 1.
 
-
-def test_estimate_orientation(mock_tilted_payslip):
+    mock_tilted_payslip = reader.read_img_as_numpy(mock_tilted_payslip)
     assert (estimate_orientation(mock_tilted_payslip) - 30.) < 1.
 
     rotated = geometry.rotate_image(mock_tilted_payslip, -30, expand=True)
     angle_rotated = estimate_orientation(rotated)
-    assert abs(angle_rotated) < 1.
+    assert abs(angle_rotated) < 1.

tests/conftest.py

@@ -3,13 +3,17 @@
 import tempfile
 from io import BytesIO
 
+import cv2
 import fitz
 import hdf5storage
 import numpy as np
 import pytest
 import requests
 import scipy.io as sio
 
+from doctr.io import reader
+from doctr.utils import geometry
+
 
 @pytest.fixture(scope="session")
 def mock_vocab():
@@ -35,6 +39,26 @@ def mock_pdf(tmpdir_factory):
     return str(fn)
 
 
+@pytest.fixture(scope="session")
+def mock_payslip(tmpdir_factory):
+    url = 'https://3.bp.blogspot.com/-Es0oHTCrVEk/UnYA-iW9rYI/AAAAAAAAAFI/hWExrXFbo9U/s1600/003.jpg'
+    file = BytesIO(requests.get(url).content)
+    folder = tmpdir_factory.mktemp("data")
+    fn = str(folder.join("mock_payslip.jpeg"))
+    with open(fn, 'wb') as f:
+        f.write(file.getbuffer())
+    return fn
+
+
+@pytest.fixture(scope="session")
+def mock_tilted_payslip(mock_payslip, tmpdir_factory):
+    image = reader.read_img_as_numpy(mock_payslip)
+    image = geometry.rotate_image(image, 30, expand=True)
+    tmp_path = str(tmpdir_factory.mktemp("data").join("mock_tilted_payslip.jpg"))
+    cv2.imwrite(tmp_path, image)
+    return tmp_path
+
+
 @pytest.fixture(scope="session")
 def mock_text_box_stream():
     url = 'https://www.pngitem.com/pimgs/m/357-3579845_love-neon-loveislove-word-text-typography-freetoedit-picsart.png'

tests/tensorflow/test_models_zoo_tf.py

@@ -5,9 +5,11 @@
 from doctr.io import Document, DocumentFile
 from doctr.models import detection, recognition
 from doctr.models.detection.predictor import DetectionPredictor
+from doctr.models.detection.zoo import detection_predictor
 from doctr.models.predictor import OCRPredictor
 from doctr.models.preprocessor import PreProcessor
 from doctr.models.recognition.predictor import RecognitionPredictor
+from doctr.models.recognition.zoo import recognition_predictor
 
 
 @pytest.mark.parametrize(
@@ -54,6 +56,36 @@ def test_ocrpredictor(mock_pdf, mock_vocab, assume_straight_pages, straighten_pa
         _ = predictor([input_page])
 
 
+def test_trained_ocr_predictor(mock_tilted_payslip):
+    doc = DocumentFile.from_images(mock_tilted_payslip)
+
+    det_predictor = detection_predictor('db_resnet50', pretrained=True, batch_size=2, assume_straight_pages=True)
+    reco_predictor = recognition_predictor('crnn_vgg16_bn', pretrained=True, batch_size=128)
+
+    predictor = OCRPredictor(
+        det_predictor,
+        reco_predictor,
+        assume_straight_pages=True,
+        straighten_pages=True,
+    )
+
+    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)
+
+    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)
+
+
 @pytest.mark.parametrize(
     "det_arch, reco_arch",
     [