Update src/streamlit_app.py

src/streamlit_app.py

@@ -49,8 +49,11 @@ def classify_input(model, preprocess, device, input_data, positive_prompts, nega
                 response = requests.get(input_data, timeout=10)
                 response.raise_for_status()  # Raise an error for bad status codes
                 image = Image.open(io.BytesIO(response.content))
-            else:  # Uploaded file - this is the key fix
-                # For uploaded files, we need to read the bytes and convert to PIL Image
+            elif isinstance(input_data, bytes):  # Uploaded file as bytes
+                # Convert bytes directly to PIL Image
+                image = Image.open(io.BytesIO(input_data))
+            else:
+                # Fallback for other formats
                 if hasattr(input_data, 'read'):
                     # It's a file-like object (UploadedFile)
                     image_bytes = input_data.read()
@@ -180,14 +183,19 @@ def main():
                     help="Upload an image file to classify"
                 )
                 if uploaded_file is not None:  # More explicit check
-                    input_data = uploaded_file
                     try:
+                        # Read the file data once and store it
+                        file_bytes = uploaded_file.getvalue()  # This gets the bytes without moving file pointer
+                        input_data = file_bytes  # Store bytes instead of file object
+                        
                         # Display the uploaded image
-                        st.image(uploaded_file, caption=f"Uploaded: {uploaded_file.name}", use_column_width=True)
+                        st.image(file_bytes, caption=f"Uploaded: {uploaded_file.name}", use_column_width=True)
                         # Show file details
-                        st.info(f"File size: {len(uploaded_file.getvalue())} bytes")
+                        st.info(f"File size: {len(file_bytes)} bytes")
+                        st.success("✅ Image uploaded successfully!")
                     except Exception as e:
-                        st.error(f"Error displaying uploaded image: {e}")
+                        st.error(f"Error processing uploaded image: {e}")
+                        input_data = None
             
             else:  # URL
                 image_url = st.text_input("Enter image URL:", placeholder="https://example.com/image.jpg")
@@ -296,6 +304,9 @@ def main():
         st.subheader("Debug Information")
         st.write(f"Device: {device}")
         st.write(f"Input type: {input_type}")
+        st.write(f"Input data: {input_data is not None}")
+        if input_data and hasattr(input_data, '__len__'):
+            st.write(f"Input data length: {len(input_data)}")
         st.write(f"Input data type: {type(input_data)}")
         st.write(f"Positive prompts count: {len(positive_prompts) if positive_prompts else 0}")
         st.write(f"Negative prompts count: {len(negative_prompts) if negative_prompts else 0}")