diff options
| author | aktersnurra <gustaf.rydholm@gmail.com> | 2020-11-18 20:56:19 +0100 |
|---|---|---|
| committer | aktersnurra <gustaf.rydholm@gmail.com> | 2020-11-18 20:56:19 +0100 |
| commit | 527bb98b191d82b308de1585047e06056258d08d (patch) | |
| tree | f33145dba398825871da3184a2735f6fb0b07268 /src/text_recognizer/line_predictor.py | |
| parent | f2cd16f340aa11afadb8fa90c29f85ca1b75a600 (diff) | |
Some minor changes.
Diffstat (limited to 'src/text_recognizer/line_predictor.py')
| -rw-r--r-- | src/text_recognizer/line_predictor.py | 4 |
1 files changed, 2 insertions, 2 deletions
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]: |