diff options
Diffstat (limited to 'text_recognizer/line_predictor.py')
-rw-r--r-- | text_recognizer/line_predictor.py | 28 |
1 files changed, 0 insertions, 28 deletions
diff --git a/text_recognizer/line_predictor.py b/text_recognizer/line_predictor.py deleted file mode 100644 index 8e348fe..0000000 --- a/text_recognizer/line_predictor.py +++ /dev/null @@ -1,28 +0,0 @@ -"""LinePredictor class.""" -import importlib -from typing import Tuple, Union - -import numpy as np -from torch import nn - -from text_recognizer import datasets, networks -from text_recognizer.models import TransformerModel -from text_recognizer.util import read_image - - -class LinePredictor: - """Given an image of a line of handwritten text, recognizes the text content.""" - - def __init__(self, dataset: str, network_fn: str) -> None: - network_fn = getattr(networks, network_fn) - dataset = getattr(datasets, dataset) - self.model = TransformerModel(network_fn=network_fn, dataset=dataset) - 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) |