diff options
| author | aktersnurra <grydholm@kth.se> | 2020-12-02 23:48:52 +0100 |
|---|---|---|
| committer | aktersnurra <grydholm@kth.se> | 2020-12-02 23:48:52 +0100 |
| commit | 5529e0fc9ca39e81fe0f08a54f257d32f0afe120 (patch) | |
| tree | f2be992554e278857db7d56786dba54a76d439c7 /src/text_recognizer/line_predictor.py | |
| parent | e3b039c9adb4bce42ede4cb682a3ae71e797539a (diff) | |
| parent | 8e3985c9cde6666e4314973312135ec1c7a025b9 (diff) | |
Merge branch 'master' of github.com:aktersnurra/text-recognizer
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]: |