Update src/streamlit_app.py

src/streamlit_app.py

@@ -39,34 +39,52 @@ def classify_input(model, preprocess, device, input_data, positive_prompts, nega
     Classify input based on positive and negative prompts using CLIP
     """
     try:
+        # Debug information
+        st.write(f"DEBUG: Input data type: {type(input_data)}")
+        st.write(f"DEBUG: Input type: {input_type}")
+        
         # Prepare text prompts
         all_prompts = positive_prompts + negative_prompts
         text_inputs = clip.tokenize(all_prompts).to(device)
         
         if input_type == "image":
-            # Process image - Fixed handling of uploaded files
+            # Process image
             if isinstance(input_data, str):  # URL
+                st.write("DEBUG: Processing URL image")
                 response = requests.get(input_data, timeout=10)
-                response.raise_for_status()  # Raise an error for bad status codes
+                response.raise_for_status()
                 image = Image.open(io.BytesIO(response.content))
-            elif isinstance(input_data, bytes):  # Uploaded file as bytes
-                # Convert bytes directly to PIL Image
+            elif isinstance(input_data, bytes):  # Raw bytes
+                st.write("DEBUG: Processing bytes 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()
-                    image = Image.open(io.BytesIO(image_bytes))
-                    # Reset file pointer for potential future reads
-                    input_data.seek(0)
-                else:
-                    # If it's already a PIL Image or other format
-                    image = input_data
+            else:  # UploadedFile object
+                st.write("DEBUG: Processing UploadedFile object")
+                # Try multiple methods to read the file
+                try:
+                    # Method 1: Use getvalue()
+                    if hasattr(input_data, 'getvalue'):
+                        image_bytes = input_data.getvalue()
+                        image = Image.open(io.BytesIO(image_bytes))
+                        st.write("DEBUG: Successfully read using getvalue()")
+                    # Method 2: Use read()
+                    elif hasattr(input_data, 'read'):
+                        input_data.seek(0)  # Reset to beginning
+                        image_bytes = input_data.read()
+                        image = Image.open(io.BytesIO(image_bytes))
+                        st.write("DEBUG: Successfully read using read()")
+                    else:
+                        st.error("DEBUG: Cannot read uploaded file")
+                        return None
+                except Exception as read_error:
+                    st.error(f"DEBUG: Error reading file: {read_error}")
+                    return None
             
-            # Convert to RGB if necessary (handles different image modes)
+            # Convert to RGB if necessary
             if image.mode != 'RGB':
                 image = image.convert('RGB')
+                st.write(f"DEBUG: Converted image from {image.mode} to RGB")
+            
+            st.write(f"DEBUG: Image size: {image.size}")
             
             image_input = preprocess(image).unsqueeze(0).to(device)
             
@@ -81,6 +99,7 @@ def classify_input(model, preprocess, device, input_data, positive_prompts, nega
         
         elif input_type == "text":
             # Process text input
+            st.write("DEBUG: Processing text input")
             input_text = clip.tokenize([input_data]).to(device)
             
             with torch.no_grad():
@@ -102,6 +121,8 @@ def classify_input(model, preprocess, device, input_data, positive_prompts, nega
         is_positive = positive_total > negative_total
         confidence = max(positive_total, negative_total)
         
+        st.write("DEBUG: Classification completed successfully")
+        
         return {
             'classification': 'Positive' if is_positive else 'Negative',
             'confidence': float(confidence),
@@ -115,10 +136,12 @@ def classify_input(model, preprocess, device, input_data, positive_prompts, nega
     
     except Exception as e:
         st.error(f"Error during classification: {e}")
+        import traceback
+        st.error(f"Traceback: {traceback.format_exc()}")
         return None
 
 def main():
-    st.title("🔍 CLIP-Based Custom Classifier")
+    st.title("CLIP-Based Custom Classifier")
     st.markdown("### Define your own positive and negative prompts to classify images or text!")
     
     # Load model
@@ -132,15 +155,15 @@ def main():
     
     # Sidebar for configuration
     with st.sidebar:
-        st.header("⚙️ Configuration")
+        st.header("Configuration")
         
         # Input type selection
         input_type = st.radio("Select input type:", ["Image", "Text"])
         
-        st.header("📝 Define Prompts")
+        st.header("Define Prompts")
         
         # Positive prompts
-        st.subheader("✅ Positive Prompts")
+        st.subheader("Positive Prompts")
         positive_prompts_text = st.text_area(
             "Enter positive prompts (one per line):",
             value="happy face\nsmiling person\njoyful expression\npositive emotion",
@@ -149,7 +172,7 @@ def main():
         )
         
         # Negative prompts
-        st.subheader("❌ Negative Prompts")
+        st.subheader("Negative Prompts")
         negative_prompts_text = st.text_area(
             "Enter negative prompts (one per line):",
             value="sad face\nangry person\nfrowning expression\nnegative emotion",
@@ -168,7 +191,7 @@ def main():
     col1, col2 = st.columns([1, 1])
     
     with col1:
-        st.header("📥 Input")
+        st.header("Input")
         
         input_data = None
         
@@ -179,42 +202,43 @@ def main():
             if image_option == "Upload":
                 uploaded_file = st.file_uploader(
                     "Choose an image file",
-                    type=['png', 'jpg', 'jpeg', 'gif', 'bmp', 'webp'],  # Added webp support
+                    type=['png', 'jpg', 'jpeg', 'gif', 'bmp', 'webp'],
                     help="Upload an image file to classify"
                 )
-                if uploaded_file is not None:  # More explicit check
+                
+                if uploaded_file is not None:
+                    st.write(f"File name: {uploaded_file.name}")
+                    st.write(f"File type: {uploaded_file.type}")
+                    st.write(f"File size: {uploaded_file.size} bytes")
+                    
+                    # Store the uploaded file directly
+                    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(file_bytes, caption=f"Uploaded: {uploaded_file.name}", use_column_width=True)
-                        # Show file details
-                        st.info(f"File size: {len(file_bytes)} bytes")
-                        st.success("✅ Image uploaded successfully!")
+                        # Display the uploaded image using the file object
+                        st.image(uploaded_file, caption=f"Uploaded: {uploaded_file.name}", use_column_width=True)
+                        st.success("Image uploaded successfully!")
                     except Exception as e:
-                        st.error(f"Error processing uploaded image: {e}")
-                        input_data = None
+                        st.error(f"Error displaying uploaded image: {e}")
+                        st.write(f"Error details: {str(e)}")
             
             else:  # URL
                 image_url = st.text_input("Enter image URL:", placeholder="https://example.com/image.jpg")
-                if image_url.strip():  # Check for non-empty URL
-                    try:
-                        # Add basic URL validation
-                        if not image_url.startswith(('http://', 'https://')):
-                            st.warning("Please enter a valid URL starting with http:// or https://")
-                        else:
+                if image_url.strip():
+                    if not image_url.startswith(('http://', 'https://')):
+                        st.warning("Please enter a valid URL starting with http:// or https://")
+                    else:
+                        try:
                             with st.spinner("Loading image..."):
                                 response = requests.get(image_url, timeout=10)
                                 response.raise_for_status()
                                 image = Image.open(io.BytesIO(response.content))
                                 input_data = image_url
                                 st.image(image, caption="Image from URL", use_column_width=True)
-                    except requests.exceptions.RequestException as e:
-                        st.error(f"Error loading image from URL: {e}")
-                    except Exception as e:
-                        st.error(f"Error processing image: {e}")
+                        except requests.exceptions.RequestException as e:
+                            st.error(f"Error loading image from URL: {e}")
+                        except Exception as e:
+                            st.error(f"Error processing image: {e}")
         
         else:  # Text input
             text_input = st.text_area(
@@ -228,16 +252,23 @@ def main():
                 st.text_area("Text to classify:", value=text_input, height=100, disabled=True)
     
     with col2:
-        st.header("📊 Results")
+        st.header("Results")
+        
+        # Show current status
+        st.write("Status Check:")
+        st.write(f"- Input data available: {input_data is not None}")
+        st.write(f"- Positive prompts: {len(positive_prompts) if positive_prompts else 0}")
+        st.write(f"- Negative prompts: {len(negative_prompts) if negative_prompts else 0}")
         
         # Check if we have all required inputs
         if not positive_prompts or not negative_prompts:
-            st.warning("⚠️ Please define both positive and negative prompts in the sidebar.")
+            st.warning("Please define both positive and negative prompts in the sidebar.")
         elif not input_data:
-            st.info("📝 Please provide input data to classify.")
+            st.info("Please provide input data to classify.")
         else:
-            if st.button("🚀 Classify", type="primary", use_container_width=True):
+            if st.button("Classify", type="primary", use_container_width=True):
                 with st.spinner("Classifying..."):
+                    st.write("Starting classification...")
                     result = classify_input(
                         model, preprocess, device, input_data, 
                         positive_prompts, negative_prompts,
@@ -245,6 +276,8 @@ def main():
                     )
                 
                 if result:
+                    st.write("Classification successful!")
+                    
                     # Main classification result
                     classification = result['classification']
                     confidence = result['confidence']
@@ -264,7 +297,7 @@ def main():
                         st.metric("Negative Score", f"{result['negative_score']:.3f}")
                     
                     # Detailed breakdown
-                    st.subheader("📈 Detailed Scores")
+                    st.subheader("Detailed Scores")
                     
                     # Positive prompts scores
                     st.write("**Positive Prompts:**")
@@ -276,10 +309,10 @@ def main():
                     for prompt, score in result['detailed_scores']['negative_prompts']:
                         st.progress(float(score), text=f"{prompt}: {score:.3f}")
                 else:
-                    st.error("Classification failed. Please check your input and try again.")
+                    st.error("Classification failed. Check the debug messages above.")
     
     # Instructions
-    with st.expander("ℹ️ How to use this app"):
+    with st.expander("How to use this app"):
         st.markdown("""
         1. **Define Prompts**: In the sidebar, enter your positive and negative prompts (one per line)
         2. **Choose Input Type**: Select whether you want to classify images or text
@@ -297,19 +330,8 @@ def main():
         - Make sure uploaded images are in supported formats (PNG, JPG, JPEG, GIF, BMP, WebP)
         - For URLs, ensure they start with http:// or https://
         - Check that both positive and negative prompts are defined
+        - Look at the debug messages for detailed error information
         """)
-    
-    # Debug information (can be removed in production)
-    if st.checkbox("Show debug info", help="Check this to see debug information"):
-        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}")
 
 if __name__ == "__main__":
     main()