blob: 7ec7307297d73b414e6c82545edc5f6dc8177606 (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
|
"""Image util functions for loading and saving images."""
import base64
from io import BytesIO
from pathlib import Path
from typing import Union
import smart_open
from PIL import Image
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: str, 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}")
|
