Update streamlit_app.py

streamlit_app.py

@@ -4,11 +4,15 @@ import torch
 from torchvision import models, transforms
 import json
 import os
-import tempfile
+import io
+import numpy as np
+from streamlit_webrtc import webrtc_streamer, WebRtcMode, RTCConfiguration
+import av
+import cv2
 
 # Configure Streamlit
 st.set_page_config(
-    page_title="Butterfly Identifier/liblika tuvastaja",
+    page_title="Butterfly Identifier/Liblikamaja ID",
     page_icon="🦋",
     layout="wide"
 )
@@ -48,45 +52,123 @@ transform = transforms.Compose([
     transforms.ToTensor(),
 ])
 
-st.title("🦋 Butterfly Identifier")
-st.write("Upload a butterfly image and I'll tell you what species it is!")
+def predict_butterfly(image):
+    """Predict butterfly species from image"""
+    if image is None:
+        return None, None
+    
+    # Convert to PIL Image if needed
+    if isinstance(image, np.ndarray):
+        image = Image.fromarray(image)
+    
+    # Preprocess
+    input_tensor = transform(image).unsqueeze(0)
+    
+    # Predict
+    with torch.no_grad():
+        output = model(input_tensor)
+        probabilities = torch.nn.functional.softmax(output[0], dim=0)
+        confidence, pred = torch.max(probabilities, 0)
+        predicted_class = class_names[pred.item()]
+        
+    return predicted_class, confidence.item()
 
-# Alternative file upload with better error handling
-try:
+# Video frame callback for live camera
+class VideoProcessor:
+    def __init__(self):
+        self.prediction_text = ""
+        self.frame_count = 0
+        
+    def recv(self, frame):
+        img = frame.to_ndarray(format="bgr24")
+        
+        # Only process every 30th frame for performance
+        self.frame_count += 1
+        if self.frame_count % 30 == 0:
+            # Convert BGR to RGB
+            rgb_img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+            predicted_class, confidence = predict_butterfly(rgb_img)
+            
+            if predicted_class and confidence > 0.7:  # Only show if confident
+                self.prediction_text = f"{predicted_class} ({confidence:.2f})"
+        
+        # Draw prediction on frame
+        if self.prediction_text:
+            cv2.putText(img, self.prediction_text, (10, 30), 
+                       cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
+        
+        return av.VideoFrame.from_ndarray(img, format="bgr24")
+
+# Streamlit UI
+st.title("🦋 Butterfly Identifier / Liblikamaja ID")
+st.write("Identify butterflies using your camera or by uploading an image!")
+
+# Create tabs for different input methods
+tab1, tab2 = st.tabs(["📷 Live Camera", "📁 Upload Image"])
+
+with tab1:
+    st.header("Live Camera Identification")
+    st.write("Point your camera at a butterfly and get real-time identification!")
+    
+    # WebRTC configuration
+    RTC_CONFIGURATION = RTCConfiguration({
+        "iceServers": [{"urls": ["stun:stun.l.google.com:19302"]}]
+    })
+    
+    # Camera streamer
+    webrtc_ctx = webrtc_streamer(
+        key="butterfly-camera",
+        mode=WebRtcMode.SENDRECV,
+        rtc_configuration=RTC_CONFIGURATION,
+        video_processor_factory=VideoProcessor,
+        media_stream_constraints={"video": True, "audio": False},
+        async_processing=True,
+    )
+    
+    st.info("💡 Tips for best results:")
+    st.write("- Hold the camera steady")
+    st.write("- Ensure good lighting")
+    st.write("- Get close to the butterfly")
+    st.write("- Wait for the green prediction text to appear")
+
+with tab2:
+    st.header("Upload Image")
+    st.write("Upload a clear photo of a butterfly for identification")
+    
     uploaded_file = st.file_uploader(
         "Choose an image...", 
         type=["jpg", "jpeg", "png"],
-        help="Upload a clear photo of a butterfly",
-        key="butterfly_image"
+        help="Upload a clear photo of a butterfly"
     )
     
     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
-        input_tensor = transform(image).unsqueeze(0)
-
-        # Predict
-        with torch.no_grad():
-            output = model(input_tensor)
-            _, pred = torch.max(output, 1)
-            predicted_class = class_names[pred.item()]
-
-        st.success(f"**Prediction: {predicted_class}**")
-
-        if predicted_class in butterfly_info:
-            st.info(butterfly_info[predicted_class]["description"])
-        
-        # Clean up temporary file
-        os.unlink(tmp_file_path)
+        try:
+            # Read file directly into memory
+            image_bytes = uploaded_file.read()
+            image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
             
-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.")
+            col1, col2 = st.columns(2)
+            
+            with col1:
+                st.image(image, caption="Uploaded Image", use_column_width=True)
+            
+            with col2:
+                predicted_class, confidence = predict_butterfly(image)
+                
+                if predicted_class:
+                    st.success(f"**Prediction: {predicted_class}**")
+                    st.info(f"Confidence: {confidence:.2%}")
+                    
+                    if predicted_class in butterfly_info:
+                        st.write("**Species Information:**")
+                        st.write(butterfly_info[predicted_class]["description"])
+                
+        except Exception as e:
+            st.error(f"Error processing image: {str(e)}")
+
+# Add footer with instructions
+st.markdown("---")
+st.markdown("### How to use:")
+st.markdown("1. **Live Camera**: Click 'START' to begin live identification")
+st.markdown("2. **Upload Image**: Choose a butterfly photo from your device")
+st.markdown("3. **Best Results**: Use clear, well-lit photos with the butterfly clearly visible")