summaryrefslogtreecommitdiff
path: root/src/text_recognizer/tests/test_character_predictor.py
blob: 7c094efe6527e24ae797a7af6031d176afa29995 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
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)