Update app.py

app.py

@@ -1,12 +1,13 @@
 import torch
 import torch.nn as nn
-from torchvision import transforms, models
+from torchvision.models import resnet18, resnet34, resnet50
+from torchvision import transforms
 from PIL import Image
 import gradio as gr
 
-# --------------------------------------------------------------
-# CLASS LABELS (17 classes)
-# --------------------------------------------------------------
+# -----------------------
+# CLASS LABELS
+# -----------------------
 CLASS_LABELS = [
     'Corn_Common_Rust', 'Corn_Gray_Leaf_Spot', 'Corn_Healthy', 'Corn_Northern_Leaf_Blight',
     'Potato_Early_Blight', 'Potato_Healthy', 'Potato_Late_Blight',
@@ -17,34 +18,50 @@ CLASS_LABELS = [
 
 NUM_CLASSES = len(CLASS_LABELS)
 
-# --------------------------------------------------------------
-# RESNET50 MODEL DEFINITION
-# --------------------------------------------------------------
+
+# -----------------------
+# MODEL ARCHITECTURE
+# -----------------------
 class ResNetPlantDisease(nn.Module):
-    def __init__(self, num_classes=17):
+    def __init__(self, num_classes=17, model_name='resnet50', pretrained=False):
         super().__init__()
-        self.model = models.resnet50(weights=None)
-        self.model.fc = nn.Sequential(
+
+        if model_name == 'resnet18':
+            self.backbone = resnet18(weights=None)
+            num_features = 512
+        elif model_name == 'resnet34':
+            self.backbone = resnet34(weights=None)
+            num_features = 512
+        elif model_name == 'resnet50':
+            self.backbone = resnet50(weights=None)
+            num_features = 2048
+        else:
+            raise ValueError("Unsupported model name")
+
+        self.backbone.fc = nn.Sequential(
             nn.Dropout(0.5),
-            nn.Linear(2048, 512),
+            nn.Linear(num_features, 512),
             nn.ReLU(),
             nn.Dropout(0.3),
             nn.Linear(512, num_classes)
         )
 
     def forward(self, x):
-        return self.model(x)
+        return self.backbone(x)
 
-# --------------------------------------------------------------
+
+# -----------------------
 # LOAD MODEL
-# --------------------------------------------------------------
-model = ResNetPlantDisease(num_classes=NUM_CLASSES)
-model.load_state_dict(torch.load("plant_disease_resnet_model.pth", map_location="cpu"))
+# -----------------------
+model = ResNetPlantDisease(num_classes=NUM_CLASSES, model_name='resnet50')
+state = torch.load("plant_disease_resnet_model.pth", map_location="cpu")
+model.load_state_dict(state)
 model.eval()
 
-# --------------------------------------------------------------
-# PREPROCESSING
-# --------------------------------------------------------------
+
+# -----------------------
+# TRANSFORMS
+# -----------------------
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor(),
@@ -54,34 +71,31 @@ transform = transforms.Compose([
     )
 ])
 
-# --------------------------------------------------------------
-# DISEASE CLASSIFICATION FUNCTION
-# --------------------------------------------------------------
-def classify(image):
-    img = Image.fromarray(image).convert("RGB")
-    x = transform(img).unsqueeze(0)
+
+# -----------------------
+# PREDICT FUNCTION
+# -----------------------
+def predict(image):
+    img = Image.fromarray(image)
+    img = transform(img).unsqueeze(0)
 
     with torch.no_grad():
-        logits = model(x)
-        probs = torch.softmax(logits[0], dim=0)
+        outputs = model(img)
+        probs = torch.softmax(outputs, dim=1)[0]
 
-    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)
-    }
+    result = {CLASS_LABELS[i]: float(probs[i]) for i in range(NUM_CLASSES)}
+    return result
 
-    return predictions
 
-# --------------------------------------------------------------
-# GRADIO INTERFACE WITH API
-# --------------------------------------------------------------
+# -----------------------
+# GRADIO UI + API
+# -----------------------
 demo = gr.Interface(
-    fn=classify,
+    fn=predict,
     inputs=gr.Image(type="numpy"),
-    outputs=gr.Label(num_top_classes=5),
-    title="Plant Disease Classification (ResNet50)",
-    description="Upload a leaf image to detect plant disease."
+    outputs=gr.Label(num_top_classes=3),
+    title="Plant Disease Detection - ResNet50",
+    description="Upload a leaf image to detect crop disease."
 )
 
 demo.launch()