""" Client-side wrapper for font classification with proper preprocessing. Works with both local models and HuggingFace Inference Endpoints. """ import base64 import io import numpy as np import requests import torch import torchvision.transforms as T from PIL import Image from transformers import AutoImageProcessor, Dinov2ForImageClassification def pad_to_square(image): """ Pad image to square while preserving aspect ratio. This is the crucial preprocessing step for font classification. """ if isinstance(image, torch.Tensor): # Convert tensor to PIL for processing if image.dim() == 4: # Batch dimension image = image.squeeze(0) image = T.ToPILImage()(image) if isinstance(image, np.ndarray): image = Image.fromarray(image) if not isinstance(image, Image.Image): raise ValueError(f"Expected PIL Image, got {type(image)}") w, h = image.size max_size = max(w, h) pad_w = (max_size - w) // 2 pad_h = (max_size - h) // 2 padding = (pad_w, pad_h, max_size - w - pad_w, max_size - h - pad_h) return T.Pad(padding, fill=0)(image) class FontClassifierClient: """ Client for font classification that ensures correct preprocessing. Works with both local models and Inference Endpoints. """ def __init__(self, model_name_or_path=None, api_url=None, api_token=None): """ Initialize font classifier client. Args: model_name_or_path: Local model path or HuggingFace model name api_url: Inference Endpoint URL (alternative to local model) api_token: API token for Inference Endpoints """ self.api_url = api_url self.api_token = api_token if api_url: # Using Inference Endpoint self.model = None self.processor = None self.headers = { "Authorization": f"Bearer {api_token}", "Content-Type": "application/json" } if api_token else {} else: # Using local model self.model = Dinov2ForImageClassification.from_pretrained(model_name_or_path) self.processor = AutoImageProcessor.from_pretrained(model_name_or_path) self.model.eval() # Set up preprocessing transform self.preprocess_transform = T.Compose([ T.Lambda(lambda x: x.convert('RGB') if hasattr(x, 'convert') else x), pad_to_square, T.Resize((224, 224)), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) def preprocess_image(self, image): """Apply correct preprocessing to image.""" if isinstance(image, str): image = Image.open(image) return self.preprocess_transform(image) def predict_local(self, image, top_k=5): """Make prediction using local model.""" if self.model is None: raise ValueError("No local model loaded") # Preprocess image processed_image = self.preprocess_image(image) pixel_values = processed_image.unsqueeze(0) # Add batch dimension # Get prediction with torch.no_grad(): outputs = self.model(pixel_values=pixel_values) logits = outputs.logits probabilities = torch.nn.functional.softmax(logits, dim=-1) # Get top-k predictions top_k_indices = torch.topk(logits, k=top_k).indices[0] top_k_labels = [self.model.config.id2label[idx.item()] for idx in top_k_indices] top_k_confidences = [probabilities[0][idx].item() for idx in top_k_indices] return list(zip(top_k_labels, top_k_confidences)) def predict_api(self, image, top_k=5): """Make prediction using Inference Endpoint API.""" if not self.api_url: raise ValueError("No API URL provided") # Preprocess image processed_image = self.preprocess_image(image) # Convert to PIL for API transmission processed_pil = T.ToPILImage()(processed_image) # Convert to bytes img_buffer = io.BytesIO() processed_pil.save(img_buffer, format='PNG') img_bytes = img_buffer.getvalue() # Encode as base64 img_base64 = base64.b64encode(img_bytes).decode() # Make API request payload = { "inputs": img_base64, "parameters": {"top_k": top_k} } response = requests.post(self.api_url, headers=self.headers, json=payload) response.raise_for_status() results = response.json() # Format results if isinstance(results, list) and len(results) > 0: predictions = [(item["label"], item["score"]) for item in results[:top_k]] return predictions else: raise ValueError(f"Unexpected API response format: {results}") def predict(self, image, top_k=5): """ Make prediction with automatic backend selection. Args: image: PIL Image, file path, or numpy array top_k: Number of top predictions to return Returns: List of (label, confidence) tuples """ if self.api_url: return self.predict_api(image, top_k) else: return self.predict_local(image, top_k) @classmethod def from_local_model(cls, model_name_or_path): """Create client for local model.""" return cls(model_name_or_path=model_name_or_path) @classmethod def from_inference_endpoint(cls, api_url, api_token=None): """Create client for Inference Endpoint.""" return cls(api_url=api_url, api_token=api_token) # Convenience functions def predict_font_local(model_name, image_path, top_k=5): """Quick prediction with local model.""" client = FontClassifierClient.from_local_model(model_name) return client.predict(image_path, top_k) def predict_font_api(api_url, image_path, api_token=None, top_k=5): """Quick prediction with Inference Endpoint.""" client = FontClassifierClient.from_inference_endpoint(api_url, api_token) return client.predict(image_path, top_k) # Example usage: if __name__ == "__main__": # Local usage # client = FontClassifierClient.from_local_model("dchen0/font-classifier-v4") # results = client.predict("test_image.png") # API usage # client = FontClassifierClient.from_inference_endpoint("https://your-endpoint.com") # results = client.predict("test_image.png") print("Font Classifier Client ready. Use FontClassifierClient.from_local_model() or FontClassifierClient.from_inference_endpoint()")