utils/image_utils.py

"""
InvoiceForge AI — utils/image_utils.py
Shared image loading and transformation helpers.

Provides:
  - load_image_from_request(): FormData or base64 JSON
  - correct_perspective(): document quad detection + warp
  - remove_shadow(): illumination normalisation
  - detect_blur(): Laplacian variance metric
  - score_quality(): composite image quality dict
"""

from __future__ import annotations

import base64
import io
import logging
import re
from typing import Optional

import cv2
import numpy as np
from flask import Request
from PIL import Image

logger = logging.getLogger(__name__)


# ─────────────────────────────────────────────────────────────────────────────
# IMAGE LOADING
# ─────────────────────────────────────────────────────────────────────────────

def load_image_from_request(request: Request) -> np.ndarray:
    """
    Load a BGR numpy image array from a Flask request.

    Accepts:
      - multipart/form-data with key 'file' or 'image'
      - application/json with key 'image' (base64 data URI or raw base64)

    Raises:
        ValueError: If no valid image source is found.

    Returns:
        BGR numpy array (uint8).
    """
    # ── 1. FormData ────────────────────────────────────────────────────────
    if request.files:
        file_obj = request.files.get("file") or request.files.get("image")
        if file_obj:
            img_bytes = file_obj.read()
            img = _decode_bytes(img_bytes, label="uploaded file")
            if img is not None:
                return img

    # ── 2. JSON body ───────────────────────────────────────────────────────
    data: dict = request.get_json(silent=True) or {}
    img_b64: str = data.get("image", "")
    if img_b64:
        # Strip data URI prefix: "data:image/jpeg;base64,..."
        img_b64 = re.sub(r"^data:image/\w+;base64,", "", img_b64)
        try:
            img_bytes = base64.b64decode(img_b64)
            img = _decode_bytes(img_bytes, label="base64 JSON")
            if img is not None:
                return img
        except Exception as exc:
            raise ValueError(f"Base64 decode failed: {exc}") from exc

    raise ValueError(
        "No image found in request. "
        "Send 'file' in FormData or 'image' (base64) in JSON body."
    )


def _decode_bytes(img_bytes: bytes, label: str = "image") -> Optional[np.ndarray]:
    """Decode raw image bytes to BGR numpy array using OpenCV."""
    nparr = np.frombuffer(img_bytes, np.uint8)
    img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
    if img is None:
        logger.warning("Failed to decode %s as image.", label)
        return None
    logger.debug("Loaded %s: shape=%s", label, img.shape)
    return img


def pil_to_bgr(pil_img: Image.Image) -> np.ndarray:
    """Convert a PIL Image (RGB) to BGR numpy array."""
    rgb = np.array(pil_img.convert("RGB"))
    return cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)


def bgr_to_pil(img_bgr: np.ndarray) -> Image.Image:
    """Convert BGR numpy array to PIL RGB Image."""
    rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
    return Image.fromarray(rgb)


# ─────────────────────────────────────────────────────────────────────────────
# PERSPECTIVE CORRECTION  (re-exported from preprocessing for convenience)
# ─────────────────────────────────────────────────────────────────────────────

def correct_perspective(img_bgr: np.ndarray) -> np.ndarray:
    """Wrapper — see ocr.preprocessing.correct_perspective for full docs."""
    from ocr.preprocessing import correct_perspective as _cp
    return _cp(img_bgr)


# ─────────────────────────────────────────────────────────────────────────────
# SHADOW REMOVAL  (re-exported from preprocessing)
# ─────────────────────────────────────────────────────────────────────────────

def remove_shadow(img_bgr: np.ndarray) -> np.ndarray:
    """Wrapper — see ocr.preprocessing.remove_shadow for full docs."""
    from ocr.preprocessing import remove_shadow as _rs
    return _rs(img_bgr)


# ─────────────────────────────────────────────────────────────────────────────
# QUALITY METRICS
# ─────────────────────────────────────────────────────────────────────────────

def detect_blur(img: np.ndarray) -> float:
    """
    Return Laplacian variance of the image.
    Values below 80 typically indicate a blurry / out-of-focus image.
    """
    gray = img if len(img.shape) == 2 else cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    return float(cv2.Laplacian(gray, cv2.CV_64F).var())


def score_quality(img: np.ndarray) -> dict:
    """
    Return a quality assessment dict.

    Keys:
      blur_score     — Laplacian variance (higher = sharper)
      brightness     — mean pixel value 0-255
      contrast       — pixel std dev
      resolution     — (width, height) tuple
      quality_score  — aggregate 0.0-1.0
      is_blurry      — bool
    """
    from ocr.preprocessing import score_image_quality
    return score_image_quality(img)