Update app.py

app.py

@@ -6,8 +6,7 @@ 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
+class_names = ['drive', 'legglance_flick', 'pullshot', 'sweep']
 
 # VGG16 Fine-tuned Model Definition
 class VGG16FineTuned(nn.Module):
@@ -100,7 +99,7 @@ def load_models():
         
         # 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.load_state_dict(torch.load('cricket_model.pth', map_location=device))
         custom_cnn_model.to(device)
         custom_cnn_model.eval()