From 7e8e54e84c63171e748bbf09516fd517e6821ace Mon Sep 17 00:00:00 2001 From: Gustaf Rydholm Date: Sat, 20 Mar 2021 18:09:06 +0100 Subject: Inital commit for refactoring to lightning --- text_recognizer/character_predictor.py | 29 +++++++++++++++++++++++++++++ 1 file changed, 29 insertions(+) create mode 100644 text_recognizer/character_predictor.py (limited to 'text_recognizer/character_predictor.py') diff --git a/text_recognizer/character_predictor.py b/text_recognizer/character_predictor.py new file mode 100644 index 0000000..ad71289 --- /dev/null +++ b/text_recognizer/character_predictor.py @@ -0,0 +1,29 @@ +"""CharacterPredictor class.""" +from typing import Dict, Tuple, Type, Union + +import numpy as np +from torch import nn + +from text_recognizer import datasets, networks +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: str, dataset: str) -> None: + """Intializes the CharacterModel and load the pretrained weights.""" + network_fn = getattr(networks, network_fn) + dataset = getattr(datasets, dataset) + self.model = CharacterModel(network_fn=network_fn, dataset=dataset) + self.model.eval() + self.model.use_swa_model() + + 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) -- cgit v1.2.3-70-g09d2