Update app.py

app.py

@@ -2,13 +2,40 @@ import gradio as gr
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from torchvision import transforms
+from torchvision import transforms, models
 from PIL import Image
 import numpy as np
 
 # Define your 4 classes
 CLASS_NAMES = ['Cover Drive', 'Pull Shot', 'Cut Shot', 'Straight Drive']  # Update with your actual class names
 
+# VGG16 Fine-tuned Model Definition
+class VGG16FineTuned(nn.Module):
+    def __init__(self, num_classes=4):
+        super(VGG16FineTuned, self).__init__()
+        # Load pre-trained VGG16 features
+        vgg16 = models.vgg16(pretrained=False)
+        self.features = vgg16.features
+        self.avgpool = vgg16.avgpool
+        
+        # Custom classifier to match your architecture
+        self.classifier = nn.Sequential(
+            nn.Linear(25088, 1024),
+            nn.ReLU(),
+            nn.Dropout(p=0.5),
+            nn.Linear(1024, 512),
+            nn.ReLU(),
+            nn.Dropout(p=0.5),
+            nn.Linear(512, num_classes)
+        )
+    
+    def forward(self, x):
+        x = self.features(x)
+        x = self.avgpool(x)
+        x = torch.flatten(x, 1)
+        x = self.classifier(x)
+        return x
+
 # Custom CNN Model Definition
 class CricketShotCNN(nn.Module):
     def __init__(self, num_classes=4):
@@ -66,7 +93,9 @@ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 def load_models():
     try:
         # Load VGG16 fine-tuned model
-        vgg16_model = torch.load('vgg16_finetuned.pth', map_location=device)
+        vgg16_model = VGG16FineTuned(num_classes=4)
+        vgg16_model.load_state_dict(torch.load('vgg16_finetuned.pth', map_location=device))
+        vgg16_model.to(device)
         vgg16_model.eval()
         
         # Load Custom CNN model