diff options
author | Gustaf Rydholm <gustaf.rydholm@gmail.com> | 2021-03-20 18:09:06 +0100 |
---|---|---|
committer | Gustaf Rydholm <gustaf.rydholm@gmail.com> | 2021-03-20 18:09:06 +0100 |
commit | 7e8e54e84c63171e748bbf09516fd517e6821ace (patch) | |
tree | 996093f75a5d488dddf7ea1f159ed343a561ef89 /text_recognizer/tests/test_character_predictor.py | |
parent | b0719d84138b6bbe5f04a4982dfca673aea1a368 (diff) |
Inital commit for refactoring to lightning
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) |