File size: 1,277 Bytes
e9dc81c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 | from typing import Dict, Any
from PIL import Image
from transformers import CLIPProcessor, CLIPModel
import torch
import base64
from io import BytesIO
class EndpointHandler:
def __init__(self, path=""):
self.model = CLIPModel.from_pretrained(path)
self.processor = CLIPProcessor.from_pretrained(path)
self.model.eval()
def _to_image(self, x) -> Image.Image:
if isinstance(x, Image.Image):
return x.convert("RGB")
if isinstance(x, (bytes, bytearray)):
return Image.open(BytesIO(x)).convert("RGB")
if isinstance(x, str):
return Image.open(BytesIO(base64.b64decode(x))).convert("RGB")
if isinstance(x, dict) and "image" in x:
return self._to_image(x["image"])
raise ValueError("Unsupported image input")
def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
inputs = data.get("inputs", data)
image = self._to_image(inputs)
proc = self.processor(images=image, return_tensors="pt")
with torch.no_grad():
feats = self.model.get_image_features(**proc)
feats = feats / feats.norm(p=2, dim=-1, keepdim=True) # L2 normalize
return {"embedding": feats[0].tolist(), "dim": int(feats.shape[-1])} |