Update app.py

app.py

@@ -3,450 +3,148 @@ import torch
 from transformers import DetrImageProcessor, DetrForObjectDetection
 from PIL import Image, ImageDraw
 import numpy as np
-from collections import Counter
-import cv2
 import time
-import tempfile
-import os
 
-# Set page config
+# Page config
 st.set_page_config(
-    page_title="Object Detection Playground",
+    page_title="Simple Object Detection",
     page_icon="🔍",
     layout="wide"
 )
 
-# Custom CSS
-st.markdown("""
-<style>
-    .main-header {
-        font-size: 2.5rem;
-        color: #1E88E5;
-        text-align: center;
-        margin-bottom: 1rem;
-        font-weight: 700;
-    }
-    .sub-header {
-        font-size: 1.2rem;
-        color: #666;
-        text-align: center;
-        margin-bottom: 2rem;
-    }
-    .stat-box {
-        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
-        color: white;
-        padding: 1.5rem;
-        border-radius: 10px;
-        margin: 0.5rem 0;
-        box-shadow: 0 4px 6px rgba(0,0,0,0.1);
-    }
-    .metric-card {
-        background: white;
-        padding: 1rem;
-        border-radius: 10px;
-        border-left: 5px solid #1E88E5;
-        box-shadow: 0 2px 4px rgba(0,0,0,0.1);
-        margin: 0.5rem 0;
-    }
-    .stButton > button {
-        background: linear-gradient(135deg, #1E88E5 0%, #0D47A1 100%);
-        color: white;
-        border: none;
-        padding: 0.5rem 2rem;
-        border-radius: 5px;
-        font-weight: 600;
-    }
-    .stButton > button:hover {
-        background: linear-gradient(135deg, #0D47A1 0%, #1565C0 100%);
-        transform: translateY(-2px);
-        transition: all 0.3s ease;
-    }
-    .confidence-slider {
-        margin: 1rem 0;
-    }
-    .model-info-box {
-        background: #f8f9fa;
-        padding: 1rem;
-        border-radius: 10px;
-        border: 1px solid #dee2e6;
-    }
-</style>
-""", unsafe_allow_html=True)
+# Title
+st.title("🔍 Simple Object Detection with DETR")
+st.markdown("Upload an image to detect objects using Facebook's DETR model")
 
-@st.cache_resource(show_spinner=True)
-def load_model():
-    """Load and cache the DETR model"""
-    try:
-        with st.spinner("Loading DETR model (first time may take a minute)..."):
-            # Load processor and model
-            processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50", revision="no_timm")
-            model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50", revision="no_timm")
-            model.eval()  # Set to evaluation mode
-        return processor, model
-    except Exception as e:
-        st.error(f"Failed to load model: {str(e)}")
-        return None, None
+# Initialize model in session state
+if 'model_loaded' not in st.session_state:
+    st.session_state.model_loaded = False
+    st.session_state.processor = None
+    st.session_state.model = None
 
-def process_image(image, processor, model, confidence_threshold):
-    """Process image and return detections"""
-    try:
-        # Convert to RGB if needed
-        if image.mode != 'RGB':
-            image = image.convert('RGB')
-        
-        # Process image
-        inputs = processor(images=image, return_tensors="pt")
-        
-        # Run inference
-        with torch.no_grad():
-            outputs = model(**inputs)
-        
-        # Process outputs
-        target_sizes = torch.tensor([image.size[::-1]])  # [height, width]
-        results = processor.post_process_object_detection(
-            outputs, 
-            target_sizes=target_sizes, 
-            threshold=0.01  # Low threshold, we'll filter later
-        )[0]
-        
-        # Filter by confidence threshold
-        mask = results["scores"] >= confidence_threshold
-        filtered_results = {
-            "scores": results["scores"][mask],
-            "labels": results["labels"][mask],
-            "boxes": results["boxes"][mask]
-        }
-        
-        return filtered_results
-    except Exception as e:
-        st.error(f"Error processing image: {str(e)}")
-        return None
-
-def draw_detections(image, results, processor, model):
-    """Draw bounding boxes on image"""
-    try:
-        # Create a copy of the image
-        img_copy = image.copy()
-        draw = ImageDraw.Draw(img_copy)
-        
-        # Color palette for different classes
-        colors = [
-            (255, 0, 0), (0, 255, 0), (0, 0, 255),
-            (255, 255, 0), (255, 0, 255), (0, 255, 255),
-            (255, 128, 0), (128, 0, 255), (0, 128, 255)
-        ]
-        
-        # Draw each detection
-        for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
-            # Get box coordinates
-            xmin, ymin, xmax, ymax = box.tolist()
-            
-            # Get label name
-            label_id = label.item()
-            label_name = model.config.id2label[label_id]
-            
-            # Choose color based on label
-            color = colors[label_id % len(colors)]
-            
-            # Draw rectangle
-            draw.rectangle([xmin, ymin, xmax, ymax], outline=color, width=3)
-            
-            # Create label text
-            label_text = f"{label_name}: {score:.2f}"
-            
-            # Draw label background
-            text_bbox = draw.textbbox((xmin, ymin), label_text)
-            draw.rectangle(text_bbox, fill=color)
-            
-            # Draw text
-            draw.text((xmin, ymin), label_text, fill="white")
-        
-        return img_copy
-    except Exception as e:
-        st.error(f"Error drawing detections: {str(e)}")
-        return image
-
-def get_statistics(results, model):
-    """Calculate and return detection statistics"""
-    if results is None or len(results["scores"]) == 0:
-        return {
-            "total_objects": 0,
-            "avg_confidence": 0,
-            "class_distribution": {},
-            "detected_classes": []
-        }
-    
-    # Count objects per class
-    class_counts = Counter()
-    confidences = []
-    
-    for score, label in zip(results["scores"], results["labels"]):
-        label_name = model.config.id2label[label.item()]
-        class_counts[label_name] += 1
-        confidences.append(score.item())
-    
-    # Prepare statistics
-    stats = {
-        "total_objects": len(results["scores"]),
-        "avg_confidence": np.mean(confidences) if confidences else 0,
-        "max_confidence": max(confidences) if confidences else 0,
-        "min_confidence": min(confidences) if confidences else 0,
-        "class_distribution": dict(class_counts),
-        "detected_classes": list(class_counts.keys())
-    }
-    
-    return stats
-
-def main():
-    # Header
-    st.markdown('<h1 class="main-header">🔍 Object Detection Playground</h1>', unsafe_allow_html=True)
-    st.markdown('<p class="sub-header">Upload images and detect objects with DETR (Detection Transformer)</p>', unsafe_allow_html=True)
-    
-    # Initialize session state
-    if 'processed_image' not in st.session_state:
-        st.session_state.processed_image = None
-    if 'detection_results' not in st.session_state:
-        st.session_state.detection_results = None
+# Sidebar
+with st.sidebar:
+    st.header("Settings")
+    
+    # Confidence threshold
+    confidence = st.slider(
+        "Confidence Threshold",
+        min_value=0.1,
+        max_value=0.99,
+        value=0.7,
+        step=0.05
+    )
     
-    # Sidebar
-    with st.sidebar:
-        st.markdown("### ⚙️ Configuration")
-        
-        # Model info
-        with st.expander("ℹ️ Model Information", expanded=True):
-            st.markdown("""
-            **Model:** facebook/detr-resnet-50  
-            **Architecture:** DETR (End-to-End Object Detection)  
-            **Backbone:** ResNet-50  
-            **Training Data:** COCO 2017  
-            **Classes:** 91 categories
-            """)
-        
-        # Confidence threshold
-        st.markdown("### 🎯 Confidence Settings")
-        confidence_threshold = st.slider(
-            "Detection Threshold",
-            min_value=0.0,
-            max_value=1.0,
-            value=0.7,
-            step=0.05,
-            help="Objects with confidence below this threshold will be filtered out"
-        )
-        
-        # Display options
-        st.markdown("### 🎨 Display Options")
-        show_labels = st.checkbox("Show labels on image", value=True)
-        show_confidence = st.checkbox("Show confidence scores", value=True)
-        
-        # Performance options
-        st.markdown("### ⚡ Performance")
-        use_gpu = st.checkbox("Use GPU if available", value=True)
-        
-        # Load model button
-        st.markdown("---")
-        if st.button("🔄 Load/Reload Model", use_container_width=True):
-            with st.spinner("Loading model..."):
-                st.cache_resource.clear()
-                processor, model = load_model()
-                if processor and model:
+    # Load model button
+    if not st.session_state.model_loaded:
+        if st.button("Load Model", type="primary"):
+            with st.spinner("Loading DETR model..."):
+                try:
+                    st.session_state.processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
+                    st.session_state.model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
+                    st.session_state.model_loaded = True
                     st.success("Model loaded successfully!")
-    
-    # Main content area
-    col1, col2 = st.columns([2, 1])
+                except Exception as e:
+                    st.error(f"Error loading model: {e}")
+    else:
+        st.success("✅ Model is loaded and ready!")
+
+# Main content
+uploaded_file = st.file_uploader(
+    "Choose an image...",
+    type=['jpg', 'jpeg', 'png']
+)
+
+if uploaded_file is not None:
+    # Display original image
+    image = Image.open(uploaded_file).convert("RGB")
+    col1, col2 = st.columns(2)
     
     with col1:
-        st.markdown("### 📤 Upload Image")
-        
-        # File uploader
-        uploaded_file = st.file_uploader(
-            "Choose an image file",
-            type=['jpg', 'jpeg', 'png', 'bmp'],
-            help="Supported formats: JPG, PNG, BMP"
-        )
-        
-        # Or use sample images
-        st.markdown("### 📸 Try Sample Images")
-        sample_col1, sample_col2, sample_col3 = st.columns(3)
-        
-        sample_images = {
-            "Street": "https://raw.githubusercontent.com/ultralytics/yolov5/master/data/images/bus.jpg",
-            "Office": "https://raw.githubusercontent.com/ultralytics/yolov5/master/data/images/zidane.jpg",
-            "Animals": "https://images.unsplash.com/photo-1564349683136-77e08dba1ef7?w=800&h=600&fit=crop"
-        }
-        
-        with sample_col1:
-            if st.button("Street Scene", use_container_width=True):
-                st.session_state.sample_image = "street"
-        with sample_col2:
-            if st.button("Office Scene", use_container_width=True):
-                st.session_state.sample_image = "office"
-        with sample_col3:
-            if st.button("Animals", use_container_width=True):
-                st.session_state.sample_image = "animals"
-        
-        # Display uploaded or sample image
-        image = None
-        if uploaded_file is not None:
-            image = Image.open(uploaded_file)
-            st.image(image, caption="Uploaded Image", use_column_width=True)
-        elif 'sample_image' in st.session_state:
-            # Note: In HuggingFace Spaces, you might need to handle sample images differently
-            # For now, we'll use placeholder
-            st.info("Sample images require internet connection. Please upload your own image for testing.")
-            image = None
-        
-        # Process button
-        if image is not None:
-            if st.button("🚀 Detect Objects", type="primary", use_container_width=True):
-                with st.spinner("Processing image..."):
-                    # Load model
-                    processor, model = load_model()
+        st.image(image, caption="Original Image", use_column_width=True)
+    
+    if st.session_state.model_loaded and st.button("Detect Objects"):
+        with st.spinner("Detecting objects..."):
+            try:
+                # Process image
+                processor = st.session_state.processor
+                model = st.session_state.model
+                
+                inputs = processor(images=image, return_tensors="pt")
+                
+                with torch.no_grad():
+                    outputs = model(**inputs)
+                
+                # Convert outputs
+                target_sizes = torch.tensor([image.size[::-1]])
+                results = processor.post_process_object_detection(
+                    outputs, 
+                    target_sizes=target_sizes, 
+                    threshold=confidence
+                )[0]
+                
+                # Draw boxes
+                draw = ImageDraw.Draw(image)
+                colors = ["red", "green", "blue", "yellow", "purple", "orange"]
+                
+                detected_objects = []
+                
+                for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
+                    box = [round(i, 2) for i in box.tolist()]
+                    label_name = model.config.id2label[label.item()]
                     
-                    if processor and model:
-                        # Process image
-                        results = process_image(image, processor, model, confidence_threshold)
-                        
-                        if results:
-                            # Draw detections
-                            if show_labels:
-                                annotated_image = draw_detections(image, results, processor, model)
-                            else:
-                                annotated_image = image
-                            
-                            # Get statistics
-                            stats = get_statistics(results, model)
-                            
-                            # Store in session state
-                            st.session_state.processed_image = annotated_image
-                            st.session_state.detection_results = results
-                            st.session_state.stats = stats
-                            st.session_state.model = model
-                            
-                            st.success(f"Detected {stats['total_objects']} objects!")
-    
-    # Display results in right column
-    with col2:
-        st.markdown("### 📊 Detection Statistics")
-        
-        if 'stats' in st.session_state and st.session_state.stats:
-            stats = st.session_state.stats
-            
-            # Metrics
-            metric_col1, metric_col2 = st.columns(2)
-            with metric_col1:
-                st.metric("Total Objects", stats['total_objects'])
-                st.metric("Avg Confidence", f"{stats['avg_confidence']:.1%}")
-            with metric_col2:
-                st.metric("Max Confidence", f"{stats['max_confidence']:.1%}")
-                st.metric("Unique Classes", len(stats['detected_classes']))
-            
-            # Class distribution
-            if stats['class_distribution']:
-                st.markdown("#### 🏷️ Detected Classes")
-                for class_name, count in sorted(stats['class_distribution'].items(), key=lambda x: x[1], reverse=True):
-                    st.markdown(f"**{class_name}**: {count} objects")
-            
-            # Confidence histogram
-            if st.session_state.detection_results and len(st.session_state.detection_results["scores"]) > 0:
-                st.markdown("#### 📈 Confidence Distribution")
-                confidences = [s.item() for s in st.session_state.detection_results["scores"]]
-                hist_values = np.histogram(confidences, bins=10, range=(0, 1))[0]
-                st.bar_chart(hist_values)
-        
-        else:
-            st.info("No detection results yet. Upload an image and click 'Detect Objects'.")
-    
-    # Display processed image below
-    if st.session_state.processed_image is not None:
-        st.markdown("---")
-        st.markdown("### 🖼️ Detection Results")
-        
-        result_col1, result_col2 = st.columns([3, 1])
-        
-        with result_col1:
-            st.image(
-                st.session_state.processed_image,
-                caption=f"Detected Objects (Threshold: {confidence_threshold})",
-                use_column_width=True
-            )
-        
-        with result_col2:
-            # Download button
-            if st.session_state.processed_image:
-                from io import BytesIO
-                buffered = BytesIO()
-                st.session_state.processed_image.save(buffered, format="PNG")
-                st.download_button(
-                    label="💾 Download Result",
-                    data=buffered.getvalue(),
-                    file_name="detection_result.png",
-                    mime="image/png",
-                    use_container_width=True
-                )
-            
-            # Reset button
-            if st.button("🔄 Clear Results", use_container_width=True):
-                st.session_state.processed_image = None
-                st.session_state.detection_results = None
-                if 'stats' in st.session_state:
-                    del st.session_state.stats
-                st.rerun()
-    
-    # Footer with model capabilities
-    st.markdown("---")
-    
-    # Model capabilities section
-    st.markdown("### 🎯 What Can DETR Detect?")
-    
-    capabilities_col1, capabilities_col2, capabilities_col3 = st.columns(3)
-    
-    with capabilities_col1:
-        st.markdown("""
-        **👥 People & Animals**
-        - person
-        - dog, cat, bird
-        - horse, sheep, cow
-        - bear, zebra, giraffe
-        """)
-    
-    with capabilities_col2:
-        st.markdown("""
-        **🚗 Vehicles**
-        - car, truck, bus
-        - bicycle, motorcycle
-        - airplane, boat
-        - train
-        """)
-    
-    with capabilities_col3:
-        st.markdown("""
-        **🏠 Everyday Objects**
-        - chair, sofa, bed
-        - dining table
-        - tv, laptop, mouse
-        - bottle, cup, fork
-        """)
-    
-    # Tips and instructions
-    with st.expander("💡 Tips for Best Results"):
-        st.markdown("""
-        1. **Use clear images** with good lighting
-        2. **Start with threshold 0.7** and adjust as needed
-        3. **For crowded scenes**, increase threshold to reduce false positives
-        4. **For small objects**, decrease threshold to catch more detections
-        5. **Images with multiple objects** work best with DETR
-        6. **Allow model to load** on first run (takes about 30 seconds)
-        """)
-    
-    # Footer
-    st.markdown("---")
-    st.markdown(
-        "<div style='text-align: center; color: #666;'>"
-        "Object Detection Playground • Powered by <a href='https://huggingface.co/facebook/detr-resnet-50' target='_blank'>DETR</a> • "
-        "Built with ❤️ using Streamlit"
-        "</div>",
-        unsafe_allow_html=True
-    )
+                    # Draw rectangle
+                    color = colors[label.item() % len(colors)]
+                    draw.rectangle(box, outline=color, width=3)
+                    
+                    # Add label
+                    label_text = f"{label_name}: {score:.2f}"
+                    draw.text((box[0], box[1]), label_text, fill=color)
+                    
+                    detected_objects.append((label_name, score.item()))
+                
+                # Display results
+                with col2:
+                    st.image(image, caption="Detected Objects", use_column_width=True)
+                
+                # Show statistics
+                st.subheader("📊 Detection Results")
+                
+                if detected_objects:
+                    col_stats1, col_stats2, col_stats3 = st.columns(3)
+                    
+                    with col_stats1:
+                        st.metric("Objects Found", len(detected_objects))
+                    
+                    with col_stats2:
+                        avg_conf = np.mean([score for _, score in detected_objects])
+                        st.metric("Average Confidence", f"{avg_conf:.1%}")
+                    
+                    with col_stats3:
+                        st.metric("Unique Classes", len(set([label for label, _ in detected_objects])))
+                    
+                    # Show details
+                    st.subheader("Detected Objects:")
+                    for label, score in detected_objects:
+                        st.write(f"- **{label}** (confidence: {score:.1%})")
+                else:
+                    st.warning("No objects detected above the confidence threshold.")
+                    
+            except Exception as e:
+                st.error(f"Error during detection: {e}")
+else:
+    st.info("👈 Please upload an image and load the model from the sidebar")
+
+# Instructions
+with st.expander("How to use this app"):
+    st.markdown("""
+    1. **Load the model** using the button in the sidebar
+    2. **Upload an image** (JPG, PNG formats)
+    3. **Adjust confidence threshold** if needed
+    4. **Click 'Detect Objects'** to run detection
+    5. **View results** and detected objects
+    """)
 
-if __name__ == "__main__":
-    main()
+# Footer
+st.markdown("---")
+st.markdown("Built with [DETR](https://huggingface.co/facebook/detr-resnet-50) • [Streamlit](https://streamlit.io)")