| """Custom inference handler for a Dedicated HF Inference Endpoint serving |
| google/siglip2-so400m-patch14-384 as a pure image embedder. |
| |
| Returns get_image_features() -> a single pooled, L2-normalized 1152-float vector, |
| which is exactly what Atlasio's monument_embeddings (vector(1152)) and the |
| recognize pipeline expect. |
| |
| Accepts either: |
| - raw image bytes with Content-Type image/jpeg|png (HF toolkit passes a PIL image), or |
| - JSON {"inputs": "<base64 image>"} or {"inputs": "<http image url>"}. |
| """ |
|
|
| import base64 |
| import io |
| from typing import Any, Dict, List, Union |
|
|
| import torch |
| from PIL import Image |
| from transformers import AutoModel, AutoProcessor |
|
|
|
|
| class EndpointHandler: |
| def __init__(self, path: str = ""): |
| self.device = "cuda" if torch.cuda.is_available() else "cpu" |
| self.processor = AutoProcessor.from_pretrained(path) |
| self.model = AutoModel.from_pretrained(path).to(self.device).eval() |
|
|
| def _to_image(self, inputs: Any) -> Image.Image: |
| if isinstance(inputs, Image.Image): |
| return inputs.convert("RGB") |
| if isinstance(inputs, (bytes, bytearray)): |
| return Image.open(io.BytesIO(bytes(inputs))).convert("RGB") |
| if isinstance(inputs, str): |
| if inputs.startswith("http://") or inputs.startswith("https://"): |
| import urllib.request |
| with urllib.request.urlopen(inputs) as r: |
| return Image.open(io.BytesIO(r.read())).convert("RGB") |
| |
| return Image.open(io.BytesIO(base64.b64decode(inputs))).convert("RGB") |
| raise ValueError(f"Unsupported input type: {type(inputs)}") |
|
|
| @torch.no_grad() |
| def __call__(self, data: Dict[str, Any]) -> Union[List[float], Dict[str, str]]: |
| inputs = data.get("inputs", data) |
| image = self._to_image(inputs) |
| pixel = self.processor(images=image, return_tensors="pt").to(self.device) |
| feats = self.model.get_image_features(**pixel) |
| feats = torch.nn.functional.normalize(feats, p=2, dim=-1) |
| return feats[0].detach().cpu().tolist() |
|
|
