You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

76 lines
1.6 KiB

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