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()