dchen0/font_classifier_v4

1"""2Standalone FontClassifierImageProcessor for HuggingFace Hub deployment.3"""4import numpy as np5import torch6import torchvision.transforms as T7from PIL import Image8from transformers import AutoImageProcessor9 10 11def pad_to_square(image):12    """13    Shared utility function to pad image to square while preserving aspect ratio.14    Works with both PIL Images and numpy arrays.15    """16    if isinstance(image, Image.Image):17        w, h = image.size18        max_size = max(w, h)19        pad_w = (max_size - w) // 220        pad_h = (max_size - h) // 221        padding = (pad_w, pad_h, max_size - w - pad_w, max_size - h - pad_h)22        return T.Pad(padding, fill=0)(image)23    elif isinstance(image, np.ndarray):24        # Convert numpy array to PIL, process, then back25        if image.ndim == 3 and image.shape[2] == 3:  # RGB26            pil_img = Image.fromarray(image.astype(np.uint8))27            padded_pil = pad_to_square(pil_img)  # Recursive call with PIL image28            return np.array(padded_pil)29    return image30 31class FontClassifierImageProcessor(AutoImageProcessor):32    """33    Custom image processor that includes pad_to_square transformation.34    This ensures that Inference Endpoints will apply the same preprocessing as training.35    """36    37    model_input_names = ["pixel_values"]38    39    def __init__(self, *args, **kwargs):40        super().__init__(*args, **kwargs)41        # Store the original preprocess method42        self._original_preprocess = super().preprocess43    44    def preprocess(self, images, **kwargs):45        """Override preprocess to include pad_to_square"""46        # Handle single image or list of images47        if isinstance(images, (Image.Image, np.ndarray)):48            images = [images]49            single_image = True50        else:51            single_image = False52        53        # Apply pad_to_square to each image using shared utility54        padded_images = [pad_to_square(img) for img in images]55        56        # Call original preprocess with padded images57        result = self._original_preprocess(padded_images, **kwargs)58        59        # If single image was passed, ensure we return the format expected60        if single_image and isinstance(result, dict) and 'pixel_values' in result:61            # Keep batch dimension for consistency62            pass63            64        return result65 66# Register the custom processor class67AutoImageProcessor.register("FontClassifierImageProcessor", FontClassifierImageProcessor)