summaryrefslogtreecommitdiff
path: root/text_recognizer/data/utils/image_utils.py
diff options
context:
space:
mode:
authorGustaf Rydholm <gustaf.rydholm@gmail.com>2022-02-11 21:57:40 +0100
committerGustaf Rydholm <gustaf.rydholm@gmail.com>2022-02-11 21:57:40 +0100
commit7d8a06f2ddebf98b724fd5940b02c5b440837601 (patch)
tree108c913b52a532cfdfd19aba4dc9a43425239255 /text_recognizer/data/utils/image_utils.py
parent343d30353b317328592ee8479d40e16f4fb2439b (diff)
fix: move image utils back into data folder
Diffstat (limited to 'text_recognizer/data/utils/image_utils.py')
-rw-r--r--text_recognizer/data/utils/image_utils.py34
1 files changed, 34 insertions, 0 deletions
diff --git a/text_recognizer/data/utils/image_utils.py b/text_recognizer/data/utils/image_utils.py
new file mode 100644
index 0000000..d374fc3
--- /dev/null
+++ b/text_recognizer/data/utils/image_utils.py
@@ -0,0 +1,34 @@
+"""Image util functions for loading and saving images."""
+import base64
+from io import BytesIO
+from pathlib import Path
+from typing import Any, Union
+
+from PIL import Image
+import smart_open
+
+
+def read_image_pil(image_uri: Union[Path, str], grayscale: bool = False) -> Image:
+ """Read image from uri."""
+ with smart_open.open(image_uri, "rb") as image_file:
+ return read_image_pil_file(image_file, grayscale)
+
+
+def read_image_pil_file(image_file: Any, grayscale: bool = False) -> Image:
+ """Return PIL image."""
+ with Image.open(image_file) as image:
+ if grayscale:
+ image = image.convert("L")
+ else:
+ image = image.convert(mode=image.mode)
+ return image
+
+
+def read_b64_image(b64_str: str, grayscale: bool = False) -> Image:
+ """Load base64-encoded images."""
+ try:
+ _, b64_data = b64_str.split(",")
+ image_file = BytesIO(base64.b64decode(b64_data))
+ return read_image_pil_file(image_file, grayscale)
+ except Exception as e:
+ raise ValueError(f"Could not load image from b64 {b64_str}: {e}")