dashboard.py

import streamlit as st
import cv2
import numpy as np
from ultralytics import YOLO
import time
import pandas as pd
import plotly.express as px
from datetime import datetime, timedelta
import json

# Page configuration
st.set_page_config(
    page_title="🏪 Retail Theft Detection Dashboard",
    page_icon="🚨",
    layout="wide"
)

# Initialize session state
if 'alerts' not in st.session_state:
    st.session_state.alerts = []
if 'detections' not in st.session_state:
    st.session_state.detections = []

@st.cache_resource
def load_model():
    """Load YOLO model (cached)"""
    return YOLO('yolov8n.pt')

def analyze_frame(frame, model):
    """Analyze a single frame"""
    zones = {
        'entrance': (0, 0, 160, 480),
        'merchandise': (160, 0, 320, 300),  
        'checkout': (160, 300, 480, 480),
        'exit': (320, 0, 640, 480)
    }
    
    results = model(frame, verbose=False)
    detections = []
    alerts = []
    
    for detection in results[0].boxes:
        if detection.conf < 0.5:
            continue
            
        cls_id = int(detection.cls)
        class_name = model.names[cls_id]
        
        if class_name == 'person':
            bbox = detection.xyxy[0].cpu().numpy()
            confidence = float(detection.conf)
            
            # Check which zones person is in
            x1, y1, x2, y2 = bbox
            center_x = (x1 + x2) / 2
            center_y = (y1 + y2) / 2
            
            current_zones = []
            for zone_name, (zx1, zy1, zx2, zy2) in zones.items():
                if zx1 <= center_x <= zx2 and zy1 <= center_y <= zy2:
                    current_zones.append(zone_name)
            
            detection_data = {
                'timestamp': datetime.now(),
                'confidence': confidence,
                'zones': current_zones,
                'center_x': center_x,
                'center_y': center_y
            }
            detections.append(detection_data)
            
            # Generate alerts
            if 'merchandise' in current_zones and len(current_zones) == 1:
                alerts.append({
                    'timestamp': datetime.now(),
                    'type': 'loitering_merchandise',
                    'description': 'Person in merchandise area',
                    'severity': 'medium',
                    'confidence': confidence
                })
            
            if 'exit' in current_zones and 'checkout' not in current_zones:
                alerts.append({
                    'timestamp': datetime.now(),
                    'type': 'bypass_checkout', 
                    'description': 'Person bypassing checkout',
                    'severity': 'high',
                    'confidence': confidence
                })
    
    return detections, alerts, zones

def draw_frame_with_zones(frame, detections, zones, alerts):
    """Draw zones and detections on frame"""
    # Zone colors
    zone_colors = {
        'entrance': (0, 255, 0),
        'merchandise': (0, 255, 255),
        'checkout': (255, 255, 0),
        'exit': (255, 0, 0)
    }
    
    # Draw zones
    for zone_name, (x1, y1, x2, y2) in zones.items():
        color = zone_colors[zone_name]
        cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
        cv2.putText(frame, zone_name.upper(), (x1, y1-10), 
                   cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)
    
    # Draw detections
    for detection in detections:
        center_x = int(detection['center_x'])
        center_y = int(detection['center_y'])
        confidence = detection['confidence']
        
        # Draw person center point
        cv2.circle(frame, (center_x, center_y), 10, (0, 255, 0), -1)
        cv2.putText(frame, f"Person: {confidence:.2f}", 
                   (center_x-50, center_y-15),
                   cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
    
    # Draw alerts on frame
    y_offset = 30
    for alert in alerts[-3:]:  # Show last 3 alerts
        severity_colors = {
            'low': (0, 255, 255),
            'medium': (0, 165, 255),
            'high': (0, 0, 255)
        }
        color = severity_colors.get(alert['severity'], (255, 255, 255))
        
        alert_text = f"⚠️ {alert['description']}"
        cv2.putText(frame, alert_text, (10, y_offset), 
                   cv2.FONT_HERSHEY_SIMPLEX, 0.7, color, 2)
        y_offset += 30
    
    return frame

def main():
    st.title("🏪 Retail Theft Detection Dashboard")
    st.markdown("---")
    
    # Load model
    model = load_model()
    
    # Sidebar
    with st.sidebar:
        st.header("🎮 Controls")
        
        camera_source = st.selectbox("📷 Camera Source", [0, 1, 2, 3], index=0)
        confidence_threshold = st.slider("🎯 Confidence Threshold", 0.1, 1.0, 0.5)
        
        # Snapshot button
        take_snapshot = st.button("📸 Take Snapshot", use_container_width=True)
        
        st.markdown("---")
        st.header("📊 Live Stats")
        
        total_alerts = len(st.session_state.alerts)
        recent_alerts = len([a for a in st.session_state.alerts 
                           if a['timestamp'] > datetime.now() - timedelta(minutes=5)])
        
        st.metric("Total Alerts", total_alerts)
        st.metric("Recent Alerts (5min)", recent_alerts)
        
        if st.session_state.alerts:
            high_alerts = len([a for a in st.session_state.alerts if a['severity'] == 'high'])
            medium_alerts = len([a for a in st.session_state.alerts if a['severity'] == 'medium'])
            
            st.metric("🚨 High Priority", high_alerts)
            st.metric("⚠️ Medium Priority", medium_alerts)
    
    # Main content
    col1, col2 = st.columns([2, 1])
    
    with col1:
        st.header("📹 Live Detection")
        
        # Create placeholder for video
        video_placeholder = st.empty()
        
        # Camera capture
        cap = cv2.VideoCapture(camera_source)
        
        if cap.isOpened():
            ret, frame = cap.read()
            if ret:
                # Analyze frame
                detections, alerts, zones = analyze_frame(frame, model)
                
                # Store data
                st.session_state.detections.extend(detections)
                st.session_state.alerts.extend(alerts)
                
                # Keep only recent data
                if len(st.session_state.detections) > 100:
                    st.session_state.detections = st.session_state.detections[-100:]
                if len(st.session_state.alerts) > 50:
                    st.session_state.alerts = st.session_state.alerts[-50:]
                
                # Draw annotations
                annotated_frame = draw_frame_with_zones(frame, detections, zones, alerts)
                
                # Convert to RGB and display
                frame_rgb = cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB)
                video_placeholder.image(frame_rgb, channels="RGB", use_column_width=True)
                
                # Save snapshot if requested
                if take_snapshot:
                    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                    filename = f"snapshot_{timestamp}.jpg"
                    cv2.imwrite(filename, annotated_frame)
                    st.success(f"📸 Snapshot saved: {filename}")
            
            cap.release()
        else:
            video_placeholder.error("❌ Could not access camera")
    
    with col2:
        st.header("🚨 Recent Alerts")
        
        if st.session_state.alerts:
            # Show last 10 alerts
            recent_alerts = sorted(st.session_state.alerts, 
                                 key=lambda x: x['timestamp'], reverse=True)[:10]
            
            for i, alert in enumerate(recent_alerts):
                time_str = alert['timestamp'].strftime('%H:%M:%S')
                
                if alert['severity'] == 'high':
                    st.error(f"🚨 **{alert['description']}**\n\n🕒 {time_str} | 🎯 {alert['confidence']:.2f}")
                elif alert['severity'] == 'medium':
                    st.warning(f"⚠️ **{alert['description']}**\n\n🕒 {time_str} | 🎯 {alert['confidence']:.2f}")
                else:
                    st.info(f"ℹ️ **{alert['description']}**\n\n🕒 {time_str} | 🎯 {alert['confidence']:.2f}")
        else:
            st.info("No alerts yet. System monitoring...")
    
    # Analytics section
    if st.session_state.detections:
        st.markdown("---")
        st.header("📈 Analytics")
        
        tab1, tab2, tab3 = st.tabs(["📊 Zone Activity", "⏰ Alert Timeline", "💾 Export"])
        
        with tab1:
            # Zone activity chart
            zone_data = []
            for detection in st.session_state.detections:
                for zone in detection['zones']:
                    zone_data.append({'zone': zone, 'timestamp': detection['timestamp']})
            
            if zone_data:
                zone_df = pd.DataFrame(zone_data)
                zone_counts = zone_df['zone'].value_counts().reset_index()
                zone_counts.columns = ['zone', 'count']
                
                fig = px.bar(zone_counts, x='zone', y='count', 
                           title="Person Detections by Zone",
                           color='zone')
                st.plotly_chart(fig, use_container_width=True)
        
        with tab2:
            # Alert timeline
            if st.session_state.alerts:
                alerts_df = pd.DataFrame(st.session_state.alerts)
                alerts_df['minute'] = alerts_df['timestamp'].dt.floor('min')
                
                timeline = alerts_df.groupby(['minute', 'severity']).size().reset_index(name='count')
                
                fig = px.line(timeline, x='minute', y='count', color='severity',
                            title="Alert Timeline")
                st.plotly_chart(fig, use_container_width=True)
        
        with tab3:
            # Export data
            col1, col2 = st.columns(2)
            
            with col1:
                if st.button("📥 Download Detection Data"):
                    df = pd.DataFrame(st.session_state.detections)
                    csv = df.to_csv(index=False)
                    st.download_button(
                        "💾 Download CSV",
                        csv,
                        f"detections_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv",
                        "text/csv"
                    )
            
            with col2:
                if st.button("📥 Download Alert Data"):
                    # Prepare alerts for JSON export
                    alerts_export = []
                    for alert in st.session_state.alerts:
                        alert_copy = alert.copy()
                        alert_copy['timestamp'] = alert_copy['timestamp'].isoformat()
                        alerts_export.append(alert_copy)
                    
                    json_data = json.dumps(alerts_export, indent=2)
                    st.download_button(
                        "💾 Download JSON",
                        json_data,
                        f"alerts_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json",
                        "application/json"
                    )
    
    # Auto-refresh button
    st.markdown("---")
    if st.button("🔄 Refresh Detection", use_container_width=True):
        st.rerun()

if __name__ == "__main__":
    main()