app.py

"""
AI-Driven Home Security Target Detection System
Main Streamlit Application
"""

import streamlit as st
import cv2
import time
import numpy as np
from datetime import datetime
from typing import Optional, List
import threading
import queue

from video_source import VideoSource
from detection_engine import DetectionEngine
from mqtt_publisher import MQTTPublisher
from utils import (
    annotate_frame,
    crop_detection,
    encode_image_base64,
    resize_frame,
    draw_fps,
    format_detection_log,
    adjust_brightness
)


# Page configuration
st.set_page_config(
    page_title="AI Home Security Detection",
    page_icon="🔒",
    layout="wide",
    initial_sidebar_state="expanded"
)


def initialize_session_state():
    """Initialize Streamlit session state variables"""
    if 'running' not in st.session_state:
        st.session_state.running = False
    if 'video_source' not in st.session_state:
        st.session_state.video_source = None
    if 'detection_engine' not in st.session_state:
        st.session_state.detection_engine = None
    if 'mqtt_publisher' not in st.session_state:
        st.session_state.mqtt_publisher = None
    if 'detection_logs' not in st.session_state:
        st.session_state.detection_logs = []
    if 'last_detection' not in st.session_state:
        st.session_state.last_detection = None
    if 'last_detection_image' not in st.session_state:
        st.session_state.last_detection_image = None
    if 'stats' not in st.session_state:
        st.session_state.stats = {'total_detections': 0, 'target_detections': 0, 'frames_processed': 0}


def sidebar_controls():
    """Render sidebar controls"""
    st.sidebar.title("🔒 Security Detection System")
    st.sidebar.markdown("---")

    # Video Source Configuration
    st.sidebar.subheader("📹 Video Source")
    source_type = st.sidebar.selectbox(
        "Source Type",
        ["Webcam", "RTSP Stream", "Video File"],
        key="source_type"
    )

    if source_type == "Webcam":
        source_input = st.sidebar.number_input(
            "Camera Index",
            min_value=0,
            max_value=10,
            value=0,
            key="webcam_index"
        )
        source_type_key = "webcam"
    elif source_type == "RTSP Stream":
        source_input = st.sidebar.text_input(
            "RTSP URL",
            value="rtsp://example.com/stream",
            key="rtsp_url"
        )
        source_type_key = "rtsp"
    else:  # Video File
        source_input = st.sidebar.text_input(
            "File Path",
            value="/path/to/video.mp4",
            key="file_path"
        )
        source_type_key = "file"

    st.sidebar.markdown("---")

    # YOLO Model Configuration
    st.sidebar.subheader("🤖 Detection Model")
    model_option = st.sidebar.selectbox(
        "YOLO Model",
        ["yolov8n.pt", "yolov8s.pt", "yolov8m.pt", "yolov8l.pt", "Custom"],
        index=1,
        key="model_option"
    )

    if model_option == "Custom":
        model_path = st.sidebar.text_input(
            "Custom Model Path",
            value="path/to/custom_model.pt",
            key="custom_model_path"
        )
    else:
        model_path = model_option

    st.sidebar.markdown("---")

    # Detection Configuration
    st.sidebar.subheader("🎯 Target Configuration")

    # Load model to get available classes
    if st.sidebar.button("Load Model & Get Classes"):
        with st.spinner("Loading model..."):
            temp_engine = DetectionEngine(model_path)
            if temp_engine.load_model():
                st.session_state.available_classes = temp_engine.get_class_names()
                st.success(f"Model loaded! {len(st.session_state.available_classes)} classes available")
            else:
                st.error("Failed to load model")

    # Target class selection
    if 'available_classes' in st.session_state:
        target_class = st.sidebar.selectbox(
            "Target Class",
            st.session_state.available_classes,
            key="target_class"
        )
    else:
        target_class = st.sidebar.text_input(
            "Target Class (load model first for list)",
            value="person",
            key="target_class_manual"
        )

    confidence_threshold = st.sidebar.slider(
        "Confidence Threshold",
        min_value=0.0,
        max_value=1.0,
        value=0.5,
        step=0.05,
        key="confidence_threshold"
    )

    st.sidebar.markdown("---")

    # Display Options
    st.sidebar.subheader("🎨 Display Options")
    auto_brightness = st.sidebar.checkbox(
        "Auto-adjust brightness (for dark cameras)",
        value=True,
        key="auto_brightness"
    )

    st.sidebar.markdown("---")

    # MQTT Configuration
    st.sidebar.subheader("📡 MQTT Configuration")
    mqtt_broker = st.sidebar.text_input(
        "Broker Address",
        value="localhost",
        key="mqtt_broker"
    )
    mqtt_port = st.sidebar.number_input(
        "Port",
        min_value=1,
        max_value=65535,
        value=1883,
        key="mqtt_port"
    )
    mqtt_topic = st.sidebar.text_input(
        "Topic",
        value="irvine/home/bedroom/item",
        key="mqtt_topic"
    )
    mqtt_username = st.sidebar.text_input(
        "Username (optional)",
        value="",
        key="mqtt_username"
    )
    mqtt_password = st.sidebar.text_input(
        "Password (optional)",
        value="",
        type="password",
        key="mqtt_password"
    )
    include_image = st.sidebar.checkbox(
        "Include image in MQTT payload",
        value=False,
        key="include_image"
    )

    st.sidebar.markdown("---")

    return {
        'source_input': source_input,
        'source_type': source_type_key,
        'model_path': model_path,
        'target_class': target_class,
        'confidence_threshold': confidence_threshold,
        'auto_brightness': auto_brightness,
        'mqtt_broker': mqtt_broker,
        'mqtt_port': mqtt_port,
        'mqtt_topic': mqtt_topic,
        'mqtt_username': mqtt_username if mqtt_username else None,
        'mqtt_password': mqtt_password if mqtt_password else None,
        'include_image': include_image
    }


def main():
    """Main application"""
    initialize_session_state()

    # Title
    st.title("🔒 AI-Driven Home Security Detection System")
    st.markdown("Real-time object detection with MQTT notifications")

    # Get configuration from sidebar
    config = sidebar_controls()

    # Control buttons
    col1, col2, col3 = st.columns([1, 1, 3])

    with col1:
        if st.button("▶️ Start Detection", disabled=st.session_state.running):
            # Ensure clean state before starting
            st.session_state.video_source = None
            st.session_state.detection_engine = None
            st.session_state.mqtt_publisher = None
            st.session_state.detection_logs = []
            st.session_state.last_detection = None
            st.session_state.last_detection_image = None
            st.session_state.stats = {'total_detections': 0, 'target_detections': 0, 'frames_processed': 0}
            st.session_state.running = True
            st.rerun()

    with col2:
        if st.button("⏹️ Stop Detection", disabled=not st.session_state.running):
            st.session_state.running = False
            # Cleanup resources
            if st.session_state.video_source:
                st.session_state.video_source.close()
                st.session_state.video_source = None
            if st.session_state.mqtt_publisher:
                st.session_state.mqtt_publisher.disconnect()
                st.session_state.mqtt_publisher = None
            if st.session_state.detection_engine:
                st.session_state.detection_engine = None
            # Reset session state data
            st.session_state.detection_logs = []
            st.session_state.last_detection = None
            st.session_state.last_detection_image = None
            st.session_state.stats = {'total_detections': 0, 'target_detections': 0, 'frames_processed': 0}
            st.rerun()

    st.markdown("---")

    # Main content area
    if st.session_state.running:
        # Initialize components if needed
        if st.session_state.video_source is None:
            st.session_state.video_source = VideoSource()
            if not st.session_state.video_source.open(config['source_input'], config['source_type']):
                st.error("Failed to open video source!")
                st.session_state.running = False
                st.rerun()

        if st.session_state.detection_engine is None:
            st.session_state.detection_engine = DetectionEngine(
                config['model_path'],
                config['confidence_threshold']
            )
            if not st.session_state.detection_engine.load_model():
                st.error("Failed to load detection model!")
                st.session_state.running = False
                st.rerun()

        if st.session_state.mqtt_publisher is None:
            st.session_state.mqtt_publisher = MQTTPublisher(
                config['mqtt_broker'],
                config['mqtt_port'],
                config['mqtt_username'],
                config['mqtt_password']
            )
            if st.session_state.mqtt_publisher.connect():
                st.success("Connected to MQTT broker")
            else:
                st.warning("Failed to connect to MQTT broker. Will queue messages.")

        # Run detection loop
        run_detection_loop(config)

    else:
        # Show placeholder when not running
        st.info("👆 Configure settings in the sidebar and click 'Start Detection' to begin")

        # Show stats and last detection even when stopped
        if st.session_state.stats['frames_processed'] > 0:
            display_stats()

        if st.session_state.last_detection:
            display_last_detection()


def run_detection_loop(config):
    """Run the main detection loop"""
    # Create layout
    col1, col2 = st.columns([2, 1])

    with col1:
        st.subheader("📹 Live Feed")
        video_placeholder = st.empty()

    with col2:
        st.subheader("📊 Statistics")
        stats_placeholder = st.empty()

        st.subheader("📝 Detection Log")
        log_placeholder = st.empty()

    # Detection loop
    fps_counter = []
    frame_count = 0
    max_frames = 1000  # Process frames continuously

    video_source = st.session_state.video_source
    detection_engine = st.session_state.detection_engine
    mqtt_publisher = st.session_state.mqtt_publisher

    for _ in range(max_frames):
        if not st.session_state.running:
            break

        loop_start = time.time()

        # Read frame
        ret, frame = video_source.read()
        if not ret:
            st.error("Failed to read from video source")
            break

        frame_count += 1
        st.session_state.stats['frames_processed'] = frame_count

        # Apply brightness adjustment if enabled
        if config['auto_brightness']:
            frame = adjust_brightness(frame)

        # Run detection
        detections = detection_engine.detect(frame, target_classes=[config['target_class']])

        # Process detections
        for detection in detections:
            # Add timestamp
            detection['timestamp'] = datetime.now().isoformat()

            # Optionally include image
            if config['include_image']:
                cropped = crop_detection(frame, detection['bbox'])
                if cropped is not None:
                    detection['image'] = encode_image_base64(cropped)

            # Publish to MQTT
            mqtt_publisher.publish_detection(detection, config['mqtt_topic'])

            # Update logs
            st.session_state.detection_logs.insert(0, detection)
            if len(st.session_state.detection_logs) > 50:
                st.session_state.detection_logs.pop()

            # Save last detection
            st.session_state.last_detection = detection
            st.session_state.last_detection_image = crop_detection(frame, detection['bbox'])

        # Update stats
        st.session_state.stats.update(detection_engine.get_stats())

        # Annotate frame
        annotated = annotate_frame(frame, detections, [config['target_class']])

        # Calculate and draw FPS
        fps_counter.append(time.time() - loop_start)
        if len(fps_counter) > 30:
            fps_counter.pop(0)
        fps = 1.0 / (sum(fps_counter) / len(fps_counter))
        annotated = draw_fps(annotated, fps)

        # Resize for display
        display_frame = resize_frame(annotated, max_width=800)

        # Convert BGR to RGB for Streamlit
        display_frame = cv2.cvtColor(display_frame, cv2.COLOR_BGR2RGB)

        # Update displays
        video_placeholder.image(display_frame, channels="RGB", width="stretch")

        # Update stats
        with stats_placeholder.container():
            display_stats()

        # Update logs
        with log_placeholder.container():
            display_logs()

        # Small delay to prevent overwhelming the UI
        time.sleep(0.01)

    # After loop ends
    if st.session_state.running:
        st.session_state.running = False
        st.info("Detection loop ended")
        st.rerun()


def display_stats():
    """Display detection statistics"""
    stats = st.session_state.stats
    col1, col2, col3 = st.columns(3)

    with col1:
        st.metric("Frames Processed", stats['frames_processed'])

    with col2:
        st.metric("Total Detections", stats['total_detections'])

    with col3:
        st.metric("Target Detections", stats['target_detections'])


def display_logs():
    """Display detection logs"""
    if st.session_state.detection_logs:
        logs_text = "\n".join([
            format_detection_log(log)
            for log in st.session_state.detection_logs[:10]
        ])
        st.text_area("Recent Detections", logs_text, height=300, disabled=True)
    else:
        st.info("No detections yet")


def display_last_detection():
    """Display last detection details"""
    st.subheader("🎯 Last Detection")

    detection = st.session_state.last_detection
    image = st.session_state.last_detection_image

    col1, col2 = st.columns(2)

    with col1:
        if image is not None:
            image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            st.image(image_rgb, caption="Detection Crop", width="stretch")

    with col2:
        st.json(detection)


if __name__ == "__main__":
    main()