Update streamlit_app.py

streamlit_app.py

@@ -9,6 +9,7 @@ import numpy as np
 from streamlit_webrtc import webrtc_streamer, WebRtcMode, RTCConfiguration
 import av
 import cv2
+import timm  
 
 # Configure Streamlit
 st.set_page_config(
@@ -36,12 +37,13 @@ def load_model():
         st.error("Model file not found. Please upload butterfly_classifier.pth to your space.")
         return None
     
-    model = models.resnet18(pretrained=False)
-    model.fc = torch.nn.Linear(model.fc.in_features, len(class_names))
+    # 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"))
     model.eval()
     return model
 
+# Load the model
 model = load_model()
 
 if model is None:
@@ -73,24 +75,6 @@ def predict_butterfly(image):
         
     return predicted_class, confidence.item()
 
-def display_prediction(predicted_class, confidence):
-    """Display prediction with confidence check"""
-    if confidence < 0.70:  # Below 70% confidence
-        st.warning("⚠️ **Image not clear - Unable to identify butterfly**")
-        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")
-    else:
-        st.success(f"**🦋 Identified: {predicted_class}**")
-        st.info(f"Confidence: {confidence:.1%}")
-        
-        if predicted_class in butterfly_info:
-            st.write("**Species Information:**")
-            st.write(butterfly_info[predicted_class]["description"])
-            
 # Video frame callback for live camera
 class VideoProcessor:
     def __init__(self):
@@ -107,7 +91,7 @@ class VideoProcessor:
             rgb_img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
             predicted_class, confidence = predict_butterfly(rgb_img)
             
-            if predicted_class and confidence > 0.7:  # Only show if confident
+            if predicted_class and confidence > 0.8:  # Only show if confidence > 80%
                 self.prediction_text = f"{predicted_class} ({confidence:.2f})"
         
         # Draw prediction on frame
@@ -144,14 +128,22 @@ with tab1:
             with col2:
                 predicted_class, confidence = predict_butterfly(image)
                 
-                if predicted_class:
+                if predicted_class and confidence >= 0.80:  # Only show if confidence >= 80%
                     st.success(f"**Prediction: {predicted_class}**")
                     st.info(f"Confidence: {confidence:.2%}")
                     
                     if predicted_class in butterfly_info:
                         st.write("**Species Information:**")
                         st.write(butterfly_info[predicted_class]["description"])
-                        
+                else:
+                    st.warning("⚠️ **Image not clear - Unable to identify butterfly**")
+                    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")
+                                    
         except Exception as e:
             st.error(f"Error processing image: {str(e)}")
 
@@ -179,13 +171,21 @@ with tab2:
             with col2:
                 predicted_class, confidence = predict_butterfly(image)
                 
-                if predicted_class:
+                if predicted_class and confidence >= 0.80:  # Only show if confidence >= 80%
                     st.success(f"**Prediction: {predicted_class}**")
                     st.info(f"Confidence: {confidence:.2%}")
                     
                     if predicted_class in butterfly_info:
                         st.write("**Species Information:**")
                         st.write(butterfly_info[predicted_class]["description"])
+                else:
+                    st.warning("⚠️ **Image not clear - Unable to identify butterfly**")
+                    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")
                 
         except Exception as e:
             st.error(f"Error processing image: {str(e)}")
@@ -193,6 +193,6 @@ with tab2:
 # Add footer with instructions
 st.markdown("---")
 st.markdown("### How to use:")
-st.markdown("1. **Live Camera**: Click 'START' to begin live identification")
+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")