From 9b8e14d89f0ef2508ed11f994f73af624155fe1d Mon Sep 17 00:00:00 2001 From: Gustaf Rydholm Date: Tue, 27 Sep 2022 01:44:49 +0200 Subject: Update data modules --- text_recognizer/data/stems/image.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) create mode 100644 text_recognizer/data/stems/image.py (limited to 'text_recognizer/data/stems/image.py') diff --git a/text_recognizer/data/stems/image.py b/text_recognizer/data/stems/image.py new file mode 100644 index 0000000..f04b3a0 --- /dev/null +++ b/text_recognizer/data/stems/image.py @@ -0,0 +1,18 @@ +from PIL import Image +import torch +from torch import Tensor +import torchvision.transforms as T + + +class ImageStem: + def __init__(self) -> None: + self.pil_transform = T.Compose([]) + self.pil_to_tensor = T.ToTensor() + self.torch_transform = torch.nn.Sequential() + + def __call__(self, img: Image) -> Tensor: + img = self.pil_transform(img) + img = self.pil_to_tensor(img) + with torch.no_grad(): + img = self.torch_transform(img) + return img -- cgit v1.2.3-70-g09d2