Update app.py

app.py

@@ -1,12 +1,12 @@
 import torch
 import torch.nn as nn
-from torchvision import models, transforms
+from torchvision import transforms, models
 from PIL import Image
 import gradio as gr
 
-# ---------------------------------------------
+# --------------------------------------------------------------
 # CLASS LABELS (17 classes)
-# ---------------------------------------------
+# --------------------------------------------------------------
 CLASS_LABELS = [
     'Corn_Common_Rust', 'Corn_Gray_Leaf_Spot', 'Corn_Healthy', 'Corn_Northern_Leaf_Blight',
     'Potato_Early_Blight', 'Potato_Healthy', 'Potato_Late_Blight',
@@ -17,27 +17,36 @@ CLASS_LABELS = [
 
 NUM_CLASSES = len(CLASS_LABELS)
 
-# ---------------------------------------------
-# Load Your ResNet50 Model
-# ---------------------------------------------
-class PlantDiseaseResNet(nn.Module):
-    def __init__(self, num_classes):
+# --------------------------------------------------------------
+# RESNET50 MODEL DEFINITION
+# --------------------------------------------------------------
+class ResNetPlantDisease(nn.Module):
+    def __init__(self, num_classes=17):
         super().__init__()
         self.model = models.resnet50(weights=None)
-        self.model.fc = nn.Linear(2048, num_classes)
+        self.model.fc = nn.Sequential(
+            nn.Dropout(0.5),
+            nn.Linear(2048, 512),
+            nn.ReLU(),
+            nn.Dropout(0.3),
+            nn.Linear(512, num_classes)
+        )
 
     def forward(self, x):
         return self.model(x)
 
-model = PlantDiseaseResNet(NUM_CLASSES)
+# --------------------------------------------------------------
+# LOAD MODEL
+# --------------------------------------------------------------
+model = ResNetPlantDisease(num_classes=NUM_CLASSES)
 model.load_state_dict(torch.load("plant_disease_resnet_model.pth", map_location="cpu"))
 model.eval()
 
-# ---------------------------------------------
-# Image Transform (224x224)
-# ---------------------------------------------
+# --------------------------------------------------------------
+# PREPROCESSING
+# --------------------------------------------------------------
 transform = transforms.Compose([
-    transforms.Resize((224, 224)),   # Your image size
+    transforms.Resize((224, 224)),
     transforms.ToTensor(),
     transforms.Normalize(
         mean=[0.485, 0.456, 0.406],
@@ -45,29 +54,34 @@ transform = transforms.Compose([
     )
 ])
 
-# ---------------------------------------------
-# Prediction Function
-# ---------------------------------------------
-def predict(image):
-    img = Image.fromarray(image)
+# --------------------------------------------------------------
+# DISEASE CLASSIFICATION FUNCTION
+# --------------------------------------------------------------
+def classify(image):
+    img = Image.fromarray(image).convert("RGB")
     x = transform(img).unsqueeze(0)
 
     with torch.no_grad():
         logits = model(x)
-        probs = logits.softmax(dim=1)[0]
+        probs = torch.softmax(logits[0], dim=0)
 
-    result = {CLASS_LABELS[i]: float(probs[i]) for i in range(NUM_CLASSES)}
-    return result
+    top_probs, top_idxs = torch.topk(probs, 5)
+    predictions = {
+        CLASS_LABELS[top_idxs[i].item()]: float(top_probs[i].item())
+        for i in range(5)
+    }
 
-# ---------------------------------------------
-# Gradio Interface + HuggingFace API
-# ---------------------------------------------
+    return predictions
+
+# --------------------------------------------------------------
+# GRADIO INTERFACE WITH API
+# --------------------------------------------------------------
 demo = gr.Interface(
-    fn=predict,
+    fn=classify,
     inputs=gr.Image(type="numpy"),
-    outputs=gr.Label(num_top_classes=3),
-    title="Plant Disease Detector (ResNet50)",
-    description="Upload a crop leaf image to predict its disease."
+    outputs=gr.Label(num_top_classes=5),
+    title="Plant Disease Classification (ResNet50)",
+    description="Upload a leaf image to detect plant disease."
 )
 
-demo.launch()
+demo.launch()