Update streamlit_app.py

streamlit_app.py

@@ -47,7 +47,7 @@ def detect_model_architecture(model_state_dict):
     if 'classifier.weight' in model_state_dict:
         classifier_input_features = model_state_dict['classifier.weight'].shape[1]
         
-        # EfficientNet feature mapping (corrected based on actual architectures)
+        # EfficientNet feature mapping (updated based on actual timm implementations)
         efficientnet_features = {
             1280: 'efficientnet_b0',
             1408: 'efficientnet_b1', 
@@ -61,7 +61,7 @@ def detect_model_architecture(model_state_dict):
         
         model_name = efficientnet_features.get(classifier_input_features, 'efficientnet_b3')
         print(f"Detected model architecture: {model_name} (classifier features: {classifier_input_features})")
-        return model_name
+        return model_name, classifier_input_features
     
     # Check bn2 layer (final batch norm before classifier)
     if 'bn2.weight' in model_state_dict:
@@ -82,11 +82,11 @@ def detect_model_architecture(model_state_dict):
         
         model_name = bn2_mapping.get(bn2_features, 'efficientnet_b3')
         print(f"Detected model architecture from bn2: {model_name}")
-        return model_name
+        return model_name, bn2_features
     
     # Default to B3 based on your error logs
     print("Could not detect model architecture, defaulting to efficientnet_b3")
-    return 'efficientnet_b3'
+    return 'efficientnet_b3', 1792
 
 @st.cache_resource
 def load_model():
@@ -105,6 +105,10 @@ def load_model():
         
         The model file appears to be a Git LFS pointer file (size: {file_size} bytes).
         This means the actual model wasn't downloaded properly.
+        
+        **To fix this:**
+        1. Run: `git lfs pull` in your repository
+        2. Or download the model file directly from your storage
         """)
         return None
     
@@ -128,7 +132,7 @@ def load_model():
                 print(f"  {key}: {model_state_dict[key].shape}")
         
         # Auto-detect the correct model architecture
-        model_name = detect_model_architecture(model_state_dict)
+        model_name, expected_features = detect_model_architecture(model_state_dict)
         
         # Get number of classes
         num_classes = len(class_names)
@@ -137,32 +141,56 @@ def load_model():
         
         print(f"Loading {model_name} with {num_classes} classes")
         
-        # Create model with correct architecture and matching parameters
-        # Use the exact same parameters as in training
-        model = timm.create_model(
-            model_name,
-            pretrained=False,  # Don't load pretrained weights
-            num_classes=num_classes,
-            drop_rate=0.4,  # Match training parameters
-            drop_path_rate=0.3
-        )
+        # Create model with correct architecture
+        # Try different parameter combinations that might have been used during training
+        model_configs = [
+            # Most likely configuration based on your checkpoint
+            {'drop_rate': 0.4, 'drop_path_rate': 0.3},
+            {'drop_rate': 0.3, 'drop_path_rate': 0.2},
+            {'drop_rate': 0.2, 'drop_path_rate': 0.1},
+            {'drop_rate': 0.0, 'drop_path_rate': 0.0},  # Default
+        ]
         
-        # Load the trained weights
-        try:
-            model.load_state_dict(model_state_dict, strict=True)
-            print("Model loaded successfully with strict=True")
-        except RuntimeError as e:
-            print(f"Strict loading failed: {e}")
-            # Try with strict=False
+        model = None
+        for config in model_configs:
+            try:
+                print(f"Trying model config: {config}")
+                model = timm.create_model(
+                    model_name,
+                    pretrained=False,
+                    num_classes=num_classes,
+                    **config
+                )
+                
+                # Try loading with strict=True first
+                model.load_state_dict(model_state_dict, strict=True)
+                print(f"Model loaded successfully with config: {config}")
+                break
+                
+            except RuntimeError as e:
+                print(f"Config {config} failed: {e}")
+                continue
+        
+        # If strict loading failed for all configs, try with strict=False
+        if model is None:
+            print("All strict loading attempts failed, trying with strict=False")
+            model = timm.create_model(
+                model_name,
+                pretrained=False,
+                num_classes=num_classes,
+                drop_rate=0.4,
+                drop_path_rate=0.3
+            )
+            
             missing_keys, unexpected_keys = model.load_state_dict(model_state_dict, strict=False)
             
             if missing_keys:
                 print(f"Missing keys: {missing_keys}")
-                st.warning(f"Some model weights were not loaded: {len(missing_keys)} missing keys")
+                st.warning(f"⚠️ Some model weights were not loaded: {len(missing_keys)} missing keys")
             
             if unexpected_keys:
                 print(f"Unexpected keys: {unexpected_keys}")
-                st.warning(f"Some checkpoint keys were not used: {len(unexpected_keys)} unexpected keys")
+                st.warning(f"⚠️ Some checkpoint keys were not used: {len(unexpected_keys)} unexpected keys")
         
         # Verify the model loaded correctly
         model.eval()
@@ -173,6 +201,12 @@ def load_model():
             try:
                 dummy_output = model(dummy_input)
                 print(f"Model test successful. Output shape: {dummy_output.shape}")
+                
+                # Verify output shape matches expected classes
+                if dummy_output.shape[1] != num_classes:
+                    st.error(f"Model output mismatch: expected {num_classes} classes, got {dummy_output.shape[1]}")
+                    return None
+                    
             except Exception as e:
                 print(f"Model test failed: {e}")
                 st.error(f"Model validation failed: {e}")
@@ -190,14 +224,19 @@ def load_model():
             checkpoint = torch.load(MODEL_PATH, map_location='cpu')
             if 'model_state_dict' in checkpoint:
                 model_keys = list(checkpoint['model_state_dict'].keys())
-                print(f"Available keys in checkpoint: {model_keys[:10]}...")  # Show first 10 keys
+                print(f"Available keys in checkpoint: {model_keys[:10]}...")
                 
                 # Show the problematic layer shapes
                 state_dict = checkpoint['model_state_dict']
                 if 'classifier.weight' in state_dict and 'bn2.weight' in state_dict:
-                    print(f"Classifier input features: {state_dict['classifier.weight'].shape[1]}")
-                    print(f"bn2 features: {state_dict['bn2.weight'].shape[0]}")
+                    classifier_features = state_dict['classifier.weight'].shape[1]
+                    bn2_features = state_dict['bn2.weight'].shape[0]
+                    print(f"Classifier input features: {classifier_features}")
+                    print(f"bn2 features: {bn2_features}")
                     
+                    if classifier_features != bn2_features:
+                        st.error(f"Architecture mismatch: classifier expects {classifier_features} features, but bn2 has {bn2_features}")
+                        
         except Exception as debug_e:
             print(f"Debug info failed: {debug_e}")
         
@@ -207,8 +246,12 @@ def load_model():
 model = load_model()
 
 if model is None:
-    st.warning("⚠️ **Model Loading Failed**")
-    st.info("Please check the logs for detailed error information.")
+    st.error("⚠️ **Model Loading Failed**")
+    st.info("**Possible solutions:**")
+    st.markdown("1. **Git LFS issue**: Run `git lfs pull` to download the actual model file")
+    st.markdown("2. **Architecture mismatch**: The model was trained with different parameters")
+    st.markdown("3. **Corrupted file**: Re-download or re-train the model")
+    st.markdown("4. **Check the console/logs** for detailed error information")
     st.stop()
 
 # Transform for preprocessing (same as training)
@@ -241,11 +284,20 @@ def predict_butterfly(image):
             probabilities = torch.nn.functional.softmax(output[0], dim=0)
             confidence, pred = torch.max(probabilities, 0)
             
+            # Get top 3 predictions for better debugging
+            top_probs, top_indices = torch.topk(probabilities, min(3, len(probabilities)))
+            
             if pred.item() < len(class_names):
                 predicted_class = class_names[pred.item()]
             else:
                 predicted_class = f"Class_{pred.item()}"
             
+            # Debug info
+            print(f"Top predictions:")
+            for i, (prob, idx) in enumerate(zip(top_probs, top_indices)):
+                class_name = class_names[idx.item()] if idx.item() < len(class_names) else f"Class_{idx.item()}"
+                print(f"  {i+1}. {class_name}: {prob.item():.3f}")
+            
         return predicted_class, confidence.item()
         
     except Exception as e:
@@ -256,6 +308,10 @@ def predict_butterfly(image):
 st.title("🦋 Butterfly Identifier / Liblikamaja ID")
 st.write("Identify butterflies using your camera or by uploading an image!")
 
+# Show model info
+if model is not None:
+    st.info(f"📊 Model loaded: {len(class_names)} butterfly species recognized")
+
 # Create tabs for different input methods
 tab1, tab2 = st.tabs(["📷 Live Camera", "📁 Upload Image"])
 
@@ -275,20 +331,27 @@ with tab1:
                 st.image(image, caption="Captured Image", use_column_width=True)
             
             with col2:
-                predicted_class, confidence = predict_butterfly(image)
+                with st.spinner("Analyzing image..."):
+                    predicted_class, confidence = predict_butterfly(image)
                 
-                if predicted_class and confidence and confidence >= 0.80:
-                    st.success(f"**Prediction: {predicted_class}**")
-                    st.info(f"Confidence: {confidence:.2%}")
+                if predicted_class and confidence:
+                    if confidence >= 0.80:
+                        st.success(f"**Prediction: {predicted_class}**")
+                        st.info(f"Confidence: {confidence:.2%}")
+                        
+                        # Show additional info if available
+                        if predicted_class in butterfly_info:
+                            st.write(f"**About:** {butterfly_info[predicted_class]['description']}")
+                    else:
+                        st.warning("⚠️ **Low confidence prediction**")
+                        st.info(f"Best guess: {predicted_class} ({confidence:.1%})")
+                        st.markdown("**Tips for better results:**")
+                        st.markdown("- Use better lighting")
+                        st.markdown("- Get closer to the butterfly")
+                        st.markdown("- Ensure the butterfly is clearly visible")
+                        st.markdown("- Avoid blurry or dark images")
                 else:
-                    st.warning("⚠️ **Image not clear - Unable to identify butterfly**")
-                    if confidence:
-                        st.info(f"Confidence too low: {confidence:.1%}")
-                    st.markdown("**Tips for better results:**")
-                    st.markdown("- Use better lighting")
-                    st.markdown("- Get closer to the butterfly")
-                    st.markdown("- Ensure the butterfly is clearly visible")
-                    st.markdown("- Avoid blurry or dark images")
+                    st.error("Unable to analyze image. Please try again.")
                                     
         except Exception as e:
             st.error(f"Error processing image: {str(e)}")
@@ -314,20 +377,27 @@ with tab2:
                 st.image(image, caption="Uploaded Image", use_column_width=True)
             
             with col2:
-                predicted_class, confidence = predict_butterfly(image)
+                with st.spinner("Analyzing image..."):
+                    predicted_class, confidence = predict_butterfly(image)
                 
-                if predicted_class and confidence and confidence >= 0.80:
-                    st.success(f"**Prediction: {predicted_class}**")
-                    st.info(f"Confidence: {confidence:.2%}")
+                if predicted_class and confidence:
+                    if confidence >= 0.80:
+                        st.success(f"**Prediction: {predicted_class}**")
+                        st.info(f"Confidence: {confidence:.2%}")
+                        
+                        # Show additional info if available
+                        if predicted_class in butterfly_info:
+                            st.write(f"**About:** {butterfly_info[predicted_class]['description']}")
+                    else:
+                        st.warning("⚠️ **Low confidence prediction**")
+                        st.info(f"Best guess: {predicted_class} ({confidence:.1%})")
+                        st.markdown("**Tips for better results:**")
+                        st.markdown("- Use better lighting")
+                        st.markdown("- Get closer to the butterfly")
+                        st.markdown("- Ensure the butterfly is clearly visible")
+                        st.markdown("- Avoid blurry or dark images")
                 else:
-                    st.warning("⚠️ **Image not clear - Unable to identify butterfly**")
-                    if confidence:
-                        st.info(f"Confidence too low: {confidence:.1%}")
-                    st.markdown("**Tips for better results:**")
-                    st.markdown("- Use better lighting")
-                    st.markdown("- Get closer to the butterfly")
-                    st.markdown("- Ensure the butterfly is clearly visible")
-                    st.markdown("- Avoid blurry or dark images")
+                    st.error("Unable to analyze image. Please try again.")
                 
         except Exception as e:
             st.error(f"Error processing image: {str(e)}")
@@ -337,4 +407,14 @@ st.markdown("---")
 st.markdown("### How to use:")
 st.markdown("1. **Camera Capture**: Take a photo using your device camera")
 st.markdown("2. **Upload Image**: Choose a butterfly photo from your device")
-st.markdown("3. **Best Results**: Use clear, well-lit photos with the butterfly clearly visible")
+st.markdown("3. **Best Results**: Use clear, well-lit photos with the butterfly clearly visible")
+
+# Debug info (only show if there are issues)
+if st.checkbox("Show debug information"):
+    st.markdown("### Debug Information")
+    st.write(f"Number of classes: {len(class_names)}")
+    st.write(f"Model loaded: {model is not None}")
+    if model:
+        st.write("Model architecture successfully detected and loaded")
+    else:
+        st.write("❌ Model failed to load - check console for details")