Add tests for prediction code

tests/data/training_page/labels.json → tests/data/training/training_dataset/labels.json

 {
     "test": {
-        "tests/data/training_page/images/fb3edb59-3678-49f8-8e16-8e32e3b0f051.png": {
+        "tests/data/training/training_dataset/images/fb3edb59-3678-49f8-8e16-8e32e3b0f051.png": {
             "text": "Both her wrists bore"
         },
-        "tests/data/training_page/images/fe498de2-ece4-4fbe-8b53-edfce1b820f0.png": {
+        "tests/data/training/training_dataset/images/fe498de2-ece4-4fbe-8b53-edfce1b820f0.png": {
             "text": "SOME years ago a contemporary"
         }
     },
     "train": {
-        "tests/data/training_page/images/0a34e13a-4ab0-4a91-8d7c-b1d8fee32628.png": {
+        "tests/data/training/training_dataset/images/0a34e13a-4ab0-4a91-8d7c-b1d8fee32628.png": {
             "text": "The latter do not regard"
         },
-        "tests/data/training_page/images/0a70e14f-feda-4607-989c-36cf581ddff5.png": {
+        "tests/data/training/training_dataset/images/0a70e14f-feda-4607-989c-36cf581ddff5.png": {
             "text": "At the beginning of"
         }
     },
     "val": {
-        "tests/data/training_page/images/0a576062-303c-4893-a729-c09c92865d31.png": {
+        "tests/data/training/training_dataset/images/0a576062-303c-4893-a729-c09c92865d31.png": {
             "text": "One can remember with"
         },
-        "tests/data/training_page/images/0b2457c8-81f1-4600-84d9-f8bf2822a991.png": {
+        "tests/data/training/training_dataset/images/0b2457c8-81f1-4600-84d9-f8bf2822a991.png": {
             "text": "The play was no more"
         }
     }

tests/test_prediction.py

+# -*- coding: utf-8 -*-
+
+import pytest
+import torch
+
+from dan.predict.prediction import DAN
+from dan.utils import read_image
+
+
+@pytest.mark.parametrize(
+    "image_name, expected_prediction",
+    (
+        (
+            "0a56e8b3-95cd-4fa5-a17b-5b0ff9e6ea84.png",
+            {"text": ["ⓈBellisson ⒻGeorges Ⓑ91 ⓁP ⒸM ⓀCh ⓄPlombier ⓅPatron?12241"]},
+        ),
+        (
+            "0dfe8bcd-ed0b-453e-bf19-cc697012296e.png",
+            {"text": ["ⓈTemplié ⒻMarcelle Ⓑ93 ⓁS Ⓚch ⓄE dactylo Ⓟ18376"]},
+        ),
+        (
+            "2c242f5c-e979-43c4-b6f2-a6d4815b651d.png",
+            {"text": ["Ⓢd ⒻCharles Ⓑ11 ⓁP ⒸC ⓀF Ⓞd Ⓟ14 31"]},
+        ),
+        (
+            "ffdec445-7f14-4f5f-be44-68d0844d0df1.png",
+            {"text": ["ⓈNaudin ⒻMarie Ⓑ53 ⓁS Ⓒv ⓀBelle mère"]},
+        ),
+    ),
+)
+def test_predict(
+    image_name,
+    expected_prediction,
+    prediction_data_path,
+):
+    device = "cpu"
+
+    dan_model = DAN(device)
+    dan_model.load(
+        prediction_data_path / "popp_line_model.pt",
+        prediction_data_path / "parameters.yml",
+        prediction_data_path / "charset.pkl",
+        mode="eval",
+    )
+
+    image = read_image(prediction_data_path / "images" / image_name)
+    image = dan_model.preprocess(image)
+
+    input_tensor = torch.tensor(image).permute(2, 0, 1).unsqueeze(0)
+    input_tensor = input_tensor.to(device)
+    input_sizes = [image.shape[:2]]
+
+    prediction = dan_model.predict(input_tensor, input_sizes)
+
+    assert prediction == expected_prediction

tests/test_training.py

@@ -57,7 +57,7 @@ def test_train_and_test(
 
     # Check that the trained model is correct
     for model_name in [expected_best_model_name, expected_last_model_name]:
-        expected_model = torch.load(FIXTURES / "training" / model_name)
+        expected_model = torch.load(FIXTURES / "training" / "models" / model_name)
         trained_model = torch.load(
             Path(
                 "outputs",