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 /src/text_recognizer/character_predictor.py | |
parent | b0719d84138b6bbe5f04a4982dfca673aea1a368 (diff) |
Inital commit for refactoring to lightning
Diffstat (limited to 'src/text_recognizer/character_predictor.py')
-rw-r--r-- | src/text_recognizer/character_predictor.py | 29 |
1 files changed, 0 insertions, 29 deletions
diff --git a/src/text_recognizer/character_predictor.py b/src/text_recognizer/character_predictor.py deleted file mode 100644 index ad71289..0000000 --- a/src/text_recognizer/character_predictor.py +++ /dev/null @@ -1,29 +0,0 @@ -"""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) |