summaryrefslogtreecommitdiff
path: root/src/text_recognizer/character_predictor.py
blob: b733a53494174ac91840259f8e3ac8c70423ea3c (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
"""CharacterPredictor class."""
from typing import Dict, Tuple, Type, Union

import numpy as np
from torch import nn

from text_recognizer.models import CharacterModel
from text_recognizer.util import read_image


class CharacterPredictor:
    """Recognizes the character in handwritten character images."""

    def __init__(self, network_fn: Type[nn.Module]) -> None:
        """Intializes the CharacterModel and load the pretrained weights."""
        self.model = CharacterModel(network_fn=network_fn)
        self.model.eval()

    def predict(self, image_or_filename: Union[np.ndarray, str]) -> Tuple[str, float]:
        """Predict on a single images contianing a handwritten character."""
        if isinstance(image_or_filename, str):
            image = read_image(image_or_filename, grayscale=True)
        else:
            image = image_or_filename
        return self.model.predict_on_image(image)