ShlokArora2709/jester

1"""2HF_server/preprocess.py — Image and text preprocessing utilities.3"""4 5from __future__ import annotations6 7import base648import binascii9from io import BytesIO10from typing import Dict11 12import numpy as np13from PIL import Image, ImageOps14 15_CLIP_IMAGE_SIZE = 22416_CLIP_MEAN = (0.48145466, 0.4578275, 0.40821073)17_CLIP_STD = (0.26862954, 0.26130258, 0.27577711)18 19 20def decode_image_b64(image_b64: str, max_bytes: int) -> Image.Image:21    """22    decode_image_b64(image_b64: str, max_bytes: int) -> Image.Image23 24    Decode a base64 image string into a PIL Image in RGB mode.25    Accepts optional data URI prefixes.26    """27    if not image_b64:28        raise ValueError("image_b64 is required")29 30    payload = image_b6431    if "," in image_b64 and image_b64.strip().lower().startswith("data:"):32        payload = image_b64.split(",", 1)[1]33 34    try:35        raw = base64.b64decode(payload, validate=True)36    except binascii.Error as exc:37        raise ValueError("invalid base64 payload") from exc38 39    if len(raw) > max_bytes:40        raise ValueError("image exceeds max byte limit")41 42    try:43        with Image.open(BytesIO(raw)) as img:44            return img.convert("RGB")45    except Exception as exc:46        raise ValueError("invalid image data") from exc47 48 49def clip_image_to_tensor(image: Image.Image) -> np.ndarray:50    """51    clip_image_to_tensor(image: Image.Image) -> np.ndarray52 53    Convert a PIL Image to a CLIP-compatible tensor: (1, 3, 224, 224).54    """55    image = ImageOps.fit(image, (_CLIP_IMAGE_SIZE, _CLIP_IMAGE_SIZE), method=Image.BICUBIC)56    array = np.asarray(image).astype(np.float32) / 255.057    array = (array - np.array(_CLIP_MEAN, dtype=np.float32)) / np.array(58        _CLIP_STD, dtype=np.float3259    )60    array = np.transpose(array, (2, 0, 1))61    return np.expand_dims(array, axis=0)62 63 64def build_text_inputs(tokenizer, text: str, max_length: int = 77) -> Dict[str, np.ndarray]:65    """66    build_text_inputs(tokenizer, text: str, max_length: int = 77) -> Dict[str, np.ndarray]67 68    Tokenize text for CLIP-style text encoders, returning numpy inputs.69    """70    if not text or not text.strip():71        raise ValueError("text is required")72 73    tokens = tokenizer(74        text,75        padding="max_length",76        truncation=True,77        max_length=max_length,78        return_tensors="np",79    )80 81    inputs: Dict[str, np.ndarray] = {}82    for key, value in tokens.items():83        if isinstance(value, np.ndarray):84            inputs[key] = value.astype(np.int64)85    return inputs86 87 88def normalize_ocr_text(text: str) -> str:89    """90    normalize_ocr_text(text: str) -> str91 92    Normalize OCR output by trimming and collapsing whitespace.93    """94    if not text:95        return ""96    return " ".join(text.split())97