pkg/util.py

import io
import numpy as np
import numpy.typing as npt
from PIL import Image
from PIL.Image import Image as pilImage
import base64
import re
from diffusers.utils import PIL_INTERPOLATION

def img_binary_data_to_pil(imgBinary: bytes) -> pilImage:
    image_data = re.sub('^data:image/.+;base64,', '', imgBinary.decode())
    decoded = base64.b64decode(image_data)
    img = Image.open(io.BytesIO(decoded))
    return img

def resizePilToMaxSide(img: pilImage, maxSideLength: int) -> pilImage:
    w, h = img.size
    maxSide = max(w, h, maxSideLength)
    scale = maxSideLength / maxSide
    targetSize = (int(scale * w), int(scale * h))
    img = img.resize(targetSize, resample=PIL_INTERPOLATION["lanczos"])
    return img

def pil_to_base64(img: pilImage | npt.NDArray[np.uint8]) -> str:
    if type(img) is pilImage:
        pass
    elif type(img) is np.ndarray:
        img = pilImage.fromarray(img.astype('uint8'))
    else:
        raise TypeError
    
    # https://stackoverflow.com/a/59583262/13797085
    file_object = io.BytesIO()
    img.save(file_object, 'PNG')
    base64img = "data:image/png;base64," \
        + base64.b64encode(file_object.getvalue()).decode('ascii')
    return base64img