Add custom handler for inference widget

README.md

@@ -1,6 +1,7 @@
 ---
 language: en
 license: mit
+pipeline_tag: image-classification
 tags:
 - agricultural-ai
 - maize

code/handler.py

@@ -0,0 +1,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()}
+        ]