diff options
Diffstat (limited to 'text_recognizer/data/transforms/image.py')
-rw-r--r-- | text_recognizer/data/transforms/image.py | 18 |
1 files changed, 18 insertions, 0 deletions
diff --git a/text_recognizer/data/transforms/image.py b/text_recognizer/data/transforms/image.py new file mode 100644 index 0000000..f04b3a0 --- /dev/null +++ b/text_recognizer/data/transforms/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 |