Update streamlit_app.py

streamlit_app.py

@@ -3,7 +3,8 @@ from PIL import Image
 import torch
 from torchvision import models, transforms
 import json
-import os  # Add this import
+import os
+import tempfile
 
 # Configure Streamlit
 st.set_page_config(
@@ -27,12 +28,10 @@ except:
 def load_model():
     MODEL_PATH = "butterfly_classifier.pth"
     
-    # Check if model exists locally
     if not os.path.exists(MODEL_PATH):
         st.error("Model file not found. Please upload butterfly_classifier.pth to your space.")
         return None
     
-    # Load the model
     model = models.resnet18(pretrained=False)
     model.fc = torch.nn.Linear(model.fc.in_features, len(class_names))
     model.load_state_dict(torch.load(MODEL_PATH, map_location="cpu"))
@@ -41,7 +40,6 @@ def load_model():
 
 model = load_model()
 
-# Add a check to ensure model loaded successfully
 if model is None:
     st.stop()
 
@@ -53,16 +51,23 @@ transform = transforms.Compose([
 st.title("🦋 Butterfly Identifier")
 st.write("Upload a butterfly image and I'll tell you what species it is!")
 
-# Use a different file uploader approach
-uploaded_file = st.file_uploader(
-    "Choose an image...", 
-    type=["jpg", "jpeg", "png"],
-    help="Upload a clear photo of a butterfly"
-)
-
-if uploaded_file is not None:
-    try:
-        image = Image.open(uploaded_file).convert("RGB")
+# Alternative file upload with better error handling
+try:
+    uploaded_file = st.file_uploader(
+        "Choose an image...", 
+        type=["jpg", "jpeg", "png"],
+        help="Upload a clear photo of a butterfly",
+        key="butterfly_image"
+    )
+    
+    if uploaded_file is not None:
+        # Save uploaded file to temporary location
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as tmp_file:
+            tmp_file.write(uploaded_file.read())
+            tmp_file_path = tmp_file.name
+        
+        # Load image from temporary file
+        image = Image.open(tmp_file_path).convert("RGB")
         st.image(image, caption="Uploaded Image", use_column_width=True)
 
         # Preprocess
@@ -78,6 +83,10 @@ if uploaded_file is not None:
 
         if predicted_class in butterfly_info:
             st.info(butterfly_info[predicted_class]["description"])
+        
+        # Clean up temporary file
+        os.unlink(tmp_file_path)
             
-    except Exception as e:
-        st.error(f"Error processing image: {str(e)}")
+except Exception as e:
+    st.error(f"Error with file upload: {str(e)}")
+    st.info("If you continue to see this error, try refreshing the page or using a different browser.")