Update streamlit_app.py

streamlit_app.py

@@ -1,18 +1,20 @@
+# Full corrected bilingual Streamlit app for Butterfly Identifier
 import streamlit as st
 from PIL import Image
 import torch
-from torchvision import models, transforms
 import json
 import os
 import io
 import numpy as np
 import timm
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
 import warnings
 warnings.filterwarnings('ignore')
 
 # Configure Streamlit
 st.set_page_config(
-    page_title="Butterfly Identifier/ Liblikamaja ID",
+    page_title="Butterfly Identifier / Liblikamaja ID",
     page_icon="🦋",
     layout="wide"
 )
@@ -40,360 +42,121 @@ def load_butterfly_info():
 
 butterfly_info = load_butterfly_info()
 
-def detect_model_architecture(model_state_dict):
-    """Detect which EfficientNet variant was used based on layer shapes"""
-    
-    # Check the classifier input features to determine model type
-    if 'classifier.weight' in model_state_dict:
-        classifier_input_features = model_state_dict['classifier.weight'].shape[1]
-        
-        # EfficientNet feature mapping (updated based on actual timm implementations)
-        efficientnet_features = {
-            1280: 'efficientnet_b0',
-            1408: 'efficientnet_b1', 
-            1536: 'efficientnet_b2',
-            1792: 'efficientnet_b3',  # Your model has 1792 features
-            1920: 'efficientnet_b4',
-            2048: 'efficientnet_b5',
-            2304: 'efficientnet_b6',
-            2560: 'efficientnet_b7'
-        }
-        
-        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, classifier_input_features
-    
-    # Check bn2 layer (final batch norm before classifier)
-    if 'bn2.weight' in model_state_dict:
-        bn2_features = model_state_dict['bn2.weight'].shape[0]
-        print(f"bn2 features: {bn2_features}")
-        
-        # This should match the classifier input features
-        bn2_mapping = {
-            1280: 'efficientnet_b0',
-            1408: 'efficientnet_b1',
-            1536: 'efficientnet_b2', 
-            1792: 'efficientnet_b3',  # Your model shows 1792
-            1920: 'efficientnet_b4',
-            2048: 'efficientnet_b5',
-            2304: 'efficientnet_b6',
-            2560: 'efficientnet_b7'
-        }
-        
-        model_name = bn2_mapping.get(bn2_features, 'efficientnet_b3')
-        print(f"Detected model architecture from bn2: {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', 1792
+# Define transform matching training pipeline
+inference_transform = A.Compose([
+    A.Resize(224, 224),
+    A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+    ToTensorV2()
+])
 
+# Load the model
 @st.cache_resource
 def load_model():
     MODEL_PATH = "butterfly_classifier.pth"
-    
-    # Check if model file exists
     if not os.path.exists(MODEL_PATH):
-        st.error(f"Model file '{MODEL_PATH}' not found!")
-        return None
-    
-    # Check file size
-    file_size = os.path.getsize(MODEL_PATH)
-    if file_size < 1000:  # Less than 1KB suggests LFS pointer file
-        st.error(f"""
-        🚨 **Git LFS Issue Detected**
-        
-        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
-        """)
+        st.error("Model file not found!")
         return None
-    
     try:
-        # Load checkpoint
-        print(f"Loading model from {MODEL_PATH} (size: {file_size} bytes)")
         checkpoint = torch.load(MODEL_PATH, map_location='cpu')
-        
-        # Extract model components
-        if 'model_state_dict' in checkpoint:
-            model_state_dict = checkpoint['model_state_dict']
-            saved_class_names = checkpoint.get('class_names', class_names)
-        else:
-            model_state_dict = checkpoint
-            saved_class_names = class_names
-        
-        # Debug: Print some key layer shapes
-        print("Key layer shapes in checkpoint:")
-        for key in ['conv_stem.weight', 'bn1.weight', 'bn2.weight', 'classifier.weight']:
-            if key in model_state_dict:
-                print(f"  {key}: {model_state_dict[key].shape}")
-        
-        # Auto-detect the correct model architecture
-        model_name, expected_features = detect_model_architecture(model_state_dict)
-        
-        # Get number of classes
+        model_state_dict = checkpoint['model_state_dict'] if 'model_state_dict' in checkpoint else checkpoint
         num_classes = len(class_names)
-        if 'classifier.weight' in model_state_dict:
-            num_classes = model_state_dict['classifier.weight'].shape[0]
-        
-        print(f"Loading {model_name} with {num_classes} classes")
-        
-        # 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
-        ]
-        
-        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")
-            
-            if unexpected_keys:
-                print(f"Unexpected keys: {unexpected_keys}")
-                st.warning(f"⚠️ Some checkpoint keys were not used: {len(unexpected_keys)} unexpected keys")
-        
-        # Verify the model loaded correctly
+        model = timm.create_model('efficientnet_b3', pretrained=False, num_classes=num_classes, drop_rate=0.4, drop_path_rate=0.3)
+        model.load_state_dict(model_state_dict, strict=False)
         model.eval()
-        
-        # Test with a dummy input to make sure everything works
-        dummy_input = torch.randn(1, 3, 224, 224)
-        with torch.no_grad():
-            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}")
-                return None
-        
-        #st.success(f"✅ Successfully loaded {model_name} with {num_classes} classes")
         return model
-        
     except Exception as e:
         st.error(f"Error loading model: {str(e)}")
-        st.error(f"Model file size: {file_size} bytes")
-        
-        # Additional debugging info
-        try:
-            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 the problematic layer shapes
-                state_dict = checkpoint['model_state_dict']
-                if 'classifier.weight' in state_dict and 'bn2.weight' in state_dict:
-                    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}")
-        
         return None
 
-# Load the model
 model = load_model()
 
-if model is None:
-    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)
-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, threshold=0.5):
-    """Predict butterfly species from image"""
     try:
         if image is None:
             return None, None
-
         if isinstance(image, np.ndarray):
             image = Image.fromarray(image)
         if image.mode != 'RGB':
             image = image.convert('RGB')
-
-        input_tensor = transform(image).unsqueeze(0)
-
+        transformed = inference_transform(image=np.array(image))
+        input_tensor = transformed['image'].unsqueeze(0)
         with torch.no_grad():
             output = model(input_tensor)
             probabilities = torch.nn.functional.softmax(output[0], dim=0)
-
             confidence, pred = torch.max(probabilities, 0)
             if confidence.item() < threshold:
                 return None, confidence.item()
-
             predicted_class = class_names[pred.item()]
             return predicted_class, confidence.item()
-    
     except Exception as e:
         st.error(f"Prediction error: {str(e)}")
         return None, None
 
-
 # UI Code
-st.title("🦋  Liblikamaja ID/ Butterfly Identifier")
-st.write("Tuvasta liblikaid oma kaamera abil või laadi üles pilt!/ 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")
+st.title("🦋  Liblikamaja ID / Butterfly Identifier")
+st.write("Tuvasta liblikaid oma kaamera abil või laadi üles pilt! / Identify butterflies using your camera or by uploading an image!")
 
-# Create tabs for different input methods
-tab1, tab2 = st.tabs(["📷 Live Camera", "📁 Upload Image"])
+tab1, tab2 = st.tabs(["📷 Live Camera / Kaamera", "📁 Upload Image / Laadi üles"])
 
 with tab1:
-    st.header("Kaamera jäädvustamine/ Camera Capture")
-    st.write("Take a photo of a butterfly for identification!")
-    
-    camera_photo = st.camera_input("Pildista liblikat/ Take a picture of a butterfly")
-    
+    st.header("Kaamera jäädvustamine / Camera Capture")
+    st.write("Tee pilt liblikast tuvastamiseks / Take a photo of a butterfly for identification.")
+    camera_photo = st.camera_input("Pildista liblikat / Capture a butterfly")
     if camera_photo is not None:
         try:
             image = Image.open(camera_photo).convert("RGB")
-            
             col1, col2 = st.columns(2)
-            
             with col1:
-                st.image(image, caption="Captured Image/Jäädvustatud pilt", use_column_width=True)
-            
+                st.image(image, caption="Jäädvustatud pilt / Captured Image", use_column_width=True)
             with col2:
-                with st.spinner("Kujutise analüüsimine..."):
+                with st.spinner("Pildi analüüsimine... / Analyzing image..."):
                     predicted_class, confidence = predict_butterfly(image)
-                
                 if predicted_class and confidence >= 0.50:
-                    st.success(f"**Liblika: {predicted_class}**")
-                    #st.info(f"Confidence: {confidence:.2%}")
+                    st.success(f"**Liblikas / Butterfly: {predicted_class}**")
+                    if predicted_class in butterfly_info:
+                        st.markdown("**Liigi kirjeldus / About this species:**")
+                        st.write(butterfly_info[predicted_class]["description"])
                 else:
-                    st.warning("⚠️ **I don't know what butterfly this is.**")
-                    st.warning("⚠️ **Low confidence prediction/ Madala usaldusväärsusega ennustus**")
-                    #st.info(f"Best guess: {predicted_class} ({confidence:.1%})")
-                    st.markdown("**Tips for better results/ Näpunäited paremate tulemuste saavutamiseks:**")
-                    st.markdown("- Use better lighting/ Kasutage paremat valgustust")
-                    st.markdown("- Ensure the butterfly is clearly visible/ Veenduge, et liblikas oleks selgelt nähtav")
-                    st.markdown("- Avoid blurry or dark images/ Vältige uduseid või tumedaid pilte")
-                    
-                        
-                                    
+                    st.warning("⚠️ Ma ei tea, mis liblikas see on / I don't know what butterfly this is")
+                    st.markdown("**Näpunäited paremate tulemuste saavutamiseks / Tips for better results:**")
+                    st.markdown("- Kasutage paremat valgustust / Use better lighting")
+                    st.markdown("- Veenduge, et liblikas oleks selgelt nähtav / Ensure the butterfly is clearly visible")
+                    st.markdown("- Vältige uduseid või tumedaid pilte / Avoid blurry or dark images")
         except Exception as e:
             st.error(f"Error processing image: {str(e)}")
 
 with tab2:
-    st.header("Upload Image/ Laadi pilt üles")
-    st.write("Upload a clear photo of a butterfly for identification/ Laadige üles liblika selge foto tuvastamiseks")
-    
-    uploaded_file = st.file_uploader(
-        "Choose an image/ Valige pilt...", 
-        type=["jpg", "jpeg", "png"],
-        help="Upload a clear photo of a butterfly/ Laadi üles selge foto liblikast"
-    )
-    
+    st.header("Laadi üles pilt / Upload Image")
+    st.write("Laadige üles liblika selge foto tuvastamiseks / Upload a clear photo of a butterfly for identification.")
+    uploaded_file = st.file_uploader("Vali pilt... / Choose an image...", type=["jpg", "jpeg", "png"])
     if uploaded_file is not None:
         try:
             image_bytes = uploaded_file.read()
             image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
-            
             col1, col2 = st.columns(2)
-            
             with col1:
-                st.image(image, caption="Uploaded Image", use_column_width=True)
-            
+                st.image(image, caption="Üleslaetud pilt / Uploaded Image", use_column_width=True)
             with col2:
-                with st.spinner("Analyzing image..."):
+                with st.spinner("Pildi analüüsimine... / Analyzing image..."):
                     predicted_class, confidence = predict_butterfly(image)
-                
-                if predicted_class and confidence:
-                    if confidence >= 0.50:
-                        st.success(f"**liblikas: {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/ Madala usaldusväärsusega ennustus**")
-                        #st.info(f"Best guess: {predicted_class} ({confidence:.1%})")
-                        st.markdown("**Tips for better results/ Näpunäited paremate tulemuste saavutamiseks:**")
-                        st.markdown("- Use better lighting/ Kasutage paremat valgustust")
-                        st.markdown("- Ensure the butterfly is clearly visible/ Veenduge, et liblikas oleks selgelt nähtav")
-                        st.markdown("- Avoid blurry or dark images/ Vältige uduseid või tumedaid pilte")
-                
-                
+                if predicted_class and confidence >= 0.50:
+                    st.success(f"**Liblikas / Butterfly: {predicted_class}**")
+                    if predicted_class in butterfly_info:
+                        st.markdown("**Liigi kirjeldus / About this species:**")
+                        st.write(butterfly_info[predicted_class]["description"])
+                else:
+                    st.warning("⚠️ Ma ei tea, mis liblikas see on / I don't know what butterfly this is")
+                    st.markdown("**Näpunäited paremate tulemuste saavutamiseks / Tips for better results:**")
+                    st.markdown("- Kasutage paremat valgustust / Use better lighting")
+                    st.markdown("- Veenduge, et liblikas oleks selgelt nähtav / Ensure the butterfly is clearly visible")
+                    st.markdown("- Vältige uduseid või tumedaid pilte / Avoid blurry or dark images")
         except Exception as e:
             st.error(f"Error processing image: {str(e)}")
 
-# Add footer
+# Footer
 st.markdown("---")
-st.markdown("### How to use/ Kuidas kasutada:")
-st.markdown("1. **Camera Capture**: Take a photo using your device camera/ **Kaamera jäädvustamine**: Tehke foto oma seadme kaameraga")
-st.markdown("2. **Upload Image**: Choose a butterfly photo from your device/ **Laadi pilt üles**: Vali oma seadmest liblika foto")
-st.markdown("3. **Best Results**: Use clear, well-lit photos with the butterfly clearly visible/ **Parimad tulemused**: Kasutage selgeid ja hästi valgustatud fotosid, kus liblikas on selgelt nähtav.")
+st.markdown("### Kuidas kasutada / How to use:")
+st.markdown("1. **Kaamera jäädvustamine / Camera Capture**: Tehke foto oma seadme kaameraga / Take a photo using your device camera")
+st.markdown("2. **Laadi pilt üles / Upload Image**: Vali oma seadmest liblika foto / Choose a butterfly photo from your device")
+st.markdown("3. **Parimad tulemused / Best Results**: Kasutage selgeid ja hästi valgustatud fotosid, kus liblikas on selgelt nähtav / 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")