midlajvalappil/Real-time_Object_Detection_with_YOLO

1"""2Real-time Object Detection with YOLO - Streamlit Application3Main application file for the web interface.4"""5 6import streamlit as st7import cv28import numpy as np9import time10from PIL import Image11import sys12import os13 14# Add src directory to path15sys.path.append(os.path.join(os.path.dirname(__file__), 'src'))16 17from detection.yolo_detector import YOLODetector18from detection.webcam_capture import WebcamCapture, get_available_cameras19from utils.config import app_config20from utils.helpers import format_detection_info, log_detection_summary21 22# Page configuration23st.set_page_config(24    page_title=app_config.ui.page_title,25    page_icon="🎯",26    layout="wide",27    initial_sidebar_state="expanded"28)29 30# Initialize session state31if 'detector' not in st.session_state:32    st.session_state.detector = None33if 'webcam' not in st.session_state:34    st.session_state.webcam = None35if 'is_running' not in st.session_state:36    st.session_state.is_running = False37if 'frame_count' not in st.session_state:38    st.session_state.frame_count = 039if 'detection_stats' not in st.session_state:40    st.session_state.detection_stats = {}41 42def initialize_detector(model_name: str, confidence_threshold: float):43    """Initialize the YOLO detector."""44    try:45        with st.spinner(f"Loading {model_name} model..."):46            detector = YOLODetector(model_name, confidence_threshold)47            if detector.model is not None:48                st.session_state.detector = detector49                st.success(f"✅ Model {model_name} loaded successfully!")50                return True51            else:52                st.error(f"❌ Failed to load model {model_name}")53                return False54    except Exception as e:55        st.error(f"❌ Error loading model: {str(e)}")56        return False57 58def initialize_webcam(camera_index: int):59    """Initialize the webcam."""60    try:61        webcam = WebcamCapture(62            camera_index=camera_index,63            width=app_config.camera.frame_width,64            height=app_config.camera.frame_height65        )66        if webcam.start_capture():67            st.session_state.webcam = webcam68            st.success(f"✅ Camera {camera_index} initialized successfully!")69            return True70        else:71            st.error(f"❌ Failed to initialize camera {camera_index}")72            return False73    except Exception as e:74        st.error(f"❌ Error initializing camera: {str(e)}")75        return False76 77def main():78    """Main application function."""79    80    # Title and description81    st.title("🎯 Real-time Object Detection with YOLO")82    st.markdown("Detect objects in real-time using your webcam and YOLO models.")83    84    # Sidebar for controls85    with st.sidebar:86        st.header("⚙️ Settings")87        88        # Model selection89        st.subheader("🤖 Model Configuration")90        selected_model = st.selectbox(91            "Select YOLO Model",92            app_config.available_models,93            index=app_config.available_models.index(app_config.detection.model_name),94            help="Choose the YOLO model variant. Nano is fastest, X is most accurate."95        )96        97        # Show model description98        model_info = app_config.get_model_info(selected_model)99        st.info(model_info['description'])100        101        # Confidence threshold102        confidence_threshold = st.slider(103            "Confidence Threshold",104            min_value=0.1,105            max_value=1.0,106            value=app_config.detection.confidence_threshold,107            step=0.05,108            help="Minimum confidence score for detections"109        )110        111        # Camera selection112        st.subheader("📹 Camera Configuration")113        available_cameras = get_available_cameras()114        115        if available_cameras:116            camera_index = st.selectbox(117                "Select Camera",118                available_cameras,119                index=0,120                help="Choose the camera to use for detection"121            )122        else:123            st.error("❌ No cameras detected!")124            camera_index = 0125        126        # Performance settings127        st.subheader("⚡ Performance")128        detection_interval = st.slider(129            "Detection Interval",130            min_value=1,131            max_value=10,132            value=app_config.performance.detection_interval,133            help="Process every N frames (higher = faster but less frequent detection)"134        )135        136        # Display settings137        st.subheader("🎨 Display Options")138        show_fps = st.checkbox("Show FPS", value=app_config.ui.show_fps)139        show_stats = st.checkbox("Show Statistics", value=app_config.ui.show_stats)140        show_confidence = st.checkbox("Show Confidence Scores", value=app_config.ui.show_confidence)141        142        # Control buttons143        st.subheader("🎮 Controls")144        145        col1, col2 = st.columns(2)146        147        with col1:148            if st.button("🚀 Start Detection", disabled=st.session_state.is_running):149                # Initialize detector if needed150                if (st.session_state.detector is None or 151                    st.session_state.detector.model_name != selected_model or152                    st.session_state.detector.confidence_threshold != confidence_threshold):153                    154                    if initialize_detector(selected_model, confidence_threshold):155                        st.session_state.detector.update_confidence_threshold(confidence_threshold)156                    else:157                        st.stop()158                159                # Initialize webcam if needed160                if (st.session_state.webcam is None or 161                    not st.session_state.webcam.is_camera_available()):162                    163                    if not initialize_webcam(camera_index):164                        st.stop()165                166                st.session_state.is_running = True167                st.session_state.frame_count = 0168                st.rerun()169        170        with col2:171            if st.button("⏹️ Stop Detection", disabled=not st.session_state.is_running):172                st.session_state.is_running = False173                if st.session_state.webcam:174                    st.session_state.webcam.stop_capture()175                    st.session_state.webcam = None176                st.rerun()177    178    # Main content area179    if st.session_state.is_running and st.session_state.webcam and st.session_state.detector:180        # Create columns for video and stats181        col1, col2 = st.columns([3, 1])182        183        with col1:184            st.subheader("📹 Live Video Feed")185            video_placeholder = st.empty()186        187        with col2:188            if show_stats:189                st.subheader("📊 Statistics")190                stats_placeholder = st.empty()191        192        # Performance metrics placeholders193        if show_fps:194            fps_placeholder = st.empty()195        196        # Detection loop197        while st.session_state.is_running:198            frame = st.session_state.webcam.get_frame()199            200            if frame is not None:201                st.session_state.frame_count += 1202                203                # Process frame for detection204                if st.session_state.frame_count % detection_interval == 0:205                    detections = st.session_state.detector.detect_objects(frame)206                    annotated_frame = st.session_state.detector.draw_detections(frame, detections)207                    208                    # Update detection stats209                    st.session_state.detection_stats = st.session_state.detector.get_detection_stats(detections)210                else:211                    annotated_frame = frame212                    detections = []213                214                # Convert BGR to RGB for display215                display_frame = cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB)216                217                # Display video218                with video_placeholder.container():219                    st.image(display_frame, channels="RGB", use_column_width=True)220                221                # Display FPS222                if show_fps:223                    current_fps = st.session_state.webcam.get_fps()224                    fps_placeholder.metric("🎯 FPS", f"{current_fps:.1f}")225                226                # Display statistics227                if show_stats and st.session_state.detection_stats:228                    with stats_placeholder.container():229                        stats = st.session_state.detection_stats230                        231                        st.metric("🎯 Objects Detected", stats['total_objects'])232                        233                        if stats['total_objects'] > 0:234                            st.metric("📈 Avg Confidence", f"{stats['avg_confidence']:.2f}")235                            st.metric("🔝 Max Confidence", f"{stats['max_confidence']:.2f}")236                            237                            # Class breakdown238                            st.write("**Object Classes:**")239                            for class_name, count in stats['class_counts'].items():240                                st.write(f"• {class_name}: {count}")241                242                # Small delay to prevent overwhelming the interface243                time.sleep(0.01)244            245            else:246                st.warning("⚠️ No frame received from camera")247                time.sleep(0.1)248    249    else:250        # Show instructions when not running251        st.info("👆 Configure your settings in the sidebar and click 'Start Detection' to begin!")252        253        # Show system information254        st.subheader("ℹ️ System Information")255        256        col1, col2, col3 = st.columns(3)257        258        with col1:259            st.write("**Available Cameras:**")260            cameras = get_available_cameras()261            if cameras:262                for cam in cameras:263                    st.write(f"• Camera {cam}")264            else:265                st.write("• No cameras detected")266        267        with col2:268            st.write("**Available Models:**")269            for model in app_config.available_models:270                st.write(f"• {model}")271        272        with col3:273            st.write("**Current Configuration:**")274            st.write(f"• Model: {app_config.detection.model_name}")275            st.write(f"• Confidence: {app_config.detection.confidence_threshold}")276            st.write(f"• Camera: {app_config.camera.camera_index}")277 278if __name__ == "__main__":279    main()280