Update app.py

app.py

@@ -5,35 +5,28 @@ import numpy as np
 from ultralytics import YOLO
 import plotly.graph_objects as go
 from collections import defaultdict
-import time
 import requests
 import pandas as pd
-from scipy.spatial import distance
 
 # --- Configuration & Initialization ---
 
-# Page configuration
 st.set_page_config(
     page_title="YOLOv8 Object Tracking & Counter",
     page_icon="🤖",
     layout="wide"
 )
 
-# Title and description
 st.title("🚦 Smart Object Traffic Analyzer (YOLOv8)")
 st.markdown("""
 A professional application for real-time **tracking and counting** of people and vehicles in video streams. 
 It uses YOLOv8 for detection and a simple tracking algorithm to count unique objects crossing a user-defined line.
 """)
 
-# COCO Class Names (Corrected for standard YOLOv8)
+# COCO Class Names (subset for demo)
 COCO_CLASS_NAMES = {
-    0: "person", 1: "bicycle", 2: "car", 3: "motorcycle", 4: "airplane", 
-    5: "bus", 6: "train", 7: "truck", 8: "boat", 9: "traffic light", 
-    # ... other classes
+    0: "person", 1: "bicycle", 2: "car", 3: "motorcycle", 5: "bus", 7: "truck"
 }
 
-# Mapping of checkboxed objects to their standard COCO Class IDs
 CLASS_MAPPING = {
     "Person": 0,
     "Bicycle": 1,
@@ -43,456 +36,190 @@ CLASS_MAPPING = {
     "Truck": 7,
 }
 
-# Initialize session state for tracking
 if 'processed_data' not in st.session_state:
     st.session_state.processed_data = {
         'total_counts': defaultdict(int),
         'frame_counts': [],
         'processed_video': None,
         'processing_complete': False,
-        'tracked_objects': {}, # Unique ID: {'class': str, 'last_centroid': (x, y), 'counted': bool}
+        'tracked_objects': {},
     }
 
-# --- Sidebar for Settings ---
+# --- Sidebar ---
 with st.sidebar:
     st.header("⚙️ Configuration Settings")
-    
-    # Model Selection
-    st.subheader("Model & Detection")
-    model_name = st.selectbox("Select YOLO Model", options=['yolov8n.pt', 'yolov8s.pt'], help="Nano (n) is fast, Small (s) is more accurate.")
-    
-    # Confidence threshold
-    confidence = st.slider(
-        "Detection Confidence Threshold",
-        min_value=0.1, max_value=1.0, value=0.40, step=0.05,
-        help="Minimum confidence to consider a detection valid."
-    )
-    
-    # Object classes to detect
+
+    model_name = st.selectbox("Select YOLO Model", options=['yolov8n.pt', 'yolov8s.pt'])
+    confidence = st.slider("Detection Confidence Threshold", 0.1, 1.0, 0.40, 0.05)
+
     st.subheader("Objects for Counting")
     selected_classes_ui = {}
-    for name, id in CLASS_MAPPING.items():
-        # Default True for Person and Car, False otherwise
+    for name in CLASS_MAPPING.keys():
         default_val = name in ["Person", "Car"]
         selected_classes_ui[name] = st.checkbox(name, value=default_val)
-    
-    # Line intersection for counting
+
     st.subheader("Counting Line Settings")
     show_line = st.checkbox("Show crossing line", value=True)
-    line_position = st.slider(
-        "Line Position (Vertical % from left)",
-        min_value=10, max_value=90, value=50,
-        help="Line position for counting objects that cross it."
-    )
-    
-    # Processing options
+    line_position = st.slider("Line Position (Vertical % from left)", 10, 90, 50)
+
     st.subheader("Performance Options")
-    process_every_nth = st.slider(
-        "Frame Skip (Process every Nth frame)",
-        min_value=1, max_value=10, value=2,
-        help="Higher values significantly speed up processing but reduce tracking smoothness."
-    )
-    
-    max_frames = st.number_input(
-        "Maximum Frames to Analyze",
-        min_value=10, max_value=5000, value=500,
-        help="Limits the processing duration for long videos. Set to a very high number (e.g., 99999) for full video."
-    )
+    process_every_nth = st.slider("Frame Skip (Process every Nth frame)", 1, 10, 2)
+    max_frames = st.number_input("Maximum Frames to Analyze", 10, 5000, 500)
 
-# --- Helper Functions ---
 
+# --- Helper Functions ---
 @st.cache_resource
 def load_model(model_path):
-    """Caches the YOLO model loading."""
     return YOLO(model_path)
 
 def get_selected_class_ids():
-    """Returns a list of COCO class IDs selected by the user."""
     return [CLASS_MAPPING[name] for name, is_selected in selected_classes_ui.items() if is_selected]
 
-# --- Core Processing Function (with simple tracking and crossing logic) ---
 
+# --- Core Processing ---
 def process_video(video_path, selected_class_ids, model_path):
-    """
-    Processes the video, performs object detection/tracking, and counts line crossings.
-    """
     model = load_model(model_path)
     cap = cv2.VideoCapture(video_path)
-    
-    # Video properties
+
     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))
     total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-    
+
     if total_frames > max_frames:
-        st.warning(f"Video is being processed for the first {max_frames} frames only (configurable in sidebar).")
-        
-    # Setup video writer (Using a smaller size for web-friendliness if possible)
+        st.warning(f"Processing limited to first {max_frames} frames.")
+
     temp_output = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
     output_path = temp_output.name
-    # mp4v or XVID is generally compatible. mp4v preferred for browser.
-    fourcc = cv2.VideoWriter_fourcc(*'mp4v') 
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
     out = cv2.VideoWriter(output_path, fourcc, fps / process_every_nth, (width, height))
-    
-    # Initialize state variables for the loop
+
     current_state = st.session_state.processed_data
     current_state['total_counts'] = defaultdict(int)
     current_state['frame_counts'] = []
-    current_state['tracked_objects'] = {} # ID: {'class': str, 'last_centroid': (x, y), 'counted': bool}
-    
-    # Define counting line
+    current_state['tracked_objects'] = {}
+
     line_x = int(width * line_position / 100)
-    
-    # UI Elements for progress
+
     progress_bar = st.progress(0)
     status_text = st.empty()
-    
+
     frame_count = 0
     processed_frames = 0
-    
+
     while cap.isOpened():
         ret, frame = cap.read()
-        
-        # Stop condition
         if not ret or processed_frames >= max_frames:
             break
-        
+
         frame_count += 1
-        
-        # Skip frames for performance (still write the frame for a continuous video)
         if frame_count % process_every_nth != 0:
-            # We don't write the skipped frame because we want the output video 
-            # to reflect the lower frame rate for smaller size and faster processing.
             continue
-        
-        # --- YOLO Detection ---
-        # NOTE: Using tracker="bytetrack.yaml" for better tracking. Requires ultralytics>=8.0.198
-        # However, for simplicity and dependency management, we will use simple centroid tracking.
+
         results = model.track(
-            frame, 
-            conf=confidence, 
-            classes=selected_class_ids, 
-            persist=True, 
-            tracker="bytetrack.yaml", # Use YOLO's built-in tracking!
+            frame,
+            conf=confidence,
+            classes=selected_class_ids,
+            persist=True,
+            tracker="bytetrack.yaml",
             verbose=False
         )
-        
+
         annotated_frame = frame.copy()
-        
-        # Current frame counts
         current_frame_counts = defaultdict(int)
-        
-        # --- Tracking and Counting Logic ---
+
         if results and results[0].boxes.id is not None:
             boxes = results[0].boxes.xyxy.cpu().numpy().astype(int)
             track_ids = results[0].boxes.id.cpu().numpy().astype(int)
             class_ids = results[0].boxes.cls.cpu().numpy().astype(int)
-            
+
             for box, track_id, class_id in zip(boxes, track_ids, class_ids):
                 x1, y1, x2, y2 = box
-                
-                # Calculate centroid
                 centroid_x = (x1 + x2) // 2
                 centroid_y = (y1 + y2) // 2
                 centroid = (centroid_x, centroid_y)
-                
+
                 class_name = COCO_CLASS_NAMES.get(class_id, "Unknown")
                 current_frame_counts[class_name] += 1
-                
-                # Update/Initialize tracked object
+
                 if track_id not in current_state['tracked_objects']:
-                    # New object detected
                     current_state['tracked_objects'][track_id] = {
-                        'class': class_name, 
-                        'last_centroid': centroid, 
+                        'class': class_name,
+                        'last_centroid': centroid,
                         'counted': False
                     }
                 else:
-                    # Existing object - Check for line crossing
                     obj_data = current_state['tracked_objects'][track_id]
                     prev_x = obj_data['last_centroid'][0]
-                    
+
                     if not obj_data['counted']:
-                        # Crossing logic: object crossed the line from one side to the other
-                        if (prev_x < line_x and centroid_x >= line_x) or \
-                           (prev_x > line_x and centroid_x <= line_x):
-                            
-                            # Object crossed the line!
+                        if (prev_x < line_x and centroid_x >= line_x) or (prev_x > line_x and centroid_x <= line_x):
                             current_state['total_counts'][class_name] += 1
-                            obj_data['counted'] = True # Count only once
-                            
-                    # Update the object's last known position
+                            obj_data['counted'] = True
+
                     obj_data['last_centroid'] = centroid
-                    
-                # Draw bounding box, track ID, and centroid
+
                 cv2.rectangle(annotated_frame, (x1, y1), (x2, y2), (255, 0, 0), 2)
                 cv2.circle(annotated_frame, centroid, 5, (0, 0, 255), -1)
-                
-                label = f"ID:{track_id} {class_name}"
-                cv2.putText(annotated_frame, label, (x1, y1 - 10), 
+                cv2.putText(annotated_frame, f"ID:{track_id} {class_name}", (x1, y1 - 10),
                             cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
-        
-        # --- Visualization & Logging ---
-        
-        # Draw counting line
+
         if show_line:
-            line_color = (0, 255, 255) # Cyan
-            cv2.line(annotated_frame, (line_x, 0), (line_x, height), line_color, 2)
-            
-            # Label for the line
-            cv2.putText(annotated_frame, "COUNTING LINE", (line_x + 5, 20), 
-                        cv2.FONT_HERSHEY_SIMPLEX, 0.6, line_color, 2)
-        
-        # Add total counter text
+            cv2.line(annotated_frame, (line_x, 0), (line_x, height), (0, 255, 255), 2)
+            cv2.putText(annotated_frame, "COUNTING LINE", (line_x + 5, 20),
+                        cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 255), 2)
+
         y_offset = 30
         for obj_type, count in current_state['total_counts'].items():
-            cv2.putText(annotated_frame, 
-                        f"TOTAL {obj_type.upper()}: {count}", 
-                        (width - 300, y_offset), 
-                        cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)
+            cv2.putText(annotated_frame, f"TOTAL {obj_type.upper()}: {count}",
+                        (width - 300, y_offset), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)
             y_offset += 35
 
-        # Store frame counts for chart (count of *objects in frame*, not crossings)
         frame_data = {'frame': processed_frames * process_every_nth}
         for name in CLASS_MAPPING.keys():
-             frame_data[name.lower()] = current_frame_counts.get(name.lower(), 0)
+            frame_data[name.lower()] = current_frame_counts.get(name.lower(), 0)
         current_state['frame_counts'].append(frame_data)
-        
-        # 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)
+
+        progress_bar.progress(min(processed_frames / max_frames, 1.0))
         status_text.text(f"Analyzing Frame {frame_count}/{total_frames} (Processed {processed_frames})")
-    
-    # --- Cleanup ---
+
     cap.release()
     out.release()
-    
-    # Update global state
+
     current_state['processing_complete'] = True
     current_state['processed_video'] = output_path
     st.session_state.processed_data = current_state
-    
+
     return output_path
 
-# Function to download video from URL
-@st.cache_data
-def download_video_from_url(url):
-    """Downloads video from URL to a temporary file."""
-    try:
-        st.info("Attempting to download video. This might take a moment...")
-        response = requests.get(url, stream=True, timeout=30)
-        response.raise_for_status() # Raise exception for bad status codes
-        
-        # Determine file extension (optional, but good practice)
-        content_type = response.headers.get('Content-Type', '')
-        suffix = '.mp4' if 'mp4' in content_type else '.mov'
-        
-        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
-        total_size = int(response.headers.get('Content-Length', 0))
-        
-        downloaded_size = 0
-        progress_placeholder = st.empty()
-        
-        for chunk in response.iter_content(chunk_size=8192):
-            temp_file.write(chunk)
-            downloaded_size += len(chunk)
-            if total_size > 0:
-                progress = downloaded_size / total_size
-                progress_placeholder.progress(progress, text=f"Downloading: {downloaded_size/(1024*1024):.2f}MB / {total_size/(1024*1024):.2f}MB")
-        
-        temp_file.close()
-        progress_placeholder.empty()
-        return temp_file.name
-        
-    except requests.exceptions.RequestException as e:
-        st.error(f"Failed to download video: {str(e)}. Check URL and file access.")
-        return None
-    except Exception as e:
-        st.error(f"An unexpected error occurred during download: {str(e)}")
-        return None
-
-# --- Main App Layout ---
-tab1, tab2, tab3 = st.tabs(["📹 Video Input", "📊 Analysis & Results", "ℹ️ Documentation"])
+
+# --- Main Layout ---
+tab1, tab2 = st.tabs(["📹 Video Input", "📊 Analysis & Results"])
 
 with tab1:
-    col1, col2 = st.columns(2)
     video_path = None
-    
-    with col1:
-        st.subheader("📁 Upload Video File")
-        uploaded_file = st.file_uploader(
-            "Choose a video file", 
-            type=['mp4', 'avi', 'mov', 'mkv'],
-            help="Supported video formats. Maximum recommended file size: 50MB."
-        )
-        
-        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
-            st.info(f"Video ready: {uploaded_file.name}")
-            st.video(uploaded_file)
-    
-    with col2:
-        st.subheader("🌐 Load from Video URL")
-        video_url = st.text_input(
-            "Enter public video URL (e.g., direct link to .mp4)",
-            placeholder="https://example.com/traffic.mp4"
-        )
-        
-        if st.button("🔗 Load from URL", use_container_width=True) and video_url:
-            video_path = download_video_from_url(video_url)
-            if video_path:
-                st.success("Video downloaded and ready for processing.")
-                # Try to display a frame if possible
-                try:
-                    cap = cv2.VideoCapture(video_path)
-                    ret, frame = cap.read()
-                    if ret:
-                        st.image(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB), caption="Video Preview", use_column_width=True)
-                    cap.release()
-                except Exception:
-                    st.warning("Could not display video preview.")
-
-    st.markdown("---")
-    
-    # Process button logic
-    if video_path:
-        if st.button("🚀 START TRACKING AND COUNTING", type="primary", use_container_width=True):
-            selected_class_ids = get_selected_class_ids()
-            
-            if not selected_class_ids:
-                st.error("Please select at least one object type to count in the sidebar.")
-            else:
-                try:
-                    with st.spinner(f"Analyzing video with {model_name}..."):
-                        process_video(video_path, selected_class_ids, model_name)
-                    st.success("Analysis complete! See results in the 'Analysis & Results' tab.")
-                    # Automatically switch to results tab on completion? (Streamlit doesn't natively support this well)
-                except Exception as e:
-                    st.error(f"An error occurred during video processing: {e}")
-                    # Optionally print traceback
-                    # import traceback; st.code(traceback.format_exc())
-    else:
-        st.info("Upload a video or provide a URL to begin.")
+    uploaded_file = st.file_uploader("Upload Video", type=['mp4', 'avi', 'mov', 'mkv'])
+    if uploaded_file:
+        tfile = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
+        tfile.write(uploaded_file.read())
+        video_path = tfile.name
+        st.video(uploaded_file)
+
+    if video_path and st.button("🚀 START TRACKING AND COUNTING"):
+        selected_class_ids = get_selected_class_ids()
+        if not selected_class_ids:
+            st.error("Please select at least one object type.")
+        else:
+            with st.spinner(f"Analyzing video with {model_name}..."):
+                process_video(video_path, selected_class_ids, model_name)
+            st.success("Analysis complete! See results tab.")
 
 with tab2:
     data = st.session_state.processed_data
     if data['processing_complete']:
-        st.header("Results Summary")
-        
-        col1, col2 = st.columns([2, 1])
-        
-        with col1:
-            st.subheader("🎥 Analyzed Video Output")
-            # Display processed video
-            with open(data['processed_video'], 'rb') as video_file:
-                video_bytes = video_file.read()
-            st.video(video_bytes)
-            
-            # Download button
-            st.download_button(
-                label="📥 Download Analyzed Video (MP4)",
-                data=video_bytes,
-                file_name="analyzed_tracking_video.mp4",
-                mime="video/mp4",
-                use_container_width=True
-            )
-            
-        with col2:
-            st.subheader("✅ Object Crossing Totals")
-            # Display total counts
-            if data['total_counts']:
-                for obj_type, count in data['total_counts'].items():
-                    st.metric(label=f"Total {obj_type.capitalize()} Crossed", value=count)
-            else:
-                st.info("No objects crossed the counting line in the analyzed section.")
-
-        st.subheader("📊 Object Presence Over Frames")
-        if data['frame_counts']:
-            df = pd.DataFrame(data['frame_counts']).fillna(0)
-            
-            # Time series chart (Plotly)
-            fig = go.Figure()
-            
-            # Add a trace for each object type (columns except 'frame')
-            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="Count of Objects Present in Frame",
-                xaxis_title="Frame Number",
-                yaxis_title="Count of Objects (Instance Count)",
-                hovermode='x unified',
-                height=400
-            )
-            
-            st.plotly_chart(fig, use_container_width=True)
-            
-            st.subheader("Data Export")
-            st.dataframe(df.tail(10), use_container_width=True, height=200)
-            
-            csv = df.to_csv(index=False).encode('utf-8')
-            st.download_button(
-                label="⬇️ Download Frame-by-Frame Data (CSV)",
-                data=csv,
-                file_name="object_count_data.csv",
-                mime="text/csv",
-            )
-            
-        else:
-            st.warning("No tracking data available. Process a video first.")
-            
-    else:
-        st.info("Process a video in the 'Video Input' tab to view analysis results.")
-
-with tab3:
-    st.header("Documentation: Smart Object Traffic Analyzer")
-    st.markdown("""
-    This application utilizes cutting-edge computer vision techniques for object tracking and crossing counting.
-    
-    ### 🔑 Core Technology
-    
-    * **YOLOv8**: The primary model for high-accuracy, real-time object detection. We recommend the `yolov8n.pt` (Nano) for speed in browser-based demos.
-    * **ByteTrack**: Used via the `ultralytics` package for robust object tracking, assigning a unique ID to each detected instance across frames.
-    * **Streamlit**: Provides the interactive, professional front-end interface.
-    
-    ---
-    
-    ### ⚙️ How Crossing Counting Works
-    
-    Unlike simple detection counters which add to a total for every frame an object is visible, this app counts **unique object crossings** of a vertical line:
-    
-    1.  **Tracking**: YOLOv8's integrated tracker assigns a persistent **Track ID** to each object (`person`, `car`, etc.).
-    2.  **Centroid Calculation**: The center-point (centroid) of the object's bounding box is calculated for every frame.
-    3.  **Crossing Logic**: The system monitors the object's horizontal position relative to the **Counting Line**. An object is counted **once** when its centroid moves from one side of the line (e.g., left) to the other (e.g., right).
-    
-    This ensures an accurate count of unique events, not redundant detections.
-    
-    ### 🚀 Deployment on Hugging Face Spaces
-    
-    This script is optimized for deployment:
-    
-    * **Caching (`@st.cache_resource`)**: The YOLO model is loaded only once, saving significant time.
-    * **Dependency List**: You will need a `requirements.txt` file in your Space with the following key libraries:
-        ```text
-        streamlit
-        ultralytics
-        opencv-python-headless
-        numpy
-        plotly
-        pandas
-        requests
-        scipy # for distance calculation, though not strictly needed with bytetrack
+        st.subheader("🎥 Processed Video")
+        with open(data['processed_video'], 'rb') as video_file:
+            video_bytes = video_file.read()