Upload folder using huggingface_hub

fe44a6e verified 3 months ago

6.03 kB

"""
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)

	"""
	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)