Update app.py

app.py

@@ -15,8 +15,11 @@ st.set_page_config(
 
 # Initialize session state
 if 'model' not in st.session_state:
-    st.session_state.model = BugClassifier()
-    st.session_state.feature_extractor = AutoFeatureExtractor.from_pretrained("google/vit-base-patch16-224")
+    try:
+        st.session_state.model = BugClassifier()
+        st.session_state.feature_extractor = AutoFeatureExtractor.from_pretrained("google/vit-base-patch16-224")
+    except Exception as e:
+        st.error(f"Error loading model: {str(e)}")
 
 def main():
     # Header
@@ -49,36 +52,41 @@ def single_bug_analysis():
     uploaded_file = st.file_uploader("Upload a bug photo", type=['png', 'jpg', 'jpeg'], key="single")
     
     if uploaded_file:
-        image = Image.open(uploaded_file)
-        col1, col2 = st.columns(2)
-        
-        with col1:
-            st.image(image, caption="Uploaded Image", use_column_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)
+        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)
 
-        # Generate and display species information
-        st.markdown("### About This Species")
-        species_info = st.session_state.model.get_species_info(prediction)
-        st.markdown(species_info)
+            # Generate and display species information
+            st.markdown("### About This Species")
+            species_info = st.session_state.model.get_species_info(prediction)
+            st.markdown(species_info)
+            
+        except Exception as e:
+            st.error(f"Error processing image: {str(e)}")
+            st.info("Please try uploading a different image.")
 
 def compare_bugs():
     col1, col2 = st.columns(2)
@@ -86,43 +94,56 @@ def compare_bugs():
     with col1:
         file1 = st.file_uploader("Upload first bug photo", type=['png', 'jpg', 'jpeg'], key="compare1")
         if file1:
-            image1 = Image.open(file1)
-            st.image(image1, caption="First Bug", use_column_width=True)
+            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 col2:
         file2 = st.file_uploader("Upload second bug photo", type=['png', 'jpg', 'jpeg'], key="compare2")
         if file2:
-            image2 = Image.open(file2)
-            st.image(image2, caption="Second Bug", use_column_width=True)
+            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:
-        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 = generate_gradcam(image1, st.session_state.model)
-            gradcam2 = generate_gradcam(image2, st.session_state.model)
-            
-            # 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_column_width=True)
+        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)
                 
-            with comp_col2:
-                st.markdown(f"**Species 2**: {pred2}")
-                st.markdown(f"Confidence: {conf2:.2f}%")
-                st.image(gradcam2, caption="Feature Highlights - Bug 2", use_column_width=True)
-            
-            # Display comparison information
-            st.markdown("### Key Differences")
-            differences = st.session_state.model.compare_species(pred1, pred2)
-            st.markdown(differences)
+                # Generate Grad-CAM visualizations
+                gradcam1 = generate_gradcam(image1, st.session_state.model)
+                gradcam2 = generate_gradcam(image2, st.session_state.model)
+                
+                # 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 __name__ == "__main__":
     main()