diff options
author | aktersnurra <gustaf.rydholm@gmail.com> | 2020-07-22 23:18:08 +0200 |
---|---|---|
committer | aktersnurra <gustaf.rydholm@gmail.com> | 2020-07-22 23:18:08 +0200 |
commit | f473456c19558aaf8552df97a51d4e18cc69dfa8 (patch) | |
tree | 0d35ce2410ff623ba5fb433d616d95b67ecf7a98 /src/text_recognizer/tests | |
parent | ad3bd52530f4800d4fb05dfef3354921f95513af (diff) |
Working training loop and testing of trained CharacterModel.
Diffstat (limited to 'src/text_recognizer/tests')
-rw-r--r-- | src/text_recognizer/tests/test_character_predictor.py | 19 |
1 files changed, 17 insertions, 2 deletions
diff --git a/src/text_recognizer/tests/test_character_predictor.py b/src/text_recognizer/tests/test_character_predictor.py index 7c094ef..c603a3a 100644 --- a/src/text_recognizer/tests/test_character_predictor.py +++ b/src/text_recognizer/tests/test_character_predictor.py @@ -1,9 +1,14 @@ """Test for CharacterPredictor class.""" +import importlib import os from pathlib import Path import unittest +import click +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" @@ -13,13 +18,23 @@ os.environ["CUDA_VISIBLE_DEVICES"] = "" class TestCharacterPredictor(unittest.TestCase): """Tests for the CharacterPredictor class.""" + # @click.command() + # @click.option( + # "--network", type=str, help="Network to load, e.g. MLP or LeNet.", default="MLP" + # ) def test_filename(self) -> None: """Test that CharacterPredictor correctly predicts on a single image, for serveral test images.""" - predictor = CharacterPredictor() + network_module = importlib.import_module("text_recognizer.networks") + network_fn_ = getattr(network_module, "MLP") + # network_args = {"input_size": [28, 28], "output_size": 62, "dropout_rate": 0} + network_args = {"input_size": 784, "output_size": 62, "dropout_rate": 0.2} + predictor = CharacterPredictor( + network_fn=network_fn_, network_args=network_args + ) for filename in SUPPORT_DIRNAME.glob("*.png"): pred, conf = predictor.predict(str(filename)) - print( + logger.info( f"Prediction: {pred} at confidence: {conf} for image with character {filename.stem}" ) self.assertEqual(pred, filename.stem) |