updated one

app.py

@@ -2,75 +2,69 @@ import gradio as gr
 import os
 import torch
 import torchvision.transforms as transforms
-import torchvision.models.resnet
+import torchvision.models as models
+from torchvision.models import ResNet18_Weights
+import torch.nn as nn
 from PIL import Image
 from timeit import default_timer as timer
 from typing import Tuple, Dict
 
-# 1. Setup class names
-class_names = ['Cabbage', 'Carrot', 'Cucumber', 'tomato']
+# 1. Class names (must match training order!)
+class_names = ['Cabbage', 'Carrot', 'Cucumber', 'Tomato']
 
+# 2. Image preprocessing pipeline from ResNet18 weights
+weights = ResNet18_Weights.DEFAULT
+manual_transforms = weights.transforms()
 
-# 2. Define transforms (adjust as needed, example here)
-manual_transforms = transforms.Compose([
-    transforms.Resize((224, 224)),
-    transforms.ToTensor(),
-    transforms.Normalize(
-        mean=[0.485, 0.456, 0.406],  # ImageNet mean
-        std=[0.229, 0.224, 0.225]    # ImageNet std
-    )
-])
-
-# 3. Prepare the ResNet model
-torch.serialization.add_safe_globals([torchvision.models.resnet.ResNet])
-
-resnet = torchvision.models.resnet18(weights=None, num_classes=len(class_names))
-
-# Load your saved state dict checkpoint safely
-resnet.load_state_dict(
-    torch.load("resnet_model.pth", map_location=torch.device("cpu"), weights_only=True)
-)
+# 3. Load model with correct architecture
+def get_resnet18():
+    model = models.resnet18(weights=weights)  # ← NEW: use correct weight object
+    num_ftrs = model.fc.in_features
+    model.fc = nn.Linear(num_ftrs, len(class_names))  # ← match training config
+    return model
 
+resnet = get_resnet18()
+resnet.load_state_dict(torch.load("resnet_model.pth", map_location=torch.device("cpu")))
 resnet.eval()
 
 # 4. Prediction function
 def predict(img: Image.Image) -> Tuple[Dict[str, float], float]:
     start_time = timer()
-    img_tensor = manual_transforms(img).unsqueeze(0)  # add batch dimension
+    img_tensor = manual_transforms(img).unsqueeze(0)
 
     with torch.inference_mode():
-        preds = torch.softmax(resnet(img_tensor), dim=1)
+        outputs = resnet(img_tensor)
+        probs = torch.softmax(outputs, dim=1)[0]
 
-    pred_labels_and_probs = {class_names[i]: float(preds[0][i]) for i in range(len(class_names))}
+    pred_labels_and_probs = {class_names[i]: float(probs[i]) for i in range(len(class_names))}
     pred_time = round(timer() - start_time, 5)
 
     return pred_labels_and_probs, pred_time
 
-# 5. Gradio app UI
+# 5. Gradio app
 with gr.Blocks() as gradio_app:
-    gr.HTML("<h1 style='text-align: center'>Rice Disease Classification - ResNet Model</h1>")
-    gr.HTML(
-        """
+    gr.HTML("<h1 style='text-align: center'>Vegetable Classification with ResNet18</h1>")
+    gr.HTML("""
         <h3 style='text-align: center'>
-        Follow me for more!
-        <a href='https://github.com/ExplorerGumel' target='_blank'>Github</a> | 
-        <a href='https://www.linkedin.com/in/munzali-alhassan/' target='_blank'>Linkedin</a>
+        Created by <a href='https://github.com/ExplorerGumel' target='_blank'>@ExplorerGumel</a>
         </h3>
-        """
-    )
+    """)
 
     with gr.Row():
         with gr.Column():
-            image = gr.Image(type="pil", label="Upload Image")
-            infer = gr.Button(value="Predict")
+            image_input = gr.Image(type="pil", label="Upload Image")
+            infer_btn = gr.Button("Predict")
 
-            example_list = [["examples/" + example] for example in os.listdir("examples")]
-            gr.Examples(examples=example_list, inputs=[image])
+            example_list = []
+            if os.path.exists("examples"):
+                example_list = [["examples/" + fname] for fname in os.listdir("examples") if fname.endswith((".jpg", ".png"))]
+            gr.Examples(examples=example_list, inputs=[image_input])
 
         with gr.Column():
-            label = gr.Label(num_top_classes=4, label="Predictions")
-            pred_time = gr.Number(label="Prediction Time (s)")
+            label_output = gr.Label(num_top_classes=4, label="Predictions")
+            time_output = gr.Number(label="Prediction Time (seconds)")
+
+    infer_btn.click(fn=predict, inputs=image_input, outputs=[label_output, time_output])
 
-        infer.click(fn=predict, inputs=[image], outputs=[label, pred_time])
+gradio_app.launch()
 
-gradio_app.launch(debug=True, share=True)