Update app.py

app.py

@@ -1,161 +1,65 @@
+# 1. app.py
+# Main Streamlit app for UI and user interaction
 import streamlit as st
 from PIL import Image
-import numpy as np
-from model_utils import BugClassifier, get_severity_prediction
-from transformers import AutoFeatureExtractor
+from dataset_utils import load_insect_dataset, preprocess_image
+from model_utils import BugClassifier
 
-# Page configuration
 st.set_page_config(
     page_title="Bug-O-Scope 🔍🐞",
     page_icon="🔍",
-    layout="wide",
-    initial_sidebar_state="expanded"
+    layout="wide"
 )
 
-# Initialize session state
 @st.cache_resource
 def load_model():
-    try:
-        return BugClassifier(), AutoFeatureExtractor.from_pretrained("google/vit-base-patch16-224")
-    except Exception as e:
-        st.error(f"Error loading model: {str(e)}")
-        return None, None
-
-if 'model' not in st.session_state:
-    st.session_state.model, st.session_state.feature_extractor = load_model()
+    return BugClassifier()
 
 def main():
-    # Header
     st.title("Bug-O-Scope 🔍🐞")
-    st.markdown("""
-    Welcome to Bug-O-Scope! Upload a picture of an insect to learn more about it.
-    This educational tool helps you identify bugs and understand their role in our ecosystem.
-    """)
+    st.markdown(
+        """
+        Welcome to Bug-O-Scope! Upload a photo of an insect to identify it and learn about its ecological importance.
+        """
+    )
 
-    # Sidebar
     st.sidebar.header("About Bug-O-Scope")
-    st.sidebar.markdown("""
-    Bug-O-Scope is an AI-powered tool that helps you:
-    * 🔍 Identify insects from photos
-    * 📚 Learn about different species
-    * 🌍 Understand their ecological impact
-    * 🔬 Compare different insects
-    """)
+    st.sidebar.markdown(
+        """
+        Bug-O-Scope is powered by AI to help you:
+        - Identify insects
+        - Learn about species
+        - Understand their ecological roles
+        """
+    )
+
+    model = load_model()
 
-    # Main content
     tab1, tab2 = st.tabs(["Single Bug Analysis", "Bug Comparison"])
 
     with tab1:
-        single_bug_analysis()
+        single_bug_analysis(model)
 
     with tab2:
-        compare_bugs()
+        st.markdown("Bug comparison feature coming soon!")
 
-def single_bug_analysis():
-    """Handle single bug analysis"""
-    uploaded_file = st.file_uploader("Upload a bug photo", type=['png', 'jpg', 'jpeg'], key="single")
-    
-    if uploaded_file:
-        try:
-            image = Image.open(uploaded_file)
-            col1, col2 = st.columns(2)
-            
-            with col1:
-                st.image(image, caption="Uploaded Image", use_container_width=True)
-            
-            with col2:
-                with st.spinner("Analyzing your bug..."):
-                    # Get predictions
-                    prediction, confidence = st.session_state.model.predict(image)
-                    severity = get_severity_prediction(prediction)
-                    
-                    st.success("Analysis Complete!")
-                    st.markdown(f"### Identified Species")
-                    st.markdown(f"**{prediction}**")
-                    st.markdown(f"Confidence: {confidence:.2f}%")
-                    
-                    st.markdown("### Ecological Impact")
-                    severity_color = {
-                        "Low": "green",
-                        "Medium": "orange",
-                        "High": "red"
-                    }
-                    st.markdown(
-                        f"Severity: <span style='color: {severity_color[severity]}'>{severity}</span>", 
-                        unsafe_allow_html=True
-                    )
+def single_bug_analysis(model):
+    uploaded_file = st.file_uploader("Upload a bug photo", type=['png', 'jpg', 'jpeg'])
 
-            # Generate and display species information
-            st.markdown("### About This Species")
-            species_info = st.session_state.model.get_species_info(prediction)
-            st.markdown(species_info)
-            
-            # Display Grad-CAM visualization
-            st.markdown("### Feature Highlights")
-            gradcam = st.session_state.model.get_gradcam(image)
-            st.image(gradcam, caption="Important Features", use_container_width=True)
-            
-        except Exception as e:
-            st.error(f"Error processing image: {str(e)}")
-            st.info("Please try uploading a different image.")
+    if uploaded_file:
+        image = Image.open(uploaded_file)
+        st.image(image, caption="Uploaded Image", use_container_width=True)
 
-def compare_bugs():
-    """Handle bug comparison"""
-    col1, col2 = st.columns(2)
-    
-    with col1:
-        file1 = st.file_uploader("Upload first bug photo", type=['png', 'jpg', 'jpeg'], key="compare1")
-        if file1:
-            try:
-                image1 = Image.open(file1)
-                st.image(image1, caption="First Bug", use_container_width=True)
-            except Exception as e:
-                st.error(f"Error loading first image: {str(e)}")
-                return
+        with st.spinner("Analyzing your bug..."):
+            prediction, confidence = model.predict(image)
+            st.success("Analysis Complete!")
+            st.markdown(f"**Identified Species**: {prediction}")
+            st.markdown(f"**Confidence**: {confidence:.2f}%")
             
-    with col2:
-        file2 = st.file_uploader("Upload second bug photo", type=['png', 'jpg', 'jpeg'], key="compare2")
-        if file2:
-            try:
-                image2 = Image.open(file2)
-                st.image(image2, caption="Second Bug", use_container_width=True)
-            except Exception as e:
-                st.error(f"Error loading second image: {str(e)}")
-                return
-    
-    if file1 and file2:
-        try:
-            with st.spinner("Generating comparison..."):
-                # Get predictions for both images
-                pred1, conf1 = st.session_state.model.predict(image1)
-                pred2, conf2 = st.session_state.model.predict(image2)
-                
-                # Generate Grad-CAM visualizations
-                gradcam1 = st.session_state.model.get_gradcam(image1)
-                gradcam2 = st.session_state.model.get_gradcam(image2)
-                
-                # Display results
-                st.markdown("### Comparison Results")
-                comp_col1, comp_col2 = st.columns(2)
-                
-                with comp_col1:
-                    st.markdown(f"**Species 1**: {pred1}")
-                    st.markdown(f"Confidence: {conf1:.2f}%")
-                    st.image(gradcam1, caption="Feature Highlights - Bug 1", use_container_width=True)
-                    
-                with comp_col2:
-                    st.markdown(f"**Species 2**: {pred2}")
-                    st.markdown(f"Confidence: {conf2:.2f}%")
-                    st.image(gradcam2, caption="Feature Highlights - Bug 2", use_container_width=True)
-                
-                # Display comparison information
-                st.markdown("### Key Differences")
-                differences = st.session_state.model.compare_species(pred1, pred2)
-                st.markdown(differences)
-                
-        except Exception as e:
-            st.error(f"Error comparing images: {str(e)}")
-            st.info("Please try uploading different images or try again.")
+            if prediction != "Unknown Insect":
+                species_info = model.get_species_info(prediction)
+                st.markdown("### About This Species")
+                st.markdown(species_info)
 
 if __name__ == "__main__":
     main()