Update app.py

app.py

@@ -1,6 +1,6 @@
 import torch
 import torch.nn as nn
-from torchvision.models import resnet18, resnet34, resnet50
+from torchvision.models import resnet50
 from torchvision import transforms
 from PIL import Image
 import gradio as gr
@@ -24,13 +24,7 @@ class ResNetPlantDisease(nn.Module):
     def __init__(self, num_classes=17, model_name='resnet50', pretrained=False):
         super().__init__()
 
-        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':
+        if model_name == 'resnet50':
             self.backbone = resnet50(weights=None)
             num_features = 2048
         else:
@@ -75,8 +69,8 @@ def predict(image):
 demo = gr.Interface(
     fn=predict,
     inputs=gr.Image(type="numpy"),
-    outputs=gr.Label(num_top_classes=3),
-    title="Plant Disease Detection - ResNet50",
+    outputs=gr.Label(num_top_classes=1),
+    title="Crop Disease Detection - ResNet",
     description="Upload a leaf image to detect crop disease."
 )