diff options
Diffstat (limited to 'src/text_recognizer/tests/test_character_predictor.py')
-rw-r--r-- | src/text_recognizer/tests/test_character_predictor.py | 26 |
1 files changed, 26 insertions, 0 deletions
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) |