Update streamlit_app.py

streamlit_app.py

@@ -1,4 +1,4 @@
-# Full corrected bilingual Streamlit app for Butterfly Identifier
+# Enhanced Butterfly Identifier Streamlit App with Better Model Loading
 import streamlit as st
 from PIL import Image
 import torch
@@ -49,136 +49,448 @@ inference_transform = A.Compose([
     ToTensorV2()
 ])
 
-# Load the model
+# Enhanced model loading function
 @st.cache_resource
 def load_model():
-    MODEL_PATH = "butterfly_classifier.pth"
-    if not os.path.exists(MODEL_PATH):
-        st.error("Model file not found!")
+    """Enhanced model loading with architecture detection and fallback options"""
+    
+    # Try different model file names
+    model_files = [
+        "butterfly_classifier.pth",
+        "best_butterfly_model_v3.pth",
+        "best_butterfly_model.pth"
+    ]
+    
+    MODEL_PATH = None
+    for model_file in model_files:
+        if os.path.exists(model_file):
+            MODEL_PATH = model_file
+            break
+    
+    if MODEL_PATH is None:
+        st.error("No model file found! Please ensure one of these files exists: " + ", ".join(model_files))
         return None
+    
+    st.info(f"Loading model from: {MODEL_PATH}")
+    
     try:
+        # Load checkpoint
         checkpoint = torch.load(MODEL_PATH, map_location='cpu')
-        model_state_dict = checkpoint['model_state_dict'] if 'model_state_dict' in checkpoint else checkpoint
-        num_classes = len(class_names)
-
-        # Attempt to auto-detect model from batch norm layer dimensions
-        bn2_shape = None
-        for key in model_state_dict:
-            if key.endswith("bn2.weight"):
-                bn2_shape = model_state_dict[key].shape[0]
-                break
-
-        feature_map = {
-            1280: 'efficientnet_b0',
-            1408: 'efficientnet_b1',
-            1536: 'efficientnet_b2',
-            1792: 'efficientnet_b3',
-            1920: 'efficientnet_b4',
-            2048: 'efficientnet_b5',
-            2304: 'efficientnet_b6',
-            2560: 'efficientnet_b7'
-        }
-
-        if bn2_shape is None:
-            st.warning("Could not detect classifier or bn2 layer in checkpoint. Defaulting to efficientnet_b3")
-            model_name = 'efficientnet_b3'
+        
+        # Extract model state dict
+        if 'model_state_dict' in checkpoint:
+            model_state_dict = checkpoint['model_state_dict']
+            if 'class_names' in checkpoint:
+                st.info(f"Model trained on {len(checkpoint['class_names'])} classes")
         else:
-            model_name = feature_map.get(bn2_shape, 'efficientnet_b3')
-            st.info(f"Detected model architecture: {model_name}")
-
-        model = timm.create_model(model_name, 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_state_dict = checkpoint
+        
+        num_classes = len(class_names)
+        
+        # Architecture detection based on model state dict
+        def detect_model_architecture(state_dict):
+            """Detect model architecture from state dict"""
+            
+            # Check for EfficientNet variants by looking at key layer dimensions
+            architecture_indicators = {
+                'conv_head.weight': 'efficientnet',
+                'head.weight': 'efficientnet_v2',
+                'classifier.weight': 'other'
+            }
+            
+            # Look for specific layer patterns
+            for key in state_dict.keys():
+                if 'conv_head.weight' in key:
+                    shape = state_dict[key].shape
+                    if len(shape) >= 2:
+                        feature_dim = shape[1]
+                        # EfficientNet feature dimensions
+                        efficientnet_map = {
+                            1280: 'efficientnet_b0',
+                            1408: 'efficientnet_b1', 
+                            1536: 'efficientnet_b2',
+                            1792: 'efficientnet_b3',
+                            1920: 'efficientnet_b4',
+                            2048: 'efficientnet_b5',
+                            2304: 'efficientnet_b6',
+                            2560: 'efficientnet_b7'
+                        }
+                        return efficientnet_map.get(feature_dim, 'efficientnet_b3')
+                
+                if 'head.weight' in key:
+                    shape = state_dict[key].shape
+                    if len(shape) >= 2:
+                        feature_dim = shape[1]
+                        # EfficientNetV2 feature dimensions
+                        efficientnetv2_map = {
+                            1280: 'tf_efficientnetv2_s',
+                            1408: 'tf_efficientnetv2_m',
+                            1792: 'tf_efficientnetv2_l'
+                        }
+                        return efficientnetv2_map.get(feature_dim, 'tf_efficientnetv2_s')
+            
+            # Fallback: check bn2 layer for EfficientNet variants
+            for key in state_dict.keys():
+                if key.endswith("bn2.weight"):
+                    bn2_shape = state_dict[key].shape[0]
+                    feature_map = {
+                        1280: 'efficientnet_b0',
+                        1408: 'efficientnet_b1',
+                        1536: 'efficientnet_b2',
+                        1792: 'efficientnet_b3',
+                        1920: 'efficientnet_b4',
+                        2048: 'efficientnet_b5',
+                        2304: 'efficientnet_b6',
+                        2560: 'efficientnet_b7'
+                    }
+                    return feature_map.get(bn2_shape, 'efficientnet_b3')
+            
+            return 'efficientnet_b3'  # Default fallback
+        
+        # Detect architecture
+        detected_arch = detect_model_architecture(model_state_dict)
+        st.info(f"Detected model architecture: {detected_arch}")
+        
+        # List of architectures to try in order
+        architectures_to_try = [
+            detected_arch,
+            'efficientnet_b3',
+            'efficientnet_b2',
+            'efficientnet_b0',
+            'efficientnet_b1',
+            'efficientnet_b4',
+            'tf_efficientnetv2_s',
+            'tf_efficientnetv2_m'
+        ]
+        
+        # Remove duplicates while preserving order
+        seen = set()
+        architectures_to_try = [x for x in architectures_to_try if not (x in seen or seen.add(x))]
+        
+        model = None
+        successful_arch = None
+        
+        # Try each architecture
+        for arch in architectures_to_try:
+            try:
+                st.info(f"Trying architecture: {arch}")
+                
+                # Create model with the detected/guessed architecture
+                model = timm.create_model(
+                    arch,
+                    pretrained=False,
+                    num_classes=num_classes,
+                    drop_rate=0.4,
+                    drop_path_rate=0.3
+                )
+                
+                # Check if the model has a custom head/classifier in the checkpoint
+                if any('head.' in key for key in model_state_dict.keys()):
+                    # Model has custom head - try to load it
+                    try:
+                        model.load_state_dict(model_state_dict, strict=False)
+                        st.success(f"✅ Successfully loaded model with architecture: {arch}")
+                        successful_arch = arch
+                        break
+                    except Exception as e:
+                        st.warning(f"Failed to load custom head for {arch}: {str(e)}")
+                        continue
+                
+                elif any('classifier.' in key for key in model_state_dict.keys()):
+                    # Model has custom classifier - try to load it
+                    try:
+                        model.load_state_dict(model_state_dict, strict=False)
+                        st.success(f"✅ Successfully loaded model with architecture: {arch}")
+                        successful_arch = arch
+                        break
+                    except Exception as e:
+                        st.warning(f"Failed to load custom classifier for {arch}: {str(e)}")
+                        continue
+                
+                else:
+                    # Try to create custom head/classifier and load backbone
+                    try:
+                        # Load backbone weights (ignore head/classifier mismatches)
+                        backbone_dict = {k: v for k, v in model_state_dict.items() 
+                                       if not (k.startswith('head.') or k.startswith('classifier.'))}
+                        
+                        model.load_state_dict(backbone_dict, strict=False)
+                        
+                        # Create new head/classifier
+                        if hasattr(model, 'classifier'):
+                            in_features = model.classifier.in_features
+                            model.classifier = torch.nn.Linear(in_features, num_classes)
+                        elif hasattr(model, 'head'):
+                            in_features = model.head.in_features
+                            model.head = torch.nn.Linear(in_features, num_classes)
+                        
+                        st.warning(f"⚠️ Loaded {arch} with new head/classifier (backbone weights only)")
+                        successful_arch = arch
+                        break
+                        
+                    except Exception as e:
+                        st.warning(f"Failed to load backbone for {arch}: {str(e)}")
+                        continue
+                        
+            except Exception as e:
+                st.warning(f"Failed to create model {arch}: {str(e)}")
+                continue
+        
+        if model is None:
+            st.error("❌ Failed to load model with any architecture!")
+            return None
+        
+        # Set model to evaluation mode
         model.eval()
+        
+        # Display model info
+        total_params = sum(p.numel() for p in model.parameters())
+        st.success(f"✅ Model loaded successfully!")
+        st.info(f"📊 Model: {successful_arch}")
+        st.info(f"🔢 Parameters: {total_params:,}")
+        st.info(f"🎯 Classes: {num_classes}")
+        
         return model
+        
     except Exception as e:
-        st.error(f"Error loading model: {str(e)}")
+        st.error(f"❌ Error loading model: {str(e)}")
+        st.error("Please check your model file and ensure it's compatible")
         return None
 
+# Load model
 model = load_model()
 
 def predict_butterfly(image, threshold=0.5):
+    """Predict butterfly species from image"""
     try:
         if model is None:
             raise ValueError("Model is not loaded.")
         if image is None:
             return None, None
+            
+        # Convert to PIL Image if needed
         if isinstance(image, np.ndarray):
             image = Image.fromarray(image)
         if image.mode != 'RGB':
             image = image.convert('RGB')
+        
+        # Apply transforms
         transformed = inference_transform(image=np.array(image))
         input_tensor = transformed['image'].unsqueeze(0)
+        
+        # Make prediction
         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
 
+def predict_with_tta(image, threshold=0.5, num_tta=5):
+    """Predict with Test Time Augmentation for better accuracy"""
+    try:
+        if model is None:
+            raise ValueError("Model is not loaded.")
+        if image is None:
+            return None, None
+            
+        # Convert to PIL Image if needed
+        if isinstance(image, np.ndarray):
+            image = Image.fromarray(image)
+        if image.mode != 'RGB':
+            image = image.convert('RGB')
+        
+        # Convert to numpy for albumentations
+        image_np = np.array(image)
+        
+        # TTA transforms
+        tta_transforms = [
+            A.Compose([
+                A.Resize(224, 224),
+                A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+                ToTensorV2()
+            ]),
+            A.Compose([
+                A.Resize(256, 256),
+                A.CenterCrop(224, 224),
+                A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+                ToTensorV2()
+            ]),
+            A.Compose([
+                A.Resize(224, 224),
+                A.HorizontalFlip(p=1.0),
+                A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+                ToTensorV2()
+            ]),
+            A.Compose([
+                A.Resize(240, 240),
+                A.Rotate(limit=10, p=1.0),
+                A.CenterCrop(224, 224),
+                A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+                ToTensorV2()
+            ]),
+            A.Compose([
+                A.Resize(224, 224),
+                A.ColorJitter(brightness=0.1, contrast=0.1, p=1.0),
+                A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+                ToTensorV2()
+            ])
+        ]
+        
+        predictions = []
+        
+        for i, transform in enumerate(tta_transforms[:num_tta]):
+            transformed = transform(image=image_np)
+            input_tensor = transformed['image'].unsqueeze(0)
+            
+            with torch.no_grad():
+                output = model(input_tensor)
+                probabilities = torch.nn.functional.softmax(output, dim=1)
+                predictions.append(probabilities)
+        
+        # Average predictions
+        avg_predictions = torch.mean(torch.stack(predictions), dim=0)
+        confidence, pred = torch.max(avg_predictions, 1)
+        
+        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"TTA Prediction error: {str(e)}")
+        return None, None
+
 # UI Code
-st.title("🦋  Liblikamaja ID / Butterfly Identifier")
+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!")
 
-tab1, tab2 = st.tabs(["📷 Live Camera / Kaamera", "📁 Upload Image / Laadi üles"])
+# Add model status indicator
+if model is not None:
+    st.success("✅ Model loaded and ready!")
+else:
+    st.error("❌ Model not loaded. Please check your model file.")
+    st.stop()
 
+# Add advanced options
+with st.expander("🔧 Advanced Options / Täpsemad seaded"):
+    confidence_threshold = st.slider(
+        "Confidence Threshold / Kindluse lävi", 
+        min_value=0.1, 
+        max_value=1.0, 
+        value=0.5, 
+        step=0.05,
+        help="Higher values = more conservative predictions"
+    )
+    
+    use_tta = st.checkbox(
+        "Use Test Time Augmentation (TTA) / Kasuta TTA", 
+        value=False,
+        help="Slower but potentially more accurate predictions"
+    )
+    
+    if use_tta:
+        tta_rounds = st.slider(
+            "TTA Rounds / TTA ringid", 
+            min_value=3, 
+            max_value=8, 
+            value=5,
+            help="More rounds = slower but potentially more accurate"
+        )
+
+tab1, tab2 = st.tabs(["📷 Live Camera / Kaamera", "📁 Upload Image / Laadi üles"])
 
 with tab1:
     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="Jäädvustatud pilt / Captured Image", use_column_width=True)
+            
             with col2:
                 with st.spinner("Pildi analüüsimine... / Analyzing image..."):
-                    predicted_class, confidence = predict_butterfly(image)
-                if predicted_class and confidence >= 0.50:
+                    if use_tta:
+                        predicted_class, confidence = predict_with_tta(image, confidence_threshold, tta_rounds)
+                    else:
+                        predicted_class, confidence = predict_butterfly(image, confidence_threshold)
+                
+                if predicted_class and confidence >= confidence_threshold:
                     st.success(f"**Liblikas / Butterfly: {predicted_class}**")
+                    st.info(f"Confidence: {confidence:.2%}")
+                    
                     if predicted_class in butterfly_info:
                         st.markdown("**Liigi kirjeldus / About this species:**")
                         st.write(butterfly_info[predicted_class]["description"])
+                    else:
+                        st.info("No additional information available for this species.")
                 else:
-                    st.warning("⚠️ Ma ei tea, mis liblikas see on / I don't know what butterfly this is")
+                    confidence_text = f" (Confidence: {confidence:.2%})" if confidence else ""
+                    st.warning(f"⚠️ Ma ei tea, mis liblikas see on / I don't know what butterfly this is{confidence_text}")
                     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")
+                    st.markdown("- Proovige madalamat kindluse läviväärtust / Try a lower confidence threshold")
+                    
         except Exception as e:
             st.error(f"Error processing image: {str(e)}")
 
 with tab2:
     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="Üleslaetud pilt / Uploaded Image", use_column_width=True)
+            
             with col2:
                 with st.spinner("Pildi analüüsimine... / Analyzing image..."):
-                    predicted_class, confidence = predict_butterfly(image)
-                if predicted_class and confidence >= 0.50:
+                    if use_tta:
+                        predicted_class, confidence = predict_with_tta(image, confidence_threshold, tta_rounds)
+                    else:
+                        predicted_class, confidence = predict_butterfly(image, confidence_threshold)
+                
+                if predicted_class and confidence >= confidence_threshold:
                     st.success(f"**Liblikas / Butterfly: {predicted_class}**")
+                    st.info(f"Confidence: {confidence:.2%}")
+                    
                     if predicted_class in butterfly_info:
                         st.markdown("**Liigi kirjeldus / About this species:**")
                         st.write(butterfly_info[predicted_class]["description"])
+                    else:
+                        st.info("No additional information available for this species.")
                 else:
-                    st.warning("⚠️ Ma ei tea, mis liblikas see on / I don't know what butterfly this is")
+                    confidence_text = f" (Confidence: {confidence:.2%})" if confidence else ""
+                    st.warning(f"⚠️ Ma ei tea, mis liblikas see on / I don't know what butterfly this is{confidence_text}")
                     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")
+                    st.markdown("- Proovige madalamat kindluse läviväärtust / Try a lower confidence threshold")
+                    
         except Exception as e:
             st.error(f"Error processing image: {str(e)}")
 
@@ -188,4 +500,12 @@ 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")
+st.markdown("4. **Täpsemad seaded / Advanced Options**: Kohandage kindluse lävi ja kasutage TTA paremate tulemuste saamiseks / Adjust confidence threshold and use TTA for better results")
 
+# Debug info
+if st.checkbox("Show Debug Info"):
+    st.write("**Class Names:**", class_names)
+    st.write("**Number of Classes:**", len(class_names))
+    st.write("**Model Status:**", "Loaded" if model else "Not Loaded")
+    if butterfly_info:
+        st.write("**Species Info Available:**", len(butterfly_info))