Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| from PIL import Image | |
| import numpy as np | |
| 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", | |
| initial_sidebar_state="expanded" | |
| ) | |
| # Initialize session state | |
| 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() | |
| 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 | |
| """) | |
| # 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: | |
| 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) | |
| # 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: | |
| 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: | |
| 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 __name__ == "__main__": | |
| main() |