From f473456c19558aaf8552df97a51d4e18cc69dfa8 Mon Sep 17 00:00:00 2001 From: aktersnurra Date: Wed, 22 Jul 2020 23:18:08 +0200 Subject: Working training loop and testing of trained CharacterModel. --- src/text_recognizer/character_predictor.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 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 69ef896..a773f36 100644 --- a/src/text_recognizer/character_predictor.py +++ b/src/text_recognizer/character_predictor.py @@ -1,8 +1,8 @@ """CharacterPredictor class.""" - -from typing import Tuple, Union +from typing import Dict, Tuple, Type, Union import numpy as np +from torch import nn from text_recognizer.models import CharacterModel from text_recognizer.util import read_image @@ -11,9 +11,9 @@ from text_recognizer.util import read_image class CharacterPredictor: """Recognizes the character in handwritten character images.""" - def __init__(self) -> None: + def __init__(self, network_fn: Type[nn.Module], network_args: Dict) -> None: """Intializes the CharacterModel and load the pretrained weights.""" - self.model = CharacterModel() + self.model = CharacterModel(network_fn=network_fn, network_args=network_args) self.model.load_weights() self.model.eval() -- cgit v1.2.3-70-g09d2