From 527bb98b191d82b308de1585047e06056258d08d Mon Sep 17 00:00:00 2001 From: aktersnurra Date: Wed, 18 Nov 2020 20:56:19 +0100 Subject: Some minor changes. --- src/text_recognizer/line_predictor.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) (limited to 'src/text_recognizer/line_predictor.py') diff --git a/src/text_recognizer/line_predictor.py b/src/text_recognizer/line_predictor.py index 981e2c9..8e348fe 100644 --- a/src/text_recognizer/line_predictor.py +++ b/src/text_recognizer/line_predictor.py @@ -6,7 +6,7 @@ import numpy as np from torch import nn from text_recognizer import datasets, networks -from text_recognizer.models import VisionTransformerModel +from text_recognizer.models import TransformerModel from text_recognizer.util import read_image @@ -16,7 +16,7 @@ class LinePredictor: def __init__(self, dataset: str, network_fn: str) -> None: network_fn = getattr(networks, network_fn) dataset = getattr(datasets, dataset) - self.model = VisionTransformerModel(network_fn=network_fn, dataset=dataset) + self.model = TransformerModel(network_fn=network_fn, dataset=dataset) self.model.eval() def predict(self, image_or_filename: Union[np.ndarray, str]) -> Tuple[str, float]: -- cgit v1.2.3-70-g09d2