From 8103ab92b9d8f1898164ded197d7ec60671e2aca Mon Sep 17 00:00:00 2001
From: Gustaf Rydholm <gustaf.rydholm@gmail.com>
Date: Sun, 6 Feb 2022 22:46:01 +0100
Subject: fix: refactor and move image utils
---
text_recognizer/data/utils/image_utils.py | 49 -------------------------------
text_recognizer/image_utils.py | 34 +++++++++++++++++++++
2 files changed, 34 insertions(+), 49 deletions(-)
delete mode 100644 text_recognizer/data/utils/image_utils.py
create mode 100644 text_recognizer/image_utils.py
(limited to 'text_recognizer')
diff --git a/text_recognizer/data/utils/image_utils.py b/text_recognizer/data/utils/image_utils.py
deleted file mode 100644
index c2b8915..0000000
--- a/text_recognizer/data/utils/image_utils.py
+++ /dev/null
@@ -1,49 +0,0 @@
-"""Image util functions for loading and saving images."""
-from pathlib import Path
-from typing import Union
-from urllib.request import urlopen
-
-import cv2
-import numpy as np
-from PIL import Image
-
-
-def read_image_pil(image_uri: Union[Path, str], grayscale: bool = False) -> Image:
- """Return PIL image."""
- image = Image.open(image_uri)
- if grayscale:
- image = image.convert("L")
- return image
-
-
-def read_image(image_uri: Union[Path, str], grayscale: bool = False) -> np.array:
- """Read image_uri."""
-
- if isinstance(image_uri, str):
- image_uri = Path(image_uri)
-
- def read_image_from_filename(image_filename: Path, imread_flag: int) -> np.array:
- return cv2.imread(str(image_filename), imread_flag)
-
- def read_image_from_url(image_url: Path, imread_flag: int) -> np.array:
- url_response = urlopen(str(image_url)) # nosec
- image_array = np.array(bytearray(url_response.read()), dtype=np.uint8)
- return cv2.imdecode(image_array, imread_flag)
-
- imread_flag = cv2.IMREAD_GRAYSCALE if grayscale else cv2.IMREAD_COLOR
- image = None
-
- if image_uri.exists():
- image = read_image_from_filename(image_uri, imread_flag)
- else:
- image = read_image_from_url(image_uri, imread_flag)
-
- if image is None:
- raise ValueError(f"Could not load image at {image_uri}")
-
- return image
-
-
-def write_image(image: np.ndarray, filename: Union[Path, str]) -> None:
- """Write image to file."""
- cv2.imwrite(str(filename), image)
diff --git a/text_recognizer/image_utils.py b/text_recognizer/image_utils.py
new file mode 100644
index 0000000..0faed79
--- /dev/null
+++ b/text_recognizer/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(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}")
--
cgit v1.2.3-70-g09d2