Update app.py

app.py

@@ -90,23 +90,48 @@ transform = transforms.Compose([
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 def load_models():
+    vgg16_model = None
+    custom_cnn_model = None
+    error_messages = []
+    
     try:
         # Load VGG16 fine-tuned model
+        print("Loading VGG16 model...")
         vgg16_model = VGG16FineTuned(num_classes=4)
-        vgg16_model.load_state_dict(torch.load('vgg16_finetuned.pth', map_location=device))
+        vgg16_state = torch.load('vgg16_finetuned.pth', map_location=device, weights_only=False)
+        vgg16_model.load_state_dict(vgg16_state)
         vgg16_model.to(device)
         vgg16_model.eval()
-        
+        print("✓ VGG16 model loaded successfully")
+    except FileNotFoundError:
+        error_messages.append("VGG16: File 'vgg16_finetuned.pth' not found")
+        print("✗ VGG16 model file not found")
+    except Exception as e:
+        error_messages.append(f"VGG16: {str(e)}")
+        print(f"✗ VGG16 loading error: {e}")
+    
+    try:
         # Load Custom CNN model
+        print("Loading Custom CNN model...")
         custom_cnn_model = CricketShotCNN(num_classes=4)
-        custom_cnn_model.load_state_dict(torch.load('cricket_model.pth', map_location=device))
+        custom_cnn_state = torch.load('custom_cnn.pth', map_location=device, weights_only=False)
+        custom_cnn_model.load_state_dict(custom_cnn_state)
         custom_cnn_model.to(device)
         custom_cnn_model.eval()
-        
-        return vgg16_model, custom_cnn_model
+        print("✓ Custom CNN model loaded successfully")
+    except FileNotFoundError:
+        error_messages.append("Custom CNN: File 'custom_cnn.pth' not found")
+        print("✗ Custom CNN model file not found")
     except Exception as e:
-        print(f"Error loading models: {e}")
-        return None, None
+        error_messages.append(f"Custom CNN: {str(e)}")
+        print(f"✗ Custom CNN loading error: {e}")
+    
+    if error_messages:
+        print("\n⚠️ Model Loading Errors:")
+        for msg in error_messages:
+            print(f"  - {msg}")
+    
+    return vgg16_model, custom_cnn_model
 
 vgg16_model, custom_cnn_model = load_models()
 
@@ -118,6 +143,9 @@ def predict(image):
     if vgg16_model is None or custom_cnn_model is None:
         return "Models not loaded properly", "Models not loaded properly"
     
+    # Define class names here to ensure they're in scope
+    class_names = ['Cover Drive', 'Pull Shot', 'Cut Shot', 'Straight Drive']
+    
     try:
         # Convert numpy array to PIL Image
         if isinstance(image, np.ndarray):
@@ -136,8 +164,8 @@ def predict(image):
             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))}
+        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