diff options
Diffstat (limited to 'text_recognizer/tests/test_character_predictor.py')
-rw-r--r-- | text_recognizer/tests/test_character_predictor.py | 31 |
1 files changed, 31 insertions, 0 deletions
diff --git a/text_recognizer/tests/test_character_predictor.py b/text_recognizer/tests/test_character_predictor.py new file mode 100644 index 0000000..01bda78 --- /dev/null +++ b/text_recognizer/tests/test_character_predictor.py @@ -0,0 +1,31 @@ +"""Test for CharacterPredictor class.""" +import importlib +import os +from pathlib import Path +import unittest + +from loguru import logger + +from text_recognizer.character_predictor import CharacterPredictor +from text_recognizer.networks import MLP + +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.""" + network_fn_ = MLP + predictor = CharacterPredictor(network_fn=network_fn_) + + for filename in SUPPORT_DIRNAME.glob("*.png"): + pred, conf = predictor.predict(str(filename)) + logger.info( + f"Prediction: {pred} at confidence: {conf} for image with character {filename.stem}" + ) + self.assertEqual(pred, filename.stem) + self.assertGreater(conf, 0.7) |