From dc28cbe2b4ed77be92ee8b2b69a20689c3bf02a4 Mon Sep 17 00:00:00 2001 From: aktersnurra Date: Sun, 8 Nov 2020 14:54:44 +0100 Subject: new updates --- src/text_recognizer/character_predictor.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) (limited to 'src/text_recognizer/character_predictor.py') 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() -- cgit v1.2.3-70-g09d2