File size: 988 Bytes
3bb5da6 2cf0517 3bb5da6 2cf0517 73e09fb ed9455f 2cf0517 ffdcfec 2cf0517 73e09fb e0cba56 2cf0517 |
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 |
from typing import Dict, List, Any
from PIL import Image
import torch
import base64
from io import BytesIO
from transformers import AutoImageProcessor, Swinv2Model
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
class EndpointHandler():
def __init__(self, path=""):
self.model = Swinv2Model.from_pretrained("microsoft/swinv2-large-patch4-window12to24-192to384-22kto1k-ft", add_pooling_layer = True).to(device)
self.processor = AutoImageProcessor.from_pretrained("microsoft/swinv2-large-patch4-window12to24-192to384-22kto1k-ft")
def __call__(self, data: Any) -> List[float]:
inputs = data.pop("inputs", data)
image = Image.open(BytesIO(base64.b64decode(inputs['image'])))
inputs = self.processor(image, return_tensors="pt").to(device)
with torch.no_grad():
outputs = self.model(**inputs)
pooler_output = outputs.pooler_output
return pooler_output[0].tolist()
|