internals/util/image.py

import io

import numpy as np
from PIL import Image


def to_bytes(image: Image.Image) -> bytes:
    with io.BytesIO() as output:
        image.save(output, format="JPEG")
        return output.getvalue()


def resize_image(image: Image.Image, dimension: int = 512) -> Image.Image:
    iw, ih = image.size
    if iw > ih:
        image = image.resize((dimension, int(dimension * ih / iw)))
    else:
        image = image.resize((int(dimension * iw / ih), dimension))
    return image


def HWC3(x):
    "x: numpy array"
    assert x.dtype == np.uint8
    if x.ndim == 2:
        x = x[:, :, None]
    assert x.ndim == 3
    H, W, C = x.shape
    assert C == 1 or C == 3 or C == 4
    if C == 3:
        return x
    if C == 1:
        return np.concatenate([x, x, x], axis=2)
    if C == 4:
        color = x[:, :, 0:3].astype(np.float32)
        alpha = x[:, :, 3:4].astype(np.float32) / 255.0
        y = color * alpha + 255.0 * (1.0 - alpha)
        y = y.clip(0, 255).astype(np.uint8)
        return y


def from_bytes(data: bytes) -> Image.Image:
    return Image.open(io.BytesIO(data))


def padd_image(image: Image.Image, to_width: int, to_height: int) -> Image.Image:
    iw, ih = image.size
    img = Image.new("RGBA", (to_width, to_height), (0, 0, 0, 0))
    img.paste(image, ((to_width - iw) // 2, (to_height - ih) // 2))
    return img


def alpha_to_white(img: Image.Image) -> Image.Image:
    if img.mode == "RGBA":
        data = img.getdata()

        new_data = []

        for item in data:
            if item[3] == 0:
                new_data.append((255, 255, 255, 255))
            else:
                new_data.append(item)

        img.putdata(new_data)

    img = img.convert("RGB")
    return img