io_tools.py

from __future__ import annotations

import io
import logging

import numpy as np
from PIL import Image, UnidentifiedImageError

log = logging.getLogger(__name__)

MAX_IMAGE_BYTES = 50 * 1024 * 1024
MIN_IMAGE_DIMENSION = 16
MAX_IMAGE_DIMENSION = 10000


def load_image_from_bytes(data: bytes) -> np.ndarray:
    if not data:
        raise ValueError("Image payload is empty")
    if len(data) > MAX_IMAGE_BYTES:
        raise ValueError("Image payload is too large")

    try:
        with Image.open(io.BytesIO(data)) as img:
            img.load()
            image = np.asarray(img.convert("RGB"))
    except UnidentifiedImageError:
        log.warning("Unidentified image. len=%s first16=%r", len(data), data[:16])
        raise

    height, width = image.shape[:2]
    if min(height, width) < MIN_IMAGE_DIMENSION:
        raise ValueError("Image is too small")
    if max(height, width) > MAX_IMAGE_DIMENSION:
        raise ValueError("Image is too large")
    return image


__all__ = ["load_image_from_bytes"]