Update app.py

app.py

@@ -1,17 +1,80 @@
 import gradio as gr
-import tensorflow as tf
-from tensorflow import keras
-import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision import transforms
 from PIL import Image
+import numpy as np
 
-class_names = ['drive', 'legglance_flick', 'pullshot', 'sweep']
+# Define your 4 classes
+CLASS_NAMES = ['Cover Drive', 'Pull Shot', 'Cut Shot', 'Straight Drive']  # Update with your actual class names
 
+# Custom CNN Model Definition
+class CricketShotCNN(nn.Module):
+    def __init__(self, num_classes=4):
+        super(CricketShotCNN, self).__init__()
+        
+        # Block 1: Input (3, 224, 224) -> Output (64, 112, 112)
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)
+        self.bn1 = nn.BatchNorm2d(64)
+        
+        # Block 2: Output (128, 56, 56)
+        self.conv2 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
+        self.bn2 = nn.BatchNorm2d(128)
+        
+        # Block 3: Output (256, 28, 28)
+        self.conv3 = nn.Conv2d(128, 256, kernel_size=3, padding=1)
+        self.bn3 = nn.BatchNorm2d(256)
+        
+        # Block 4: Output (512, 14, 14)
+        self.conv4 = nn.Conv2d(256, 512, kernel_size=3, padding=1)
+        self.bn4 = nn.BatchNorm2d(512)
+        
+        self.pool = nn.MaxPool2d(2, 2)
+        self.dropout = nn.Dropout(0.5)
+        
+        # Fully Connected Layers
+        self.fc1 = nn.Linear(512 * 14 * 14, 512)
+        self.fc2 = nn.Linear(512, 128)
+        self.fc3 = nn.Linear(128, num_classes)
+
+    def forward(self, x):
+        x = self.pool(F.relu(self.bn1(self.conv1(x))))
+        x = self.pool(F.relu(self.bn2(self.conv2(x))))
+        x = self.pool(F.relu(self.bn3(self.conv3(x))))
+        x = self.pool(F.relu(self.bn4(self.conv4(x))))
+        
+        x = x.view(-1, 512 * 14 * 14)
+        
+        x = F.relu(self.fc1(x))
+        x = self.dropout(x)
+        x = F.relu(self.fc2(x))
+        x = self.fc3(x)
+        
+        return x
+
+# Image 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])
+])
 
 # Load models
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
 def load_models():
     try:
-        vgg16_model = keras.models.load_model('vgg16_finetuned.pth')
-        custom_cnn_model = keras.models.load_model('cricket_model.pth')
+        # Load VGG16 fine-tuned model
+        vgg16_model = torch.load('vgg16_finetuned.pth', map_location=device)
+        vgg16_model.eval()
+        
+        # Load Custom CNN model
+        custom_cnn_model = CricketShotCNN(num_classes=4)
+        custom_cnn_model.load_state_dict(torch.load('custom_cnn.pth', map_location=device))
+        custom_cnn_model.to(device)
+        custom_cnn_model.eval()
+        
         return vgg16_model, custom_cnn_model
     except Exception as e:
         print(f"Error loading models: {e}")
@@ -19,33 +82,6 @@ def load_models():
 
 vgg16_model, custom_cnn_model = load_models()
 
-def preprocess_image(image, target_size=(224, 224)):
-    """Preprocess image for model prediction"""
-    if image is None:
-        return None
-    
-    # Convert to PIL Image if needed
-    if not isinstance(image, Image.Image):
-        image = Image.fromarray(image)
-    
-    # Resize image
-    image = image.resize(target_size)
-    
-    # Convert to array and normalize
-    img_array = np.array(image)
-    
-    # Handle grayscale images
-    if len(img_array.shape) == 2:
-        img_array = np.stack([img_array] * 3, axis=-1)
-    
-    # Add batch dimension
-    img_array = np.expand_dims(img_array, axis=0)
-    
-    # Normalize to [0, 1]
-    img_array = img_array.astype('float32') / 255.0
-    
-    return img_array
-
 def predict(image):
     """Make predictions with both models"""
     if image is None:
@@ -54,56 +90,69 @@ def predict(image):
     if vgg16_model is None or custom_cnn_model is None:
         return "Models not loaded properly", "Models not loaded properly"
     
-    # Preprocess image
-    processed_img = preprocess_image(image)
-    
-    # Get predictions from both models
-    vgg16_pred = vgg16_model.predict(processed_img, verbose=0)[0]
-    custom_cnn_pred = custom_cnn_model.predict(processed_img, verbose=0)[0]
-    
-    # Create confidence dictionaries
-    vgg16_confidence = {CLASS_NAMES[i]: float(vgg16_pred[i]) for i in range(len(CLASS_NAMES))}
-    custom_cnn_confidence = {CLASS_NAMES[i]: float(custom_cnn_pred[i]) for i in range(len(CLASS_NAMES))}
+    try:
+        # Convert numpy array to PIL Image
+        if isinstance(image, np.ndarray):
+            image = Image.fromarray(image.astype('uint8'), 'RGB')
+        
+        # Preprocess image
+        img_tensor = transform(image).unsqueeze(0).to(device)
+        
+        # Get predictions from both models
+        with torch.no_grad():
+            vgg16_output = vgg16_model(img_tensor)
+            custom_cnn_output = custom_cnn_model(img_tensor)
+            
+            # Apply softmax to get probabilities
+            vgg16_probs = F.softmax(vgg16_output, dim=1)[0]
+            custom_cnn_probs = F.softmax(custom_cnn_output, dim=1)[0]
+        
+        # Create confidence dictionaries
+        vgg16_confidence = {CLASS_NAMES[i]: float(vgg16_probs[i]) for i in range(len(CLASS_NAMES))}
+        custom_cnn_confidence = {CLASS_NAMES[i]: float(custom_cnn_probs[i]) for i in range(len(CLASS_NAMES))}
+        
+        return vgg16_confidence, custom_cnn_confidence
     
-    return vgg16_confidence, custom_cnn_confidence
+    except Exception as e:
+        print(f"Prediction error: {e}")
+        return f"Error: {str(e)}", f"Error: {str(e)}"
 
 # Create Gradio interface
-with gr.Blocks(title="Dual Model Comparison") as demo:
+with gr.Blocks(title="Cricket Shot Classification - Dual Model Comparison", theme=gr.themes.Soft()) as demo:
     gr.Markdown(
         """
-        # 🔍 Dual Model Image Classification
+        # 🏏 Cricket Shot Classification - Dual Model Comparison
         
-        Compare predictions from two models trained on the same dataset:
+        Compare predictions from two models trained on the same cricket shot dataset:
         - **VGG16 Fine-tuned**: Transfer learning model based on VGG16
         - **Custom CNN**: CNN trained from scratch
         
-        Upload an image to see predictions and confidence scores from both models.
+        Upload an image of a cricket shot to see predictions and confidence scores from both models.
         """
     )
     
     with gr.Row():
         with gr.Column():
-            input_image = gr.Image(label="Upload Image", type="numpy")
-            predict_btn = gr.Button("Predict", variant="primary")
+            input_image = gr.Image(label="Upload Cricket Shot Image", type="numpy")
+            predict_btn = gr.Button("🔍 Predict", variant="primary", size="lg")
     
     with gr.Row():
         with gr.Column():
-            gr.Markdown("### VGG16 Fine-tuned Model")
+            gr.Markdown("### 📊 VGG16 Fine-tuned Model")
             vgg16_output = gr.Label(label="Predictions", num_top_classes=4)
         
         with gr.Column():
-            gr.Markdown("### Custom CNN Model")
+            gr.Markdown("### 📊 Custom CNN Model")
             custom_cnn_output = gr.Label(label="Predictions", num_top_classes=4)
     
-    # Examples section (optional - add your example images)
-    gr.Markdown("### Examples")
-    gr.Examples(
-        examples=[
-            # Add paths to example images here
-            # ["example1.jpg"],
-            # ["example2.jpg"],
-        ],
-        inputs=input_image,
+    gr.Markdown(
+        """
+        ---
+        ### 📝 About the Models
+        - Both models are trained on the same cricket shot dataset with 4 classes
+        - Input image size: 224x224 pixels
+        - The predictions show probability scores for each cricket shot type
+        """
     )
     
     # Connect the prediction function