Update streamlit_app.py

streamlit_app.py

@@ -37,10 +37,25 @@ def load_model():
         st.error("Model file not found. Please upload butterfly_classifier.pth to your space.")
         return None
     
-    # Use EfficientNet-B0 (same as training)
-    model = timm.create_model('efficientnet_b0', pretrained=False, num_classes=len(class_names))
-    model.load_state_dict(torch.load(MODEL_PATH, map_location="cpu"))
+    # Load the checkpoint first to check the actual number of classes
+    checkpoint = torch.load(MODEL_PATH, map_location="cpu")
+    
+    # Get the number of classes from the saved model weights
+    if 'classifier.weight' in checkpoint:
+        num_classes_in_model = checkpoint['classifier.weight'].shape[0]
+    elif 'fc.weight' in checkpoint:  # Alternative naming
+        num_classes_in_model = checkpoint['fc.weight'].shape[0]
+    else:
+        # Fallback: assume it matches class_names
+        num_classes_in_model = len(class_names)
+    
+    st.info(f"Model has {num_classes_in_model} classes, class_names.txt has {len(class_names)} classes")
+    
+    # Create model with the correct number of classes from the saved model
+    model = timm.create_model('efficientnet_b0', pretrained=False, num_classes=num_classes_in_model)
+    model.load_state_dict(checkpoint)
     model.eval()
+    
     return model
 
 # Load the model