File size: 1,967 Bytes
0dd652c 6431538 0dd652c 6431538 0dd652c | 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 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 | from typing import Dict, Any, List, Union
import base64
from io import BytesIO
import torch
import numpy as np
from PIL import Image
from transformers import XCLIPProcessor, XCLIPModel
class EndpointHandler:
def __init__(self, path=""):
self.model = XCLIPModel.from_pretrained(path)
self.processor = XCLIPProcessor.from_pretrained(path)
self.model.eval()
try:
self.num_frames = self.model.config.vision_config.num_frames
except AttributeError:
self.num_frames = 8
def _decode(self, b64: str) -> Image.Image:
return Image.open(BytesIO(base64.b64decode(b64))).convert("RGB")
def _sample(self, frames: List[Image.Image]) -> List[np.ndarray]:
idx = np.linspace(0, len(frames) - 1, self.num_frames).round().astype(int)
return [np.array(frames[i]) for i in idx]
def __call__(self, data: Dict[str, Any]) -> Union[Dict, List]:
inputs = data.get("inputs", data)
frames_b64 = inputs["frames"] if isinstance(inputs, dict) else inputs
labels = inputs.get("candidate_labels") if isinstance(inputs, dict) else None
frames = [self._decode(f) for f in frames_b64]
if not frames:
raise ValueError("No frames provided")
video = self._sample(frames)
if labels:
proc = self.processor(text=labels, videos=video,
return_tensors="pt", padding=True)
with torch.no_grad():
probs = self.model(**proc).logits_per_video.softmax(dim=1)[0]
return [{"label": l, "score": float(s)} for l, s in zip(labels, probs)]
proc = self.processor(videos=video, return_tensors="pt")
with torch.no_grad():
feats = self.model.get_video_features(pixel_values=proc["pixel_values"])
feats = feats / feats.norm(p=2, dim=-1, keepdim=True)
return {"embedding": feats[0].tolist(), "dim": int(feats.shape[-1])} |