Update streamlit_app.py

streamlit_app.py

@@ -43,26 +43,48 @@ def load_model():
     # Extract the actual model state dict
     if 'model_state_dict' in checkpoint:
         model_state_dict = checkpoint['model_state_dict']
+        # Get class names from checkpoint if available
+        if 'class_names' in checkpoint:
+            saved_class_names = checkpoint['class_names']
+            print(f"Loaded class names from checkpoint: {len(saved_class_names)} classes")
     else:
         # If it's just the state dict directly
         model_state_dict = checkpoint
+        saved_class_names = class_names
     
     # Get the number of classes from the model weights
     if 'classifier.weight' in model_state_dict:
         num_classes_in_model = model_state_dict['classifier.weight'].shape[0]
+    elif 'head.weight' in model_state_dict:  # Alternative naming in some timm versions
+        num_classes_in_model = model_state_dict['head.weight'].shape[0]
     elif 'fc.weight' in model_state_dict:
         num_classes_in_model = model_state_dict['fc.weight'].shape[0]
     else:
         # Fallback: assume it matches class_names
         num_classes_in_model = len(class_names)
     
-    # Create model with the correct number of classes
-    model = timm.create_model('efficientnet_b0', pretrained=False, num_classes=num_classes_in_model)
+    print(f"Creating model with {num_classes_in_model} classes")
     
-    # Load the model state dict (not the entire checkpoint)
-    model.load_state_dict(model_state_dict)
-    model.eval()
+    # Create model exactly as in training - with dropout and drop_path
+    model = timm.create_model(
+        'efficientnet_b0',
+        pretrained=False,  # Don't load pretrained weights
+        num_classes=num_classes_in_model,
+        drop_rate=0.3,      # Match training parameters
+        drop_path_rate=0.2  # Match training parameters
+    )
+    
+    # Load the model state dict
+    try:
+        model.load_state_dict(model_state_dict, strict=True)
+        print("Model loaded successfully!")
+    except RuntimeError as e:
+        print(f"Error loading model: {e}")
+        # Try with strict=False as fallback
+        model.load_state_dict(model_state_dict, strict=False)
+        print("Model loaded with some missing/unexpected keys")
     
+    model.eval()
     return model
 
 # Load the model
@@ -71,9 +93,11 @@ model = load_model()
 if model is None:
     st.stop()
 
+# Use the exact same transforms as in training validation
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
 ])
 
 def predict_butterfly(image):
@@ -85,6 +109,10 @@ def predict_butterfly(image):
     if isinstance(image, np.ndarray):
         image = Image.fromarray(image)
     
+    # Ensure RGB format
+    if image.mode != 'RGB':
+        image = image.convert('RGB')
+    
     # Preprocess
     input_tensor = transform(image).unsqueeze(0)