summaryrefslogtreecommitdiff
path: root/src/text_recognizer/character_predictor.py
diff options
context:
space:
mode:
authorGustaf Rydholm <gustaf.rydholm@gmail.com>2021-03-20 18:09:06 +0100
committerGustaf Rydholm <gustaf.rydholm@gmail.com>2021-03-20 18:09:06 +0100
commit7e8e54e84c63171e748bbf09516fd517e6821ace (patch)
tree996093f75a5d488dddf7ea1f159ed343a561ef89 /src/text_recognizer/character_predictor.py
parentb0719d84138b6bbe5f04a4982dfca673aea1a368 (diff)
Inital commit for refactoring to lightning
Diffstat (limited to 'src/text_recognizer/character_predictor.py')
-rw-r--r--src/text_recognizer/character_predictor.py29
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)