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()