Your commit message

app.py

@@ -1,90 +1,76 @@
-### 1. Imports and class names setup ###
 import gradio as gr
 import os
 import torch
-from torchvision import models, transforms
+import torchvision.transforms as transforms
+import torchvision.models.resnet
 from PIL import Image
 from timeit import default_timer as timer
 from typing import Tuple, Dict
 
-# Setup class names
+# 1. Setup class names
 class_names = ['Cabbage', 'Carrot', 'Cucumber', 'tomato']
 
-### 2. Model and transforms preparation ###
 
-# Load a pretrained ResNet model (e.g. ResNet18) and adjust final layer for 4 classes
-resnet = models.resnet18(pretrained=False)
-num_features = resnet.fc.in_features
-resnet.fc = torch.nn.Linear(num_features, len(class_names))
-
-# Load your trained weights
-resnet.load_state_dict(torch.load("resnet_model.pth", map_location=torch.device("cpu")))
-
-# Define the transforms to match ResNet input requirements
+# 2. Define transforms (adjust as needed, example here)
 manual_transforms = transforms.Compose([
-    transforms.Resize(256),
-    transforms.CenterCrop(224),
+    transforms.Resize((224, 224)),
     transforms.ToTensor(),
     transforms.Normalize(
-        mean=[0.485, 0.456, 0.406],  # Standard ImageNet mean
-        std=[0.229, 0.224, 0.225]    # Standard ImageNet std
+        mean=[0.485, 0.456, 0.406],  # ImageNet mean
+        std=[0.229, 0.224, 0.225]    # ImageNet std
     )
 ])
 
-### 3. Predict function ###
-def predict(img: Image.Image) -> Tuple[Dict, float]:
+# 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)
+)
+
+resnet.eval()
+
+# 4. Prediction function
+def predict(img: Image.Image) -> Tuple[Dict[str, float], float]:
     start_time = timer()
-    
-    img = manual_transforms(img).unsqueeze(0)  # Add batch dimension
-    
-    resnet.eval()
+    img_tensor = manual_transforms(img).unsqueeze(0)  # add batch dimension
+
     with torch.inference_mode():
-        pred_logits = resnet(img)
-        pred_probs = torch.softmax(pred_logits, dim=1)
-    
-    pred_labels_and_probs = {class_names[i]: float(pred_probs[0][i]) for i in range(len(class_names))}
-    
+        preds = torch.softmax(resnet(img_tensor), dim=1)
+
+    pred_labels_and_probs = {class_names[i]: float(preds[0][i]) for i in range(len(class_names))}
     pred_time = round(timer() - start_time, 5)
-    
-    return pred_labels_and_probs, pred_time
 
-### 4. Gradio app ###
+    return pred_labels_and_probs, pred_time
 
+# 5. Gradio app UI
 with gr.Blocks() as gradio_app:
-    gr.HTML(
-        """
-        <h1 style='text-align: center'>
-        Rice Disease Classification - ResNet Model
-        </h1>
-        """
-    )
-    
+    gr.HTML("<h1 style='text-align: center'>Rice Disease Classification - ResNet Model</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>  
+        <a href='https://www.linkedin.com/in/munzali-alhassan/' target='_blank'>Linkedin</a>
         </h3>
         """
     )
-    
+
     with gr.Row():
         with gr.Column():
             image = gr.Image(type="pil", label="Upload Image")
             infer = gr.Button(value="Predict")
-            
+
             example_list = [["examples/" + example] for example in os.listdir("examples")]
-            gr.Examples(
-                examples=example_list,
-                inputs=[image]
-            )
-        
+            gr.Examples(examples=example_list, inputs=[image])
+
         with gr.Column():
             label = gr.Label(num_top_classes=4, label="Predictions")
             pred_time = gr.Number(label="Prediction Time (s)")
-        
+
         infer.click(fn=predict, inputs=[image], outputs=[label, pred_time])
 
 gradio_app.launch(debug=True, share=True)
-