src/egg_damage/preprocessing.py

from __future__ import annotations

from pathlib import Path
from typing import Any

import numpy as np
from PIL import Image, ImageOps


def load_pil_image(image: str | Path | Image.Image, mode: str = "RGB") -> Image.Image:
    if isinstance(image, Image.Image):
        img = image.copy()
    else:
        img = Image.open(image)
    img = ImageOps.exif_transpose(img)
    return img.convert(mode)


def resize_image(image: Image.Image, size: int | tuple[int, int]) -> Image.Image:
    if isinstance(size, int):
        size = (size, size)
    return image.resize(size, Image.Resampling.BILINEAR)


def pil_to_uint8_array(image: Image.Image) -> np.ndarray:
    arr = np.asarray(image)
    if arr.dtype != np.uint8:
        arr = np.clip(arr, 0, 255).astype(np.uint8)
    return arr


def image_size_from_config(config: dict[str, Any]) -> int:
    return int(config.get("preprocessing", {}).get("image_size", 224))