Update src/streamlit_app.py

src/streamlit_app.py

@@ -5,13 +5,10 @@ from PIL import Image
 import numpy as np
 import io
 import requests
+import tempfile
 import os
 from typing import List, Tuple
 
-# Set cache directories to writable locations
-os.environ['TORCH_HOME'] = '/tmp/torch_cache'
-os.environ['HF_HOME'] = '/tmp/hf_cache'
-
 # Configure page
 st.set_page_config(
     page_title="CLIP Classifier",
@@ -23,92 +20,49 @@ st.set_page_config(
 def load_clip_model():
     """Load CLIP model and preprocessing function"""
     try:
-        # Ensure cache directories exist
-        os.makedirs('/tmp/torch_cache', exist_ok=True)
-        os.makedirs('/tmp/clip_models', exist_ok=True)
-        
         device = "cuda" if torch.cuda.is_available() else "cpu"
-        model, preprocess = clip.load("ViT-B/32", device=device, download_root="/tmp/clip_models")
+        model, preprocess = clip.load("ViT-B/32", device=device)
         return model, preprocess, device
     except Exception as e:
         st.error(f"Error loading CLIP model: {e}")
         return None, None, None
 
-def classify_input(model, preprocess, device, input_data, positive_prompts, negative_prompts, input_type="image"):
+def classify_input(model, preprocess, device, image_data, positive_prompts, negative_prompts):
     """
-    Classify input based on positive and negative prompts using CLIP
+    Classify image 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
-            if isinstance(input_data, str):  # URL
-                st.write("DEBUG: Processing URL image")
-                response = requests.get(input_data, timeout=10)
-                response.raise_for_status()
-                image = Image.open(io.BytesIO(response.content))
-            elif isinstance(input_data, bytes):  # Raw bytes
-                st.write("DEBUG: Processing bytes image")
-                image = Image.open(io.BytesIO(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
-            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)
-            
-            # Get features
-            with torch.no_grad():
-                image_features = model.encode_image(image_input)
-                text_features = model.encode_text(text_inputs)
-                
-                # Calculate similarities
-                similarities = (100.0 * image_features @ text_features.T).softmax(dim=-1)
-                similarities = similarities[0].cpu().numpy()
+        # Process image
+        if isinstance(image_data, str):  # URL
+            response = requests.get(image_data, timeout=10)
+            response.raise_for_status()
+            image = Image.open(io.BytesIO(response.content))
+        else:  # PIL Image or uploaded file
+            if hasattr(image_data, 'read'):
+                # Handle Streamlit UploadedFile
+                image_bytes = image_data.read()
+                image = Image.open(io.BytesIO(image_bytes))
+            else:
+                image = image_data
         
-        elif input_type == "text":
-            # Process text input
-            st.write("DEBUG: Processing text input")
-            input_text = clip.tokenize([input_data]).to(device)
+        # Convert to RGB if necessary
+        if image.mode != 'RGB':
+            image = image.convert('RGB')
+        
+        image_input = preprocess(image).unsqueeze(0).to(device)
+        
+        # Get features
+        with torch.no_grad():
+            image_features = model.encode_image(image_input)
+            text_features = model.encode_text(text_inputs)
             
-            with torch.no_grad():
-                input_features = model.encode_text(input_text)
-                text_features = model.encode_text(text_inputs)
-                
-                # Calculate similarities
-                similarities = (100.0 * input_features @ text_features.T).softmax(dim=-1)
-                similarities = similarities[0].cpu().numpy()
+            # Calculate similarities
+            similarities = (100.0 * image_features @ text_features.T).softmax(dim=-1)
+            similarities = similarities[0].cpu().numpy()
         
         # Calculate scores for positive and negative categories
         positive_scores = similarities[:len(positive_prompts)]
@@ -121,8 +75,6 @@ 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),
@@ -136,16 +88,15 @@ 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.markdown("### Define your own positive and negative prompts to classify images or text!")
+    st.markdown("### Define your own positive and negative prompts to classify images!")
     
     # Load model
-    model, preprocess, device = load_clip_model()
+    with st.spinner("Loading CLIP model..."):
+        model, preprocess, device = load_clip_model()
     
     if model is None:
         st.error("Failed to load CLIP model. Please check your installation.")
@@ -157,9 +108,6 @@ def main():
     with st.sidebar:
         st.header("Configuration")
         
-        # Input type selection
-        input_type = st.radio("Select input type:", ["Image", "Text"])
-        
         st.header("Define Prompts")
         
         # Positive prompts
@@ -191,93 +139,76 @@ def main():
     col1, col2 = st.columns([1, 1])
     
     with col1:
-        st.header("Input")
+        st.header("Input Image")
         
-        input_data = None
+        # Tabs for different input methods
+        tab1, tab2 = st.tabs(["Upload Image", "Image URL"])
         
-        if input_type == "Image":
-            # Image input options
-            image_option = st.radio("Choose image source:", ["Upload", "URL"])
-            
-            if image_option == "Upload":
-                uploaded_file = st.file_uploader(
-                    "Choose an image file",
-                    type=['png', 'jpg', 'jpeg', 'gif', 'bmp', 'webp'],
-                    help="Upload an image file to classify"
-                )
-                
-                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:
-                        # 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 displaying uploaded image: {e}")
-                        st.write(f"Error details: {str(e)}")
+        image_data = None
+        
+        with tab1:
+            # File uploader - simplified for HF Spaces
+            uploaded_file = st.file_uploader(
+                "Choose an image file",
+                type=['png', 'jpg', 'jpeg', 'gif', 'bmp', 'webp'],
+                help="Upload an image file to classify",
+                key="image_uploader"  # Add explicit key
+            )
             
-            else:  # URL
-                image_url = st.text_input("Enter image URL:", placeholder="https://example.com/image.jpg")
-                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}")
+            if uploaded_file is not None:
+                image_data = uploaded_file
+                # Display image
+                st.image(uploaded_file, caption=f"Uploaded: {uploaded_file.name}", use_column_width=True)
+                st.success("Image uploaded successfully!")
         
-        else:  # Text input
-            text_input = st.text_area(
-                "Enter text to classify:",
-                height=150,
-                placeholder="Type your text here...",
-                help="Enter the text you want to classify"
+        with tab2:
+            # URL input
+            image_url = st.text_input(
+                "Enter image URL:", 
+                placeholder="https://example.com/image.jpg",
+                help="Enter a direct link to an image"
             )
-            if text_input.strip():
-                input_data = text_input.strip()
-                st.text_area("Text to classify:", value=text_input, height=100, disabled=True)
+            
+            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))
+                            image_data = image_url
+                            st.image(image, caption="Image from URL", use_column_width=True)
+                            st.success("Image loaded successfully!")
+                    except Exception as e:
+                        st.error(f"Error loading image: {e}")
     
     with col2:
-        st.header("Results")
+        st.header("Classification 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}")
+        # Status check
+        ready_to_classify = (
+            image_data is not None and 
+            len(positive_prompts) > 0 and 
+            len(negative_prompts) > 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.")
-        elif not input_data:
-            st.info("Please provide input data to classify.")
+        elif image_data is None:
+            st.info("Please provide an image to classify.")
         else:
-            if st.button("Classify", type="primary", use_container_width=True):
+            st.success("Ready to classify!")
+            
+            if st.button("Classify Image", 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,
-                        input_type.lower()
+                        model, preprocess, device, image_data, 
+                        positive_prompts, negative_prompts
                     )
                 
                 if result:
-                    st.write("Classification successful!")
-                    
                     # Main classification result
                     classification = result['classification']
                     confidence = result['confidence']
@@ -287,10 +218,10 @@ def main():
                     st.markdown(f"### Classification: <span style='color: {color}'>{classification}</span>", 
                               unsafe_allow_html=True)
                     
-                    # Confidence and scores
-                    st.metric("Confidence", f"{confidence:.3f}")
-                    
-                    col_pos, col_neg = st.columns(2)
+                    # Metrics
+                    col_conf, col_pos, col_neg = st.columns(3)
+                    with col_conf:
+                        st.metric("Confidence", f"{confidence:.3f}")
                     with col_pos:
                         st.metric("Positive Score", f"{result['positive_score']:.3f}")
                     with col_neg:
@@ -300,37 +231,34 @@ def main():
                     st.subheader("Detailed Scores")
                     
                     # Positive prompts scores
-                    st.write("**Positive Prompts:**")
-                    for prompt, score in result['detailed_scores']['positive_prompts']:
-                        st.progress(float(score), text=f"{prompt}: {score:.3f}")
+                    with st.expander("Positive Prompts Scores", expanded=True):
+                        for prompt, score in result['detailed_scores']['positive_prompts']:
+                            st.progress(float(score), text=f"{prompt}: {score:.3f}")
                     
                     # Negative prompts scores
-                    st.write("**Negative Prompts:**")
-                    for prompt, score in result['detailed_scores']['negative_prompts']:
-                        st.progress(float(score), text=f"{prompt}: {score:.3f}")
+                    with st.expander("Negative Prompts Scores", expanded=True):
+                        for prompt, score in result['detailed_scores']['negative_prompts']:
+                            st.progress(float(score), text=f"{prompt}: {score:.3f}")
                 else:
-                    st.error("Classification failed. Check the debug messages above.")
+                    st.error("Classification failed. Please try again.")
     
     # Instructions
     with st.expander("How to use this app"):
         st.markdown("""
+        **Instructions:**
         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
-        3. **Provide Input**: 
-           - For images: Upload a file or provide a URL
-           - For text: Type or paste your text
-        4. **Classify**: Click the "Classify" button to see results
+        2. **Upload Image**: Use either the file uploader or paste an image URL
+        3. **Classify**: Click the "Classify Image" button to see results
         
-        **Examples of prompts:**
-        - **Image classification**: "happy dog, playful pet" vs "aggressive dog, angry animal"
-        - **Text sentiment**: "positive review, good experience" vs "negative review, bad experience"
-        - **Content moderation**: "safe content, family friendly" vs "inappropriate content, offensive material"
+        **Example prompts:**
+        - **Emotion detection**: "happy, smiling, joy" vs "sad, crying, anger"
+        - **Object detection**: "dog, puppy, canine" vs "cat, kitten, feline"
+        - **Content type**: "food, meal, cooking" vs "vehicle, car, transportation"
         
-        **Troubleshooting:**
-        - 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
+        **Tips for Hugging Face Spaces:**
+        - Use common image formats (JPG, PNG, WebP)
+        - For URLs, make sure they're publicly accessible
+        - Keep image sizes reasonable for faster processing
         """)
 
 if __name__ == "__main__":