Update app.py

app.py

@@ -13,7 +13,7 @@ import os
 print("🚀 Loading Ugandan Sign Language Model...")
 
 # ============================================================================
-# MODEL SETUP
+# MODEL SETUP - FIXED VERSION
 # ============================================================================
 
 class MinimalClassifier(nn.Module):
@@ -32,24 +32,66 @@ processor = XCLIPProcessor.from_pretrained("microsoft/xclip-base-patch32")
 xclip_model = XCLIPModel.from_pretrained("microsoft/xclip-base-patch32").to(device)
 xclip_model.eval()
 
-# Load your trained model
+# Load your trained model - WITH ERROR HANDLING
 try:
     checkpoint = torch.load("best_xclip_model.pth", map_location=device, weights_only=False)
+    
+    # DEBUG: Check what's in the checkpoint
+    print(f"🔍 Checkpoint keys: {list(checkpoint.keys())}")
+    
+    # FIX: Handle missing 'num_classes' key
+    if 'num_classes' in checkpoint:
+        num_classes = checkpoint['num_classes']
+    else:
+        # Try to infer number of classes
+        if 'id_to_label' in checkpoint:
+            num_classes = len(checkpoint['id_to_label'])
+        elif 'label_to_id' in checkpoint:
+            num_classes = len(checkpoint['label_to_id'])
+        else:
+            # Count from model weights
+            for key in checkpoint.keys():
+                if 'model_state_dict' in checkpoint:
+                    weight_key = [k for k in checkpoint['model_state_dict'].keys() if 'classifier' in k and 'weight' in k]
+                    if weight_key:
+                        num_classes = checkpoint['model_state_dict'][weight_key[0]].shape[0]
+                        break
+            else:
+                num_classes = 85  # Default fallback
+    
+    print(f"✅ Using num_classes: {num_classes}")
+    
+    # Initialize model
     model = MinimalClassifier(
         input_dim=512,
-        num_classes=checkpoint['num_classes'],
+        num_classes=num_classes,
         dropout=0.5
     ).to(device)
-    model.load_state_dict(checkpoint['model_state_dict'])
-    model.eval()
     
-    id_to_label = checkpoint['id_to_label']
+    # Load state dict
+    if 'model_state_dict' in checkpoint:
+        model.load_state_dict(checkpoint['model_state_dict'])
+    else:
+        # If checkpoint IS the state dict
+        model.load_state_dict(checkpoint)
+    
+    # Load label mappings
+    if 'id_to_label' in checkpoint:
+        id_to_label = checkpoint['id_to_label']
+    else:
+        # Create default mapping
+        id_to_label = {i: f"class_{i}" for i in range(num_classes)}
+        print("⚠️ Created default label mapping")
+    
     label_to_id = {v: k for k, v in id_to_label.items()}
     
+    model.eval()
     print(f"✅ Model loaded! Can recognize {len(id_to_label)} signs")
+    print(f"📊 Sample classes: {list(id_to_label.values())[:5]}")
     
 except Exception as e:
     print(f"❌ Error loading model: {e}")
+    print("💡 TIP: Make sure your model file has 'num_classes' or 'id_to_label' key")
     exit(1)
 
 # ============================================================================
@@ -121,6 +163,7 @@ def predict_sign(video_path):
         return predicted_label, confidence_value
         
     except Exception as e:
+        print(f"❌ Prediction error: {e}")
         return "Unknown", 0.0
 
 # ============================================================================