data_utils.py

fromioimportBytesIO
importbase64
from .pngimportWriter, from_array

try:
fromPILimportImage

pil_imported=True
exceptImportError:
pil_imported=False


defimage_array_to_data_uri(img, backend="pil", compression=4, ext="png"):
"""Converts a numpy array of uint8 into a base64 png or jpg string.

 Parameters
 ----------
 img: ndarray of uint8
 array image
 backend: str
 'auto', 'pil' or 'pypng'. If 'auto', Pillow is used if installed,
 otherwise pypng.
 compression: int, between 0 and 9
 compression level to be passed to the backend
 ext: str, 'png' or 'jpg'
 compression format used to generate b64 string
 """
# PIL and pypng error messages are quite obscure so we catch invalid compression values
ifcompression<0orcompression>9:
raiseValueError("compression level must be between 0 and 9.")
alpha=False
ifimg.ndim==2:
mode="L"
elifimg.ndim==3andimg.shape[-1] ==3:
mode="RGB"
elifimg.ndim==3andimg.shape[-1] ==4:
mode="RGBA"
alpha=True
else:
raiseValueError("Invalid image shape")
ifbackend=="auto":
backend="pil"ifpil_importedelse"pypng"
ifext!="png"andbackend!="pil":
raiseValueError("jpg binary strings are only available with PIL backend")

ifbackend=="pypng":
ndim=img.ndim
sh=img.shape
ifndim==3:
img=img.reshape((sh[0], sh[1] *sh[2]))
w=Writer(
sh[1], sh[0], greyscale=(ndim==2), alpha=alpha, compression=compression
 )
img_png=from_array(img, mode=mode)
prefix="data:image/png;base64,"
withBytesIO() asstream:
w.write(stream, img_png.rows)
base64_string=prefix+base64.b64encode(stream.getvalue()).decode("utf-8")
else: # pil
ifnotpil_imported:
raiseImportError(
"pillow needs to be installed to use `backend='pil'. Please"
"install pillow or use `backend='pypng'."
 )
pil_img=Image.fromarray(img)
ifext=="jpg"orext=="jpeg":
prefix="data:image/jpeg;base64,"
ext="jpeg"
else:
prefix="data:image/png;base64,"
ext="png"
withBytesIO() asstream:
pil_img.save(stream, format=ext, compress_level=compression)
base64_string=prefix+base64.b64encode(stream.getvalue()).decode("utf-8")
returnbase64_string