image_authenticity/utils/preprocessing.py

"""
utils/preprocessing.py

Image loading and preprocessing utilities.
"""

import io
import requests
import numpy as np
from PIL import Image
from typing import Union


def load_image(source: Union[str, Image.Image, np.ndarray, bytes]) -> Image.Image:
    """
    Universal image loader — accepts:
      - str path to a local file
      - str URL to fetch remotely
      - PIL Image (pass-through)
      - numpy array
      - raw bytes
    Always returns an RGB PIL Image.
    """
    if isinstance(source, Image.Image):
        img = source

    elif isinstance(source, np.ndarray):
        img = Image.fromarray(source)

    elif isinstance(source, bytes):
        img = Image.open(io.BytesIO(source))

    elif isinstance(source, str):
        if source.startswith("http://") or source.startswith("https://"):
            response = requests.get(source, timeout=15)
            response.raise_for_status()
            img = Image.open(io.BytesIO(response.content))
        else:
            img = Image.open(source)
    else:
        raise ValueError(f"Unsupported image source type: {type(source)}")

    if img.mode != "RGB":
        img = img.convert("RGB")
    return img


def preprocess_for_display(image: Image.Image, max_size: int = 512) -> Image.Image:
    """
    Resize image for display (keep aspect ratio, cap at max_size px).
    """
    w, h = image.size
    if max(w, h) > max_size:
        scale = max_size / max(w, h)
        new_w = int(w * scale)
        new_h = int(h * scale)
        image = image.resize((new_w, new_h), Image.LANCZOS)
    return image