updated one

app.py

@@ -1,48 +1,47 @@
 import gradio as gr
 import os
 import torch
+import torch.nn as nn
 import torchvision.transforms as transforms
 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. Class names (must match training order!)
+# 1. Class names (order must match your training)
 class_names = ['Cabbage', 'Carrot', 'Cucumber', 'Tomato']
 
-# 2. Image preprocessing pipeline from ResNet18 weights
+# 2. Load default preprocessing from pretrained weights
 weights = ResNet18_Weights.DEFAULT
-manual_transforms = weights.transforms()
+transform = weights.transforms()
 
-# 3. Load model with correct architecture
+# 3. Define and load the model
 def get_resnet18():
-    model = models.resnet18(weights=weights)  # ← NEW: use correct weight object
+    model = models.resnet18(weights=weights)
     num_ftrs = model.fc.in_features
-    model.fc = nn.Linear(num_ftrs, len(class_names))  # ← match training config
+    model.fc = nn.Linear(num_ftrs, len(class_names))
     return model
 
-resnet = get_resnet18()
-resnet.load_state_dict(torch.load("resnet_model.pth", map_location=torch.device("cpu")))
-resnet.eval()
+model = get_resnet18()
+model.load_state_dict(torch.load("resnet_model.pth", map_location=torch.device("cpu"), weights_only=False))
+model.eval()
 
 # 4. Prediction function
 def predict(img: Image.Image) -> Tuple[Dict[str, float], float]:
-    start_time = timer()
-    img_tensor = manual_transforms(img).unsqueeze(0)
+    start = timer()
+    img_tensor = transform(img).unsqueeze(0)
 
     with torch.inference_mode():
-        outputs = resnet(img_tensor)
-        probs = torch.softmax(outputs, dim=1)[0]
-
-    pred_labels_and_probs = {class_names[i]: float(probs[i]) for i in range(len(class_names))}
-    pred_time = round(timer() - start_time, 5)
+        output = model(img_tensor)
+        probs = torch.softmax(output, dim=1)[0]
 
-    return pred_labels_and_probs, pred_time
+    result = {class_names[i]: float(probs[i]) for i in range(len(class_names))}
+    elapsed = round(timer() - start, 5)
+    return result, elapsed
 
-# 5. Gradio app
-with gr.Blocks() as gradio_app:
+# 5. Build the Gradio interface
+with gr.Blocks() as demo:
     gr.HTML("<h1 style='text-align: center'>Vegetable Classification with ResNet18</h1>")
     gr.HTML("""
         <h3 style='text-align: center'>
@@ -53,18 +52,19 @@ with gr.Blocks() as gradio_app:
     with gr.Row():
         with gr.Column():
             image_input = gr.Image(type="pil", label="Upload Image")
-            infer_btn = gr.Button("Predict")
+            predict_button = gr.Button("Predict")
 
-            example_list = []
+            examples = []
             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])
+                examples = [["examples/" + fname] for fname in os.listdir("examples") if fname.lower().endswith((".png", ".jpg", ".jpeg"))]
+            gr.Examples(examples=examples, inputs=[image_input])
 
         with gr.Column():
-            label_output = gr.Label(num_top_classes=4, label="Predictions")
-            time_output = gr.Number(label="Prediction Time (seconds)")
+            output_label = gr.Label(num_top_classes=4, label="Predictions")
+            output_time = gr.Number(label="Prediction Time (seconds)")
 
-    infer_btn.click(fn=predict, inputs=image_input, outputs=[label_output, time_output])
+    predict_button.click(fn=predict, inputs=image_input, outputs=[output_label, output_time])
 
-gradio_app.launch()
+# 6. Launch the app
+demo.launch()