Update app.py

app.py

@@ -1,25 +1,28 @@
 import streamlit as st
-import torch
 from PIL import Image
-import io
 import numpy as np
-from model_utils import BugClassifier, generate_gradcam, get_severity_prediction
+from model_utils import BugClassifier, get_severity_prediction
 from transformers import AutoFeatureExtractor
 
 # Page configuration
 st.set_page_config(
     page_title="Bug-O-Scope 🔍🐞",
     page_icon="🔍",
-    layout="wide"
+    layout="wide",
+    initial_sidebar_state="expanded"
 )
 
 # Initialize session state
-if 'model' not in st.session_state:
+@st.cache_resource
+def load_model():
     try:
-        st.session_state.model = BugClassifier()
-        st.session_state.feature_extractor = AutoFeatureExtractor.from_pretrained("google/vit-base-patch16-224")
+        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()
 
 def main():
     # Header
@@ -49,6 +52,7 @@ def main():
         compare_bugs()
 
 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:
@@ -76,19 +80,27 @@ def single_bug_analysis():
                         "Medium": "orange",
                         "High": "red"
                     }
-                    st.markdown(f"Severity: <span style='color: {severity_color[severity]}'>{severity}</span>", 
-                               unsafe_allow_html=True)
+                    st.markdown(
+                        f"Severity: <span style='color: {severity_color[severity]}'>{severity}</span>", 
+                        unsafe_allow_html=True
+                    )
 
             # 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.")
 
 def compare_bugs():
+    """Handle bug comparison"""
     col1, col2 = st.columns(2)
     
     with col1:
@@ -119,8 +131,8 @@ def compare_bugs():
                 pred2, conf2 = st.session_state.model.predict(image2)
                 
                 # Generate Grad-CAM visualizations
-                gradcam1 = generate_gradcam(image1, st.session_state.model)
-                gradcam2 = generate_gradcam(image2, st.session_state.model)
+                gradcam1 = st.session_state.model.get_gradcam(image1)
+                gradcam2 = st.session_state.model.get_gradcam(image2)
                 
                 # Display results
                 st.markdown("### Comparison Results")