summaryrefslogtreecommitdiff
path: root/src/text_recognizer/character_predictor.py
diff options
context:
space:
mode:
Diffstat (limited to 'src/text_recognizer/character_predictor.py')
-rw-r--r--src/text_recognizer/character_predictor.py7
1 files changed, 5 insertions, 2 deletions
diff --git a/src/text_recognizer/character_predictor.py b/src/text_recognizer/character_predictor.py
index df37e68..ad71289 100644
--- a/src/text_recognizer/character_predictor.py
+++ b/src/text_recognizer/character_predictor.py
@@ -4,6 +4,7 @@ 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
@@ -11,9 +12,11 @@ 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:
+ def __init__(self, network_fn: str, dataset: str) -> None:
"""Intializes the CharacterModel and load the pretrained weights."""
- self.model = CharacterModel(network_fn=network_fn)
+ 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()