utils.py

"""
Utility functions for the Qwen Image Editor application.
"""

import io
import base64
import numpy as np
from PIL import Image
from typing import Optional, Tuple


def image_to_bytes(img: Image.Image, format: str = "PNG") -> bytes:
    """
    Convert PIL Image to bytes.
    
    Args:
        img: PIL Image object
        format: Output format (PNG, JPEG, etc.)
    
    Returns:
        Image as bytes
    """
    buffer = io.BytesIO()
    img.save(buffer, format=format)
    buffer.seek(0)
    return buffer.getvalue()


def bytes_to_image(data: bytes) -> Image.Image:
    """
    Convert bytes to PIL Image.
    
    Args:
        data: Image data as bytes
    
    Returns:
        PIL Image object
    """
    return Image.open(io.BytesIO(data))


def image_to_base64(img: Image.Image) -> str:
    """
    Convert PIL Image to base64 string.
    
    Args:
        img: PIL Image object
    
    Returns:
        Base64 encoded string
    """
    buffer = io.BytesIO()
    img.save(buffer, format="PNG")
    return base64.b64encode(buffer.getvalue()).decode()


def base64_to_image(b64: str) -> Image.Image:
    """
    Convert base64 string to PIL Image.
    
    Args:
        b64: Base64 encoded string
    
    Returns:
        PIL Image object
    """
    return Image.open(io.BytesIO(base64.b64decode(b64)))


def get_image_info(img: Image.Image) -> dict:
    """
    Get information about an image.
    
    Args:
        img: PIL Image object
    
    Returns:
        Dictionary with image information
    """
    return {
        "width": img.width,
        "height": img.height,
        "mode": img.mode,
        "format": img.format,
        "size_bytes": len(image_to_bytes(img))
    }


def create_thumbnail(img: Image.Image, max_size: int = 256) -> Image.Image:
    """
    Create a thumbnail of the image.
    
    Args:
        img: PIL Image object
        max_size: Maximum dimension size
    
    Returns:
        Thumbnail PIL Image
    """
    img_copy = img.copy()
    img_copy.thumbnail((max_size, max_size), Image.Resampling.LANCZOS)
    return img_copy


def validate_image(img: Optional[Image.Image]) -> bool:
    """
    Validate that the image is valid.
    
    Args:
        img: PIL Image object to validate
    
    Returns:
        True if valid, False otherwise
    """
    if img is None:
        return False
    try:
        return img.width > 0 and img.height > 0
    except:
        return False


def np_to_pil(arr: np.ndarray) -> Image.Image:
    """
    Convert numpy array to PIL Image.
    
    Args:
        arr: Numpy array
    
    Returns:
        PIL Image
    """
    if arr.dtype != np.uint8:
        arr = arr.astype(np.uint8)
    return Image.fromarray(arr)


def pil_to_np(img: Image.Image) -> np.ndarray:
    """
    Convert PIL Image to numpy array.
    
    Args:
        img: PIL Image
    
    Returns:
        Numpy array
    """
    return np.array(img)