diff options
author | Gustaf Rydholm <gustaf.rydholm@gmail.com> | 2022-02-11 21:57:40 +0100 |
---|---|---|
committer | Gustaf Rydholm <gustaf.rydholm@gmail.com> | 2022-02-11 21:57:40 +0100 |
commit | 7d8a06f2ddebf98b724fd5940b02c5b440837601 (patch) | |
tree | 108c913b52a532cfdfd19aba4dc9a43425239255 /text_recognizer/data/utils | |
parent | 343d30353b317328592ee8479d40e16f4fb2439b (diff) |
fix: move image utils back into data folder
Diffstat (limited to 'text_recognizer/data/utils')
-rw-r--r-- | text_recognizer/data/utils/image_utils.py | 34 |
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}") |