From dc28cbe2b4ed77be92ee8b2b69a20689c3bf02a4 Mon Sep 17 00:00:00 2001
From: aktersnurra <gustaf.rydholm@gmail.com>
Date: Sun, 8 Nov 2020 14:54:44 +0100
Subject: new updates

---
 src/text_recognizer/line_predictor.py | 28 ++++++++++++++++++++++++++++
 1 file changed, 28 insertions(+)
 create mode 100644 src/text_recognizer/line_predictor.py

(limited to 'src/text_recognizer/line_predictor.py')

diff --git a/src/text_recognizer/line_predictor.py b/src/text_recognizer/line_predictor.py
new file mode 100644
index 0000000..981e2c9
--- /dev/null
+++ b/src/text_recognizer/line_predictor.py
@@ -0,0 +1,28 @@
+"""LinePredictor class."""
+import importlib
+from typing import Tuple, Union
+
+import numpy as np
+from torch import nn
+
+from text_recognizer import datasets, networks
+from text_recognizer.models import VisionTransformerModel
+from text_recognizer.util import read_image
+
+
+class LinePredictor:
+    """Given an image of a line of handwritten text, recognizes the text content."""
+
+    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.eval()
+
+    def predict(self, image_or_filename: Union[np.ndarray, str]) -> Tuple[str, float]:
+        """Predict on a single images contianing a handwritten character."""
+        if isinstance(image_or_filename, str):
+            image = read_image(image_or_filename, grayscale=True)
+        else:
+            image = image_or_filename
+        return self.model.predict_on_image(image)
-- 
cgit v1.2.3-70-g09d2