image_utils.py

from typing import Optional, Tuple

import gradio as gr
import numpy as np
from PIL import Image, ImageChops, ImageOps

NEGATIVE_DEFAULT = (
    "blurry, ugly, bad anatomy, deformed, distorted, extra fingers, "
    "extra limbs, poorly drawn hands, poorly drawn face, low quality"
)

def round_to_multiple_of_8(x: int) -> int:
    x = int(x)
    return max(64, x - (x % 8))

def resize_keep_aspect(img: Image.Image, max_side: int, resample) -> Image.Image:
    w, h = img.size
    longest = max(w, h)
    if longest <= max_side:
        return img

    scale = max_side / float(longest)
    new_w = round_to_multiple_of_8(int(w * scale))
    new_h = round_to_multiple_of_8(int(h * scale))
    return img.resize((new_w, new_h), resample)

def merge_layer_alphas(layers) -> Optional[Image.Image]:
    if not layers:
        return None

    merged = None
    for layer in layers:
        if layer is None:
            continue

        if not isinstance(layer, Image.Image):
            layer = Image.fromarray(np.array(layer))

        rgba = layer.convert("RGBA")
        alpha = rgba.getchannel("A")

        if merged is None:
            merged = alpha
        else:
            merged = ImageChops.lighter(merged, alpha)

    return merged

def threshold_mask(mask_l: Image.Image, threshold: int = 10) -> Image.Image:
    arr = np.array(mask_l, dtype=np.uint8)
    arr = np.where(arr > threshold, 255, 0).astype(np.uint8)
    return Image.fromarray(arr, mode="L")

def extract_image_and_mask(editor_value) -> Tuple[Image.Image, Image.Image]:
    if editor_value is None:
        raise gr.Error("Upload an image first.")

    bg = editor_value.get("background", None)
    layers = editor_value.get("layers", None)

    if bg is None:
        raise gr.Error("Upload an image first.")

    if not isinstance(bg, Image.Image):
        bg = Image.fromarray(np.array(bg))

    bg = ImageOps.exif_transpose(bg).convert("RGB")

    mask = merge_layer_alphas(layers)
    if mask is None:
        raise gr.Error("Paint a mask first.")

    if mask.size != bg.size:
        mask = mask.resize(bg.size, Image.NEAREST)

    mask = threshold_mask(mask.convert("L"))

    if mask.getbbox() is None:
        raise gr.Error("Paint a mask first.")

    return bg, mask

def prepare_inputs_for_cpu(
    image: Image.Image,
    mask: Image.Image,
    max_side: int,
) -> Tuple[Image.Image, Image.Image]:
    image = resize_keep_aspect(image, max_side=max_side, resample=Image.LANCZOS)
    mask = mask.resize(image.size, Image.NEAREST)

    w, h = image.size
    w = round_to_multiple_of_8(w)
    h = round_to_multiple_of_8(h)

    if image.size != (w, h):
        image = image.resize((w, h), Image.LANCZOS)

    if mask.size != (w, h):
        mask = mask.resize((w, h), Image.NEAREST)

    return image, mask