inference/preprocessing.py

"""
Standalone PaddleOCR-VL Image Preprocessing (No PyTorch / No Transformers)
==========================================================================
Replicates the official PaddleOCRVLImageProcessor pipeline using only
Pillow + NumPy.

Pipeline:
  1. Convert to RGB
  2. smart_resize → both dims divisible by 28, within pixel budget
  3. Rescale to [0, 1]
  4. Normalize (CLIP mean/std)
  5. Patchify → (num_patches, 3, 14, 14)
"""

from __future__ import annotations

import math
from typing import Tuple

import numpy as np
from PIL import Image

# ---------------------------------------------------------------------------
# Constants (from PaddleOCR-VL config)
# ---------------------------------------------------------------------------

PATCH_SIZE = 14
MERGE_SIZE = 2
FACTOR = PATCH_SIZE * MERGE_SIZE  # 28

MIN_PIXELS = 28 * 28 * 130   # 101,920
MAX_PIXELS = 28 * 28 * 1280  # 1,003,520

# CLIP mean / std (OpenAI variant used by PaddleOCR-VL)
IMAGE_MEAN = np.array([0.48145466, 0.4578275, 0.40821073], dtype=np.float32)
IMAGE_STD = np.array([0.26862954, 0.26130258, 0.27577711], dtype=np.float32)


# ---------------------------------------------------------------------------
# smart_resize
# ---------------------------------------------------------------------------

def smart_resize(
    height: int,
    width: int,
    factor: int = FACTOR,
    min_pixels: int = MIN_PIXELS,
    max_pixels: int = MAX_PIXELS,
) -> Tuple[int, int]:
    """
    Rescale so that:
      1. Both dimensions are divisible by `factor`.
      2. Total pixels ∈ [min_pixels, max_pixels].
      3. Aspect ratio is preserved as closely as possible.
    """
    if height < factor:
        width = round((width * factor) / height)
        height = factor
    if width < factor:
        height = round((height * factor) / width)
        width = factor

    if max(height, width) / min(height, width) > 200:
        raise ValueError(
            f"Aspect ratio too extreme: {max(height, width) / min(height, width)}"
        )

    h_bar = round(height / factor) * factor
    w_bar = round(width / factor) * factor

    if h_bar * w_bar > max_pixels:
        beta = math.sqrt((height * width) / max_pixels)
        h_bar = math.floor(height / beta / factor) * factor
        w_bar = math.floor(width / beta / factor) * factor
    elif h_bar * w_bar < min_pixels:
        beta = math.sqrt(min_pixels / (height * width))
        h_bar = math.ceil(height * beta / factor) * factor
        w_bar = math.ceil(width * beta / factor) * factor

    return h_bar, w_bar


# ---------------------------------------------------------------------------
# Preprocessing
# ---------------------------------------------------------------------------

def preprocess(
    image: Image.Image,
) -> Tuple[np.ndarray, Tuple[int, int, int]]:
    """
    Preprocess a PIL image into patch tensor for ONNX inference.

    Args:
        image: PIL RGB image (any size).

    Returns:
        pixel_values: (num_patches, 3, 14, 14) float32 array
        grid_thw: (grid_t, grid_h, grid_w) — temporal=1 always
    """
    # 1. Convert to RGB
    img = image.convert("RGB")

    # 2. smart_resize
    width, height = img.size
    new_h, new_w = smart_resize(height, width)

    if (new_w, new_h) != (width, height):
        img = img.resize((new_w, new_h), Image.BICUBIC)

    # 3. To numpy, rescale to [0, 1]
    arr = np.array(img, dtype=np.float32) / 255.0

    # 4. Normalize (CHW)
    arr = (arr - IMAGE_MEAN.reshape(1, 1, 3)) / IMAGE_STD.reshape(1, 1, 3)
    arr = arr.transpose(2, 0, 1)  # HWC → CHW

    # 5. Patchify: (C, H, W) → (num_patches, C, 14, 14)
    c, h, w = arr.shape
    grid_h = h // PATCH_SIZE
    grid_w = w // PATCH_SIZE
    grid_t = 1  # temporal patches = 1 for images

    # Reshape into grid of patches
    patches = arr.reshape(c, grid_h, PATCH_SIZE, grid_w, PATCH_SIZE)
    patches = patches.transpose(1, 3, 0, 2, 4)  # → (grid_h, grid_w, c, 14, 14)
    patches = patches.reshape(-1, c, PATCH_SIZE, PATCH_SIZE)  # → (N, 3, 14, 14)

    # Handle temporal dim (always 1 for images)
    patches = np.tile(patches, (grid_t, 1, 1, 1))

    grid_thw = (grid_t, grid_h, grid_w)

    return patches.astype(np.float32), grid_thw


def preprocess_for_onnx(image: Image.Image) -> Tuple[np.ndarray, np.ndarray]:
    """
    Preprocess image and return ONNX-ready inputs.

    Args:
        image: PIL RGB image.

    Returns:
        pixel_values: (1, num_patches, 3, 14, 14) float32
        position_ids: (1, 1) int64
    """
    patches, _grid_thw = preprocess(image)
    # Add batch dimension
    pixel_values = patches[np.newaxis, ...]  # (1, N, 3, 14, 14)
    position_ids = np.zeros((1, 1), dtype=np.int64)
    return pixel_values.astype(np.float32), position_ids


# ---------------------------------------------------------------------------
# Reverse: patches → image (for verification)
# ---------------------------------------------------------------------------

def patches_to_image(
    pixel_values: np.ndarray,
    grid_h: int,
    grid_w: int,
) -> Image.Image:
    """
    Reconstruct an image from patch tensor (for debugging).

    Args:
        pixel_values: (N, 3, 14, 14) or (1, N, 3, 14, 14)
        grid_h, grid_w: grid dimensions

    Returns:
        PIL Image (approx reconstruction of preprocessed input)
    """
    if pixel_values.ndim == 5:
        pixel_values = pixel_values.squeeze(0)  # remove batch

    n_patches = grid_h * grid_w
    patches = pixel_values[:n_patches]  # (N, 3, 14, 14)

    # Un-patchify
    c = patches.shape[1]
    ps = patches.shape[2]
    patches = patches.reshape(grid_h, grid_w, c, ps, ps)
    patches = patches.transpose(2, 0, 3, 1, 4)  # → (c, grid_h, ps, grid_w, ps)
    img = patches.reshape(c, grid_h * ps, grid_w * ps)  # (c, H, W)

    # De-normalize
    img = img.transpose(1, 2, 0)  # CHW → HWC
    img = img * IMAGE_STD.reshape(1, 1, 3) + IMAGE_MEAN.reshape(1, 1, 3)
    img = np.clip(img * 255, 0, 255).astype(np.uint8)

    return Image.fromarray(img)