| from typing import Dict, Any |
| import io, base64, torch, torchvision |
| from PIL import Image |
| from torchvision import transforms as T |
| from safetensors.torch import load_file |
|
|
| class EndpointHandler: |
| def __init__(self, path: str = "."): |
| self.labels = ['battery', 'biological', 'cardboard', 'clothes', 'glass', 'metal', 'paper', 'plastic', 'shoes', 'trash'] |
| self.model = torchvision.models.efficientnet_v2_s(weights=None) |
| nf = self.model.classifier[1].in_features |
| self.model.classifier = torch.nn.Sequential( |
| torch.nn.Linear(nf, 256), |
| torch.nn.ReLU(inplace=True), |
| torch.nn.Dropout(0.5), |
| torch.nn.Linear(256, len(self.labels)), |
| ) |
| state = load_file(f"{path}/model.safetensors", device="cpu") |
| self.model.load_state_dict(state) |
| self.model.eval() |
|
|
| self.preprocess = T.Compose([ |
| T.Resize((224, 224)), |
| T.ToTensor(), |
| T.Normalize(mean=[0.485,0.456,0.406], |
| std =[0.229,0.224,0.225]) |
| ]) |
|
|
| def __call__(self, data: Dict[str, Any]): |
| img_bytes = data["inputs"] |
| if isinstance(img_bytes, str): |
| img_bytes = base64.b64decode(img_bytes) |
| img = Image.open(io.BytesIO(img_bytes)).convert("RGB") |
| x = self.preprocess(img).unsqueeze(0) |
| with torch.no_grad(): |
| probs = self.model(x).softmax(1)[0] |
| topk = probs.topk(5) |
| return [{"label": self.labels[i], "score": float(topk.values[j])} |
| for j,i in enumerate(topk.indices)] |
|
|
