latest

7 files changed, 62 insertions, 0 deletions

diff --git a/src/text_recognizer/tests/__init__.py b/src/text_recognizer/tests/__init__.py
new file mode 100644
index 0000000..18ff212
--- /dev/null
+++ b/src/text_recognizer/tests/__init__.py
@@ -0,0 +1 @@
+"""Test modules for the text text recognizer."""
diff --git a/src/text_recognizer/tests/support/__init__.py b/src/text_recognizer/tests/support/__init__.py
new file mode 100644
index 0000000..a265ede
--- /dev/null
+++ b/src/text_recognizer/tests/support/__init__.py
@@ -0,0 +1,2 @@
+"""Support file modules."""
+from .create_emnist_support_files import create_emnist_support_files
diff --git a/src/text_recognizer/tests/support/create_emnist_support_files.py b/src/text_recognizer/tests/support/create_emnist_support_files.py
new file mode 100644
index 0000000..5dd1a81
--- /dev/null
+++ b/src/text_recognizer/tests/support/create_emnist_support_files.py
@@ -0,0 +1,33 @@
+"""Module for creating EMNIST test support files."""
+from pathlib import Path
+import shutil
+
+from text_recognizer.datasets.emnist_dataset import (
+    fetch_emnist_dataset,
+    load_emnist_mapping,
+)
+from text_recognizer.util import write_image
+
+SUPPORT_DIRNAME = Path(__file__).parents[0].resolve() / "emnist"
+
+
+def create_emnist_support_files() -> None:
+    """Create support images for test of CharacterPredictor class."""
+    shutil.rmtree(SUPPORT_DIRNAME, ignore_errors=True)
+    SUPPORT_DIRNAME.mkdir()
+
+    dataset = fetch_emnist_dataset(split="byclass", train=False)
+    mapping = load_emnist_mapping()
+
+    for index in [5, 7, 9]:
+        image, label = dataset[index]
+        if len(image.shape) == 3:
+            image = image.squeeze(0)
+        image = image.numpy()
+        label = mapping[int(label)]
+        print(index, label)
+        write_image(image, str(SUPPORT_DIRNAME / f"{label}.png"))
+
+
+if __name__ == "__main__":
+    create_emnist_support_files()
diff --git a/src/text_recognizer/tests/support/emnist/8.png b/src/text_recognizer/tests/support/emnist/8.png
new file mode 100644
index 0000000..faa29aa
--- /dev/null
+++ b/src/text_recognizer/tests/support/emnist/8.png
diff --git a/src/text_recognizer/tests/support/emnist/U.png b/src/text_recognizer/tests/support/emnist/U.png
new file mode 100644
index 0000000..304eaec
--- /dev/null
+++ b/src/text_recognizer/tests/support/emnist/U.png
diff --git a/src/text_recognizer/tests/support/emnist/e.png b/src/text_recognizer/tests/support/emnist/e.png
new file mode 100644
index 0000000..a03ecd4
--- /dev/null
+++ b/src/text_recognizer/tests/support/emnist/e.png
diff --git a/src/text_recognizer/tests/test_character_predictor.py b/src/text_recognizer/tests/test_character_predictor.py
new file mode 100644
index 0000000..7c094ef
--- /dev/null
+++ b/src/text_recognizer/tests/test_character_predictor.py
@@ -0,0 +1,26 @@
+"""Test for CharacterPredictor class."""
+import os
+from pathlib import Path
+import unittest
+
+from text_recognizer.character_predictor import CharacterPredictor
+
+SUPPORT_DIRNAME = Path(__file__).parents[0].resolve() / "support" / "emnist"
+
+os.environ["CUDA_VISIBLE_DEVICES"] = ""
+
+
+class TestCharacterPredictor(unittest.TestCase):
+    """Tests for the CharacterPredictor class."""
+
+    def test_filename(self) -> None:
+        """Test that CharacterPredictor correctly predicts on a single image, for serveral test images."""
+        predictor = CharacterPredictor()
+
+        for filename in SUPPORT_DIRNAME.glob("*.png"):
+            pred, conf = predictor.predict(str(filename))
+            print(
+                f"Prediction: {pred} at confidence: {conf} for image with character {filename.stem}"
+            )
+            self.assertEqual(pred, filename.stem)
+            self.assertGreater(conf, 0.7)