Update model_utils.py

model_utils.py

@@ -1,176 +1,41 @@
-import streamlit as st
+# 3. model_utils.py
+# Model management (loading, prediction, and species information)
+from transformers import ViTForImageClassification
 from PIL import Image
-from model_utils import BugClassifier, get_severity_prediction
-
-# Page configuration
-st.set_page_config(
-    page_title="Bug-O-Scope 🔍🐞",
-    page_icon="🔍",
-    layout="wide"
-)
-
-# 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"])
-
-    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:
-            # Load and display image
-            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)
-                    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:
-            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:
-        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 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:
+import torch
+from dataset_utils import load_species_descriptions
+
+class BugClassifier:
+    def __init__(self, model_path="google/vit-base-patch16-224"):
+        self.model = ViTForImageClassification.from_pretrained(model_path)
+        self.model.eval()
+        self.labels = [
+            "Seven-spotted Ladybug", "Monarch Butterfly", "Carpenter Ant",
+            "Japanese Beetle", "Garden Spider", "Green Grasshopper",
+            "Luna Moth", "Common Dragonfly", "Honey Bee", "Paper Wasp"
+        ]
+        # Dynamically load species descriptions
+        self.species_descriptions = load_species_descriptions()
+
+    def predict(self, image):
         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.")
-                
+            processor = ViTImageProcessor.from_pretrained("google/vit-base-patch16-224")
+            inputs = processor(images=image, return_tensors="pt")
+            with torch.no_grad():
+                outputs = self.model(**inputs)
+                probabilities = torch.nn.functional.softmax(outputs.logits, dim=1)
+                confidence, predicted_idx = probabilities.max(dim=1)
+                confidence = confidence.item() * 100
+                predicted_label = self.labels[predicted_idx.item()]
+
+                if confidence < 30:
+                    return "Unknown Insect", confidence
+
+                return predicted_label, confidence
         except Exception as e:
-            st.error(f"Error comparing images: {str(e)}")
-            st.info("Please try uploading different images or try again.")
+            return "Error Processing Image", 0.0
 
-if __name__ == "__main__":
-    main()
+    def get_species_info(self, species):
+        return self.species_descriptions.get(
+            species, "Information not available. Consider updating your dataset for this species."
+        )