Update app.py

app.py

@@ -1,12 +1,11 @@
 import gradio as gr
 import torch
-from torchvision import transforms
-from PIL import Image
 import torch.nn as nn
+from torchvision import models, transforms
+from PIL import Image
 import os
-from torchvision import models
 
-# Custom Residual Block
+# Define the same custom residual block and EfficientNetWithNovelty model
 class ResidualBlock(nn.Module):
     def __init__(self, in_channels, out_channels):
         super(ResidualBlock, self).__init__()
@@ -28,7 +27,6 @@ class ResidualBlock(nn.Module):
         x = self.relu(x)
         return x
 
-# EfficientNet Model with Novelty (Residual Block)
 class EfficientNetWithNovelty(nn.Module):
     def __init__(self, num_classes):
         super(EfficientNetWithNovelty, self).__init__()
@@ -55,49 +53,58 @@ class EfficientNetWithNovelty(nn.Module):
         
         return x
 
-# Load the model and weights
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-# Update this path with your model path
-model_path = 'best_model.pth'
-num_classes = 10  # Update based on your dataset
-
+# Load the model checkpoint
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+num_classes = 10  # Number of classes as per your dataset
 model = EfficientNetWithNovelty(num_classes)
-
-# Load checkpoint instead of raw state_dict
-checkpoint = torch.load(model_path, map_location=device)
-model.load_state_dict(checkpoint["model_state_dict"])  # Correct key for model weights
+checkpoint = torch.load('best_model2.pth')
+model.load_state_dict(checkpoint['model_state_dict'])
 model.to(device)
 model.eval()
 
-# Define image transformations (same as during training)
+# Define image transformations for preprocessing
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor(),
     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
 ])
 
-# Define the prediction function for Gradio
+# Define the class labels explicitly
+class_labels = [
+    "KNUCKLE",
+    "LEGSPIN",
+    "OFFSPIN",
+    "OUTSWING",
+    "STRAIGHT",
+    "BACK_OF_HAND",
+    "CARROM",
+    "CROSSSEAM",
+    "GOOGLY",
+    "INSWING"
+]
+
+# Prediction function
 def predict(image):
-    image = Image.fromarray(image)  # Convert numpy array to PIL Image
-    image = transform(image).unsqueeze(0)  # Apply transformations and add batch dimension
-    image = image.to(device)
-
+    # Preprocess image
+    image = Image.fromarray(image)  # Convert numpy array to PIL Image if it's from Gradio
+    image = transform(image).unsqueeze(0).to(device)
+    
+    # Get model predictions
     with torch.no_grad():
-        outputs = model(image)
-        _, predicted = torch.max(outputs, 1)
-
-    # Class names for your classification
-    class_names = ['OUTSWING', 'STRAIGHT', 'BACK_OF_HAND', 'CARROM', 'CROSSSEAM', 
-                   'GOOGLY', 'INSWING', 'KNUCKLE', 'LEGSPIN', 'OFFSPIN']
-    predicted_label = class_names[predicted.item()]
+        output = model(image)
+        _, predicted = torch.max(output, 1)
+    
+    # Get predicted class label
+    predicted_label = class_labels[predicted.item()]
     return predicted_label
 
-# Create the Gradio Interface
-iface = gr.Interface(fn=predict,
-                     inputs=gr.Image(type="numpy"),  # Accepts image input
-                     outputs=gr.Text(),  # Output the predicted class label
-                     live=True)  # live=True enables prediction while image is being uploaded
+# Gradio interface
+iface = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="numpy", label="Upload Cricket Grip Image"),
+    outputs=gr.Textbox(label="Predicted Grip Type"),
+    live=True
+)
 
-# Launch the interface
-iface.launch()
+if __name__ == "__main__":
+    iface.launch()