File size: 2,024 Bytes
1032eea |
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 50 51 52 53 |
import torch
import torch.nn as nn
from torchvision import models, transforms
from PIL import Image
import io
class EndpointHandler():
def __init__(self, path=""):
# 1. Define device
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 2. Define class names (Matches alphabetical order used in training)
self.class_names = ['Gray Leaf Spot', 'Healthy']
# 3. Initialize Model Architecture (Update if using EfficientNet)
# Note: You can make this dynamic or hardcode it to your best model
self.model = models.resnet50(weights=None)
self.model.fc = nn.Linear(self.model.fc.in_features, len(self.class_names))
# 4. Load weights (Hugging Face passes the folder path in 'path')
# Ensure 'model.pth' is the name of your file in the root
state_dict = torch.load(f"{path}/model.pth", map_location=self.device)
self.model.load_state_dict(state_dict)
self.model.to(self.device)
self.model.eval()
# 5. Define Preprocessing
self.transform = transforms.Compose([
transforms.Resize((256, 256)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
def __call__(self, data):
# Data is a dictionary containing the image bytes
inputs = data.pop("inputs", data)
# Convert bytes to PIL Image
image = Image.open(io.BytesIO(inputs)).convert("RGB")
# Preprocess
tensor = self.transform(image).unsqueeze(0).to(self.device)
# Inference
with torch.no_grad():
outputs = self.model(tensor)
probs = torch.nn.functional.softmax(outputs, dim=1)
conf, pred_idx = torch.max(probs, 1)
# Return formatted result for the widget
return [
{"label": self.class_names[pred_idx.item()], "score": conf.item()}
] |