76 lines
1.6 KiB
Python
76 lines
1.6 KiB
Python
import PIL
|
|
import cv2
|
|
import base64
|
|
import numpy as np
|
|
from PIL.Image import Image as ImageFile
|
|
from PIL.JpegImagePlugin import JpegImageFile
|
|
from PIL.PngImagePlugin import PngImageFile
|
|
from PIL import Image
|
|
from io import BytesIO
|
|
|
|
|
|
def _encode_numpy_array_image(image):
|
|
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
|
|
|
if image.shape[-1] == 3:
|
|
_, buffer = cv2.imencode(".jpg", image)
|
|
|
|
elif image.shape[-1] == 4:
|
|
_, buffer = cv2.imencode(".png", image)
|
|
|
|
image_as_text = base64.b64encode(buffer)
|
|
|
|
return image_as_text
|
|
|
|
|
|
def _encode_pil_image(image):
|
|
opencv_image = np.array(image)
|
|
image_as_text = _encode_numpy_array_image(image=opencv_image)
|
|
|
|
return image_as_text
|
|
|
|
|
|
def _encode_image_file(image):
|
|
pil_image = Image.open(image)
|
|
|
|
return _encode_pil_image(pil_image)
|
|
|
|
|
|
def encode(image):
|
|
|
|
if (
|
|
type(image) == np.ndarray
|
|
or type(image) == str
|
|
or isinstance(
|
|
image,
|
|
(
|
|
JpegImageFile,
|
|
PngImageFile,
|
|
ImageFile,
|
|
),
|
|
)
|
|
):
|
|
|
|
if type(image) == np.ndarray:
|
|
image_as_text = _encode_numpy_array_image(image)
|
|
|
|
elif type(image) == str:
|
|
image_as_text = _encode_image_file(image)
|
|
|
|
else:
|
|
image_as_text = _encode_pil_image(image)
|
|
|
|
return image_as_text.decode("ascii")
|
|
|
|
else:
|
|
raise Exception(
|
|
"expected numpy.array, PIL.Image or str, not: ", str(type(image))
|
|
)
|
|
|
|
|
|
def decode(jpg_as_text):
|
|
if jpg_as_text is None:
|
|
return None
|
|
pil_image = Image.open(BytesIO(base64.b64decode(jpg_as_text)))
|
|
return pil_image
|