Update model_utils.py

model_utils.py

@@ -1,158 +1,176 @@
-import torch
-from transformers import ViTImageProcessor, ViTForImageClassification
-import numpy as np
+import streamlit as st
 from PIL import Image
-import cv2
+from model_utils import BugClassifier, get_severity_prediction
 
-class BugClassifier:
-    def __init__(self):
-        """Initialize the bug classifier with a simple ViT model"""
-        try:
-            # Load the model and processor
-            model_name = "google/vit-base-patch16-224"
-            self.processor = ViTImageProcessor.from_pretrained(model_name)
-            self.model = ViTForImageClassification.from_pretrained(
-                model_name,
-                num_labels=10,
-                ignore_mismatched_sizes=True
-            )
-            
-            # Set model to evaluation mode
-            self.model.eval()
-            
-            # Define class labels
-            self.labels = [
-                "Seven-spotted Ladybug", "Monarch Butterfly", "Carpenter Ant", 
-                "Japanese Beetle", "Garden Spider", "Green Grasshopper",
-                "Luna Moth", "Common Dragonfly", "Honey Bee", "Paper Wasp"
-            ]
-            
-            # Species information
-            self.species_info = {
-                "Seven-spotted Ladybug": """
-                The Seven-spotted Ladybug (Coccinella septempunctata) is a beneficial garden insect.
-                Key characteristics:
-                - Red wing covers with seven black spots
-                - Natural pest controller, eating aphids and other small insects
-                - Typically 7-8mm in length
-                - Can eat up to 5,000 aphids in their lifetime
-                """,
-                "Monarch Butterfly": """
-                The Monarch Butterfly (Danaus plexippus) is known for its migration patterns.
-                Key characteristics:
-                - Orange wings with black veins and white spots
-                - Wingspan of 93-105mm
-                - Feeds on milkweed as caterpillars
-                - Makes annual migrations of up to 3,000 miles
-                """,
-                # Add more species info as needed
-            }
-            
-            print("Model initialized successfully")
-            
-        except Exception as e:
-            print(f"Error initializing model: {str(e)}")
-            raise RuntimeError(f"Failed to initialize BugClassifier: {str(e)}")
+# Page configuration
+st.set_page_config(
+    page_title="Bug-O-Scope 🔍🐞",
+    page_icon="🔍",
+    layout="wide"
+)
 
-    def preprocess_image(self, image):
-        """Preprocess the image for model input"""
-        try:
-            # Ensure image is in RGB format
-            if image.mode != 'RGB':
-                image = image.convert('RGB')
-            
-            # Process image using the ViT processor
-            inputs = self.processor(images=image, return_tensors="pt")
-            return inputs.pixel_values
-            
-        except Exception as e:
-            print(f"Error preprocessing image: {str(e)}")
-            raise
+# Initialize session state for model
+@st.cache_resource
+def load_model():
+    try:
+        print("Loading model...")
+        model = BugClassifier()
+        print("Model loaded successfully")
+        return model
+    except Exception as e:
+        print(f"Error loading model: {str(e)}")
+        return None
+
+# Ensure model is loaded
+if 'model' not in st.session_state:
+    st.session_state.model = load_model()
+
+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.
+    """)
+
+    # 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
+    """)
+
+    # Check if model loaded successfully
+    if st.session_state.model is None:
+        st.error("Error: Model failed to load. Please try refreshing the page.")
+        return
+
+    # Main content
+    tab1, tab2 = st.tabs(["Single Bug Analysis", "Bug Comparison"])
 
-    def predict(self, image):
-        """Predict the insect species from an image"""
+    with tab1:
+        single_bug_analysis()
+
+    with tab2:
+        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:
         try:
-            print("Starting prediction...")
+            # Load and display image
+            image = Image.open(uploaded_file)
+            col1, col2 = st.columns(2)
             
-            # Preprocess image
-            inputs = self.preprocess_image(image)
-            print("Image preprocessed successfully")
+            with col1:
+                st.image(image, caption="Uploaded Image", use_container_width=True)
             
-            # Make prediction
-            with torch.no_grad():
-                outputs = self.model(inputs)
-                probabilities = torch.nn.functional.softmax(outputs.logits, dim=1)
-                confidence, predicted = torch.max(probabilities, 1)
-                
-                confidence_score = confidence.item() * 100
-                predicted_label = self.labels[predicted.item()]
-                
-                print(f"Prediction complete: {predicted_label} with confidence {confidence_score}%")
-                
-                # Return Unknown if confidence is too low
-                if confidence_score < 40:
-                    return "Unknown Insect", confidence_score
-                
-                return predicted_label, confidence_score
+            with col2:
+                with st.spinner("Analyzing your bug..."):
+                    # Get predictions
+                    prediction, confidence = st.session_state.model.predict(image)
+                    print(f"Prediction: {prediction}, Confidence: {confidence}")
+                    
+                    st.success("Analysis Complete!")
+                    st.markdown("### Identified Species")
+                    st.markdown(f"**{prediction}**")
+                    st.markdown(f"Confidence: {confidence:.2f}%")
+                    
+                    # Only show ecological impact for known insects
+                    if prediction != "Unknown Insect" and prediction != "Error Processing Image":
+                        severity = get_severity_prediction(prediction)
+                        st.markdown("### Ecological Impact")
+                        severity_color = {
+                            "Low": "green",
+                            "Medium": "orange",
+                            "High": "red",
+                            "Unknown": "gray"
+                        }
+                        st.markdown(
+                            f"Severity: <span style='color: {severity_color[severity]}'>{severity}</span>", 
+                            unsafe_allow_html=True
+                        )
+
+            # Display species information
+            if prediction != "Unknown Insect" and prediction != "Error Processing Image":
+                st.markdown("### About This Species")
+                species_info = st.session_state.model.get_species_info(prediction)
+                st.markdown(species_info)
                 
+                # Display 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:
-            print(f"Prediction error: {str(e)}")
-            return "Error Processing Image", 0.0
+            st.error(f"Error processing image: {str(e)}")
+            st.info("Please try uploading a different image.")
 
-    def get_species_info(self, species):
-        """Get information about a species"""
-        default_info = f"""
-        This appears to be a {species}.
-        While we don't have detailed information about this specific species,
-        all insects play important roles in their ecosystems.
-        Consider taking another photo from a different angle or in better lighting
-        for more accurate identification.
-        """
-        return self.species_info.get(species, default_info)
-
-    def get_gradcam(self, image):
-        """Generate attention visualization"""
-        try:
-            # Basic attention visualization
-            inputs = self.preprocess_image(image)
+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 torch.no_grad():
-                outputs = self.model(inputs, output_attentions=True)
-                # Get attention from last layer
-                attention = outputs.attentions[-1].mean(dim=1).mean(dim=1)
+    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
+                pred1, conf1 = st.session_state.model.predict(image1)
+                pred2, conf2 = st.session_state.model.predict(image2)
+                
+                if pred1 not in ["Unknown Insect", "Error Processing Image"] and \
+                   pred2 not in ["Unknown Insect", "Error Processing Image"]:
+                    
+                    # 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}%")
+                        gradcam1 = st.session_state.model.get_gradcam(image1)
+                        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}%")
+                        gradcam2 = st.session_state.model.get_gradcam(image2)
+                        st.image(gradcam2, caption="Feature Highlights - Bug 2", use_container_width=True)
+                    
+                    # Display comparison
+                    st.markdown("### Key Differences")
+                    st.markdown(st.session_state.model.get_species_info(pred1))
+                    st.markdown(st.session_state.model.get_species_info(pred2))
+                else:
+                    st.warning("Unable to generate meaningful comparison due to low confidence predictions.")
                 
-            # Process attention map
-            attention_map = attention.numpy()[0]
-            attention_map = cv2.resize(attention_map, (224, 224))
-            attention_map = np.uint8(255 * attention_map)
-            heatmap = cv2.applyColorMap(attention_map, cv2.COLORMAP_JET)
-            
-            # Prepare original image
-            img_array = np.array(image.resize((224, 224)))
-            
-            # Combine heatmap with original image
-            output = cv2.addWeighted(img_array, 0.7, heatmap, 0.3, 0)
-            
-            return Image.fromarray(output)
-            
         except Exception as e:
-            print(f"Visualization error: {str(e)}")
-            return image
+            st.error(f"Error comparing images: {str(e)}")
+            st.info("Please try uploading different images or try again.")
 
-def get_severity_prediction(species):
-    """Get ecological severity prediction"""
-    severity_map = {
-        "Seven-spotted Ladybug": "Low",
-        "Monarch Butterfly": "Low",
-        "Carpenter Ant": "Medium",
-        "Japanese Beetle": "High",
-        "Garden Spider": "Low",
-        "Green Grasshopper": "Medium",
-        "Luna Moth": "Low",
-        "Common Dragonfly": "Low",
-        "Honey Bee": "Low",
-        "Paper Wasp": "Medium",
-        "Unknown Insect": "Unknown",
-        "Error Processing Image": "Unknown"
-    }
-    return severity_map.get(species, "Unknown")
+if __name__ == "__main__":
+    main()