Update app.py

app.py

@@ -1,150 +1,401 @@
 import streamlit as st
-import torch
-from transformers import DetrImageProcessor, DetrForObjectDetection
-from PIL import Image, ImageDraw
+import cv2
+import tempfile
 import numpy as np
+from ultralytics import YOLO
+import plotly.graph_objects as go
+from collections import defaultdict
 import time
+import requests
+from PIL import Image
+import io
+import pandas as pd
 
-# Page config
+# Page configuration
 st.set_page_config(
-    page_title="Simple Object Detection",
-    page_icon="🔍",
+    page_title="People & Vehicle Counter",
+    page_icon="🚗",
     layout="wide"
 )
 
-# Title
-st.title("🔍 Simple Object Detection with DETR")
-st.markdown("Upload an image to detect objects using Facebook's DETR model")
+# Title and description
+st.title("🚗 People & Vehicle Counter")
+st.markdown("""
+Upload a video or provide a video URL to count people and vehicles in real-time using YOLOv8.
+This app is useful for traffic monitoring, retail analytics, and crowd management.
+""")
 
-# 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
-
-# Sidebar
+# Sidebar for settings
 with st.sidebar:
-    st.header("Settings")
+    st.header("⚙️ Settings")
     
     # Confidence threshold
     confidence = st.slider(
-        "Confidence Threshold",
+        "Detection Confidence",
         min_value=0.1,
-        max_value=0.99,
-        value=0.7,
-        step=0.05
+        max_value=1.0,
+        value=0.25,
+        step=0.05,
+        help="Higher values reduce false positives but might miss some objects"
     )
     
-    # 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!")
-                except Exception as e:
-                    st.error(f"Error loading model: {e}")
-    else:
-        st.success("✅ Model is loaded and ready!")
+    # Object classes to detect
+    st.subheader("Objects to Count")
+    count_person = st.checkbox("Person", value=True)
+    count_car = st.checkbox("Car", value=True)
+    count_bus = st.checkbox("Bus", value=False)
+    count_truck = st.checkbox("Truck", value=False)
+    count_motorcycle = st.checkbox("Motorcycle", value=False)
+    count_bicycle = st.checkbox("Bicycle", value=False)
+    
+    # Line intersection for counting
+    st.subheader("Counting Line")
+    show_line = st.checkbox("Show counting line", value=True)
+    line_position = st.slider(
+        "Line position (%)",
+        min_value=0,
+        max_value=100,
+        value=50,
+        help="Vertical line position for counting object crossings"
+    )
+    
+    # Processing options
+    st.subheader("Processing Options")
+    process_every_nth = st.slider(
+        "Process every Nth frame",
+        min_value=1,
+        max_value=10,
+        value=3,
+        help="Higher values speed up processing but reduce accuracy"
+    )
+    
+    max_frames = st.number_input(
+        "Maximum frames to process",
+        min_value=10,
+        max_value=1000,
+        value=200,
+        help="Limit processing for long videos"
+    )
 
-# Main content
-uploaded_file = st.file_uploader(
-    "Choose an image...",
-    type=['jpg', 'jpeg', 'png']
-)
+# Initialize session state
+if 'total_counts' not in st.session_state:
+    st.session_state.total_counts = defaultdict(int)
+if 'frame_counts' not in st.session_state:
+    st.session_state.frame_counts = []
+if 'processing_complete' not in st.session_state:
+    st.session_state.processing_complete = False
+if 'processed_video' not in st.session_state:
+    st.session_state.processed_video = None
+
+# COCO class names for YOLO (common objects)
+CLASS_NAMES = {
+    0: "person", 1: "bicycle", 2: "car", 3: "motorcycle", 
+    5: "bus", 7: "truck", 64: "chair"  # Note: YOLOv8 uses different indices
+}
 
-if uploaded_file is not None:
-    # Display original image
-    image = Image.open(uploaded_file).convert("RGB")
+# Map our checkboxes to class IDs
+def get_selected_classes():
+    selected_classes = []
+    class_mapping = {
+        "person": 0,
+        "bicycle": 2,
+        "car": 2,
+        "motorcycle": 3,
+        "bus": 5,
+        "truck": 7
+    }
+    
+    if count_person:
+        selected_classes.append(0)
+    if count_bicycle:
+        selected_classes.append(1)
+    if count_car:
+        selected_classes.append(2)
+    if count_motorcycle:
+        selected_classes.append(3)
+    if count_bus:
+        selected_classes.append(5)
+    if count_truck:
+        selected_classes.append(7)
+    
+    return selected_classes
+
+# Load YOLO model with caching
+@st.cache_resource
+def load_model():
+    return YOLO('yolov8n.pt')  # Using nano model for speed
+
+# Function to process video
+def process_video(video_path, selected_classes):
+    model = load_model()
+    
+    cap = cv2.VideoCapture(video_path)
+    fps = int(cap.get(cv2.CAP_PROP_FPS))
+    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    
+    # Initialize video writer
+    temp_output = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
+    output_path = temp_output.name
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
+    
+    # Reset counts
+    st.session_state.total_counts = defaultdict(int)
+    st.session_state.frame_counts = []
+    
+    # Progress bar
+    progress_bar = st.progress(0)
+    status_text = st.empty()
+    
+    frame_count = 0
+    processed_frames = 0
+    
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret or processed_frames >= max_frames:
+            break
+        
+        frame_count += 1
+        
+        # Process only every nth frame
+        if frame_count % process_every_nth != 0:
+            out.write(frame)
+            continue
+        
+        # Run YOLO inference
+        results = model(frame, conf=confidence, classes=selected_classes, verbose=False)
+        
+        # Draw results
+        annotated_frame = results[0].plot()
+        
+        # Count objects
+        frame_counts = defaultdict(int)
+        for box in results[0].boxes:
+            cls_id = int(box.cls[0])
+            class_name = CLASS_NAMES.get(cls_id, f"class_{cls_id}")
+            frame_counts[class_name] += 1
+            st.session_state.total_counts[class_name] += 1
+        
+        # Store frame counts for chart
+        st.session_state.frame_counts.append({
+            'frame': processed_frames,
+            **frame_counts
+        })
+        
+        # Draw counting line if enabled
+        if show_line:
+            line_x = int(width * line_position / 100)
+            cv2.line(annotated_frame, (line_x, 0), (line_x, height), (0, 255, 255), 2)
+        
+        # Add counter text
+        y_offset = 30
+        for obj_type, count in frame_counts.items():
+            cv2.putText(annotated_frame, 
+                       f"{obj_type}: {count}", 
+                       (10, y_offset), 
+                       cv2.FONT_HERSHEY_SIMPLEX, 
+                       0.7, (0, 255, 0), 2)
+            y_offset += 25
+        
+        # Write frame to output video
+        out.write(annotated_frame)
+        processed_frames += 1
+        
+        # Update progress
+        progress = min(processed_frames / max_frames, 1.0)
+        progress_bar.progress(progress)
+        status_text.text(f"Processing frame {processed_frames}/{max_frames}")
+    
+    cap.release()
+    out.release()
+    
+    st.session_state.processing_complete = True
+    st.session_state.processed_video = output_path
+    
+    return output_path
+
+# Function to download video from URL
+def download_video_from_url(url):
+    try:
+        response = requests.get(url, stream=True)
+        if response.status_code == 200:
+            temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
+            for chunk in response.iter_content(chunk_size=8192):
+                temp_file.write(chunk)
+            temp_file.close()
+            return temp_file.name
+        else:
+            st.error(f"Failed to download video. Status code: {response.status_code}")
+            return None
+    except Exception as e:
+        st.error(f"Error downloading video: {str(e)}")
+        return None
+
+# Main app layout
+tab1, tab2, tab3 = st.tabs(["📹 Video Input", "📊 Results", "ℹ️ About"])
+
+with tab1:
     col1, col2 = st.columns(2)
     
     with col1:
-        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()]
-                    
-                    # 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")
+        st.subheader("Upload Video")
+        uploaded_file = st.file_uploader(
+            "Choose a video file", 
+            type=['mp4', 'avi', 'mov', 'mkv']
+        )
+        
+        if uploaded_file is not None:
+            # Save uploaded file to temp location
+            tfile = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
+            tfile.write(uploaded_file.read())
+            video_path = tfile.name
+            
+            # Display video info
+            st.video(uploaded_file)
+            st.info(f"Uploaded: {uploaded_file.name}")
+    
+    with col2:
+        st.subheader("Video URL")
+        video_url = st.text_input(
+            "Enter video URL",
+            placeholder="https://example.com/video.mp4"
+        )
+        
+        if st.button("Load from URL") and video_url:
+            with st.spinner("Downloading video..."):
+                video_path = download_video_from_url(video_url)
+                if video_path:
+                    st.success("Video downloaded successfully!")
+                    # Display first frame
+                    cap = cv2.VideoCapture(video_path)
+                    ret, frame = cap.read()
+                    if ret:
+                        st.image(frame, caption="First frame of video", use_column_width=True)
+                    cap.release()
+    
+    # Process button
+    if ('video_path' in locals() and video_path) or ('uploaded_file' in locals() and uploaded_file):
+        if st.button("🚀 Start Counting", type="primary"):
+            selected_classes = get_selected_classes()
+            
+            if not selected_classes:
+                st.warning("Please select at least one object type to count.")
+            else:
+                with st.spinner("Processing video with YOLOv8..."):
+                    output_path = process_video(video_path, selected_classes)
+                    st.success("Processing complete!")
 
-# Instructions
-with st.expander("How to use this app"):
+with tab2:
+    if st.session_state.processing_complete:
+        col1, col2 = st.columns([2, 1])
+        
+        with col1:
+            st.subheader("Processed Video")
+            # Display processed video
+            video_file = open(st.session_state.processed_video, 'rb')
+            video_bytes = video_file.read()
+            st.video(video_bytes)
+            
+            # Download button
+            st.download_button(
+                label="📥 Download Processed Video",
+                data=video_bytes,
+                file_name="processed_video.mp4",
+                mime="video/mp4"
+            )
+        
+        with col2:
+            st.subheader("📈 Total Counts")
+            
+            # Display total counts
+            if st.session_state.total_counts:
+                for obj_type, count in st.session_state.total_counts.items():
+                    st.metric(label=obj_type.capitalize(), value=count)
+            else:
+                st.info("No objects detected")
+            
+            # Summary statistics
+            st.subheader("📊 Summary")
+            if st.session_state.frame_counts:
+                df = pd.DataFrame(st.session_state.frame_counts)
+                st.dataframe(df.tail(10), use_container_width=True)
+        
+        # Time series chart
+        st.subheader("📈 Objects Over Time")
+        if st.session_state.frame_counts:
+            df = pd.DataFrame(st.session_state.frame_counts)
+            
+            fig = go.Figure()
+            
+            # Add a trace for each object type
+            for column in df.columns:
+                if column != 'frame':
+                    fig.add_trace(go.Scatter(
+                        x=df['frame'],
+                        y=df[column],
+                        name=column.capitalize(),
+                        mode='lines+markers'
+                    ))
+            
+            fig.update_layout(
+                title="Object Counts per Frame",
+                xaxis_title="Frame Number",
+                yaxis_title="Count",
+                hovermode='x unified',
+                height=400
+            )
+            
+            st.plotly_chart(fig, use_container_width=True)
+    else:
+        st.info("Process a video first to see results here.")
+
+with tab3:
     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
+    ## About This App
+    
+    ### 🔧 Technology Stack
+    - **YOLOv8**: State-of-the-art object detection model
+    - **Streamlit**: Interactive web app framework
+    - **OpenCV**: Computer vision library for video processing
+    - **Plotly**: Interactive visualizations
+    
+    ### 📋 Features
+    1. **Multiple Input Sources**: Upload videos or use URLs
+    2. **Customizable Detection**: Select specific object classes
+    3. **Real-time Counting**: Track objects frame by frame
+    4. **Visual Analytics**: Interactive charts and statistics
+    5. **Export Results**: Download processed videos and data
+    
+    ### 🎯 Use Cases
+    - **Traffic Monitoring**: Count vehicles on roads
+    - **Retail Analytics**: Track customer movements
+    - **Crowd Management**: Monitor people density
+    - **Security**: Detect and count objects of interest
+    
+    ### ⚠️ Limitations
+    - Processing speed depends on video length and resolution
+    - Maximum 200 frames processed in this demo
+    - Accuracy depends on model confidence settings
+    
+    ### 📚 How to Use
+    1. Upload a video or provide a URL
+    2. Configure detection settings in the sidebar
+    3. Click "Start Counting"
+    4. View results in the Results tab
+    5. Download processed video and data
+    
+    ---
+    
+    **Note**: This app runs on Hugging Face Spaces with limited resources. 
+    For heavy processing, consider running locally with GPU support.
     """)
 
 # Footer
 st.markdown("---")
-st.markdown("Built with [DETR](https://huggingface.co/facebook/detr-resnet-50) • [Streamlit](https://streamlit.io)")
+st.markdown(
+    "<div style='text-align: center'>"
+    "Built with ❤️ using YOLOv8, Streamlit, and Hugging Face Spaces"
+    "</div>",
+    unsafe_allow_html=True
+)