File size: 10,642 Bytes
409c18f
 
 
 
 
58ede7f
409c18f
 
 
 
 
 
58ede7f
409c18f
 
58ede7f
409c18f
 
 
 
58ede7f
409c18f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
58ede7f
409c18f
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
"""
Real-time Object Detection with YOLO - Streamlit Application
Main application file for the web interface.
"""

import streamlit as st
import cv2
import numpy as np
import time
from PIL import Image
import sys
import os

# Add src directory to path
sys.path.append(os.path.join(os.path.dirname(__file__), 'src'))

from detection.yolo_detector import YOLODetector
from detection.webcam_capture import WebcamCapture, get_available_cameras
from utils.config import app_config
from utils.helpers import format_detection_info, log_detection_summary

# Page configuration
st.set_page_config(
    page_title=app_config.ui.page_title,
    page_icon="๐ŸŽฏ",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Initialize session state
if 'detector' not in st.session_state:
    st.session_state.detector = None
if 'webcam' not in st.session_state:
    st.session_state.webcam = None
if 'is_running' not in st.session_state:
    st.session_state.is_running = False
if 'frame_count' not in st.session_state:
    st.session_state.frame_count = 0
if 'detection_stats' not in st.session_state:
    st.session_state.detection_stats = {}

def initialize_detector(model_name: str, confidence_threshold: float):
    """Initialize the YOLO detector."""
    try:
        with st.spinner(f"Loading {model_name} model..."):
            detector = YOLODetector(model_name, confidence_threshold)
            if detector.model is not None:
                st.session_state.detector = detector
                st.success(f"โœ… Model {model_name} loaded successfully!")
                return True
            else:
                st.error(f"โŒ Failed to load model {model_name}")
                return False
    except Exception as e:
        st.error(f"โŒ Error loading model: {str(e)}")
        return False

def initialize_webcam(camera_index: int):
    """Initialize the webcam."""
    try:
        webcam = WebcamCapture(
            camera_index=camera_index,
            width=app_config.camera.frame_width,
            height=app_config.camera.frame_height
        )
        if webcam.start_capture():
            st.session_state.webcam = webcam
            st.success(f"โœ… Camera {camera_index} initialized successfully!")
            return True
        else:
            st.error(f"โŒ Failed to initialize camera {camera_index}")
            return False
    except Exception as e:
        st.error(f"โŒ Error initializing camera: {str(e)}")
        return False

def main():
    """Main application function."""
    
    # Title and description
    st.title("๐ŸŽฏ Real-time Object Detection with YOLO")
    st.markdown("Detect objects in real-time using your webcam and YOLO models.")
    
    # Sidebar for controls
    with st.sidebar:
        st.header("โš™๏ธ Settings")
        
        # Model selection
        st.subheader("๐Ÿค– Model Configuration")
        selected_model = st.selectbox(
            "Select YOLO Model",
            app_config.available_models,
            index=app_config.available_models.index(app_config.detection.model_name),
            help="Choose the YOLO model variant. Nano is fastest, X is most accurate."
        )
        
        # Show model description
        model_info = app_config.get_model_info(selected_model)
        st.info(model_info['description'])
        
        # Confidence threshold
        confidence_threshold = st.slider(
            "Confidence Threshold",
            min_value=0.1,
            max_value=1.0,
            value=app_config.detection.confidence_threshold,
            step=0.05,
            help="Minimum confidence score for detections"
        )
        
        # Camera selection
        st.subheader("๐Ÿ“น Camera Configuration")
        available_cameras = get_available_cameras()
        
        if available_cameras:
            camera_index = st.selectbox(
                "Select Camera",
                available_cameras,
                index=0,
                help="Choose the camera to use for detection"
            )
        else:
            st.error("โŒ No cameras detected!")
            camera_index = 0
        
        # Performance settings
        st.subheader("โšก Performance")
        detection_interval = st.slider(
            "Detection Interval",
            min_value=1,
            max_value=10,
            value=app_config.performance.detection_interval,
            help="Process every N frames (higher = faster but less frequent detection)"
        )
        
        # Display settings
        st.subheader("๐ŸŽจ Display Options")
        show_fps = st.checkbox("Show FPS", value=app_config.ui.show_fps)
        show_stats = st.checkbox("Show Statistics", value=app_config.ui.show_stats)
        show_confidence = st.checkbox("Show Confidence Scores", value=app_config.ui.show_confidence)
        
        # Control buttons
        st.subheader("๐ŸŽฎ Controls")
        
        col1, col2 = st.columns(2)
        
        with col1:
            if st.button("๐Ÿš€ Start Detection", disabled=st.session_state.is_running):
                # Initialize detector if needed
                if (st.session_state.detector is None or 
                    st.session_state.detector.model_name != selected_model or
                    st.session_state.detector.confidence_threshold != confidence_threshold):
                    
                    if initialize_detector(selected_model, confidence_threshold):
                        st.session_state.detector.update_confidence_threshold(confidence_threshold)
                    else:
                        st.stop()
                
                # Initialize webcam if needed
                if (st.session_state.webcam is None or 
                    not st.session_state.webcam.is_camera_available()):
                    
                    if not initialize_webcam(camera_index):
                        st.stop()
                
                st.session_state.is_running = True
                st.session_state.frame_count = 0
                st.rerun()
        
        with col2:
            if st.button("โน๏ธ Stop Detection", disabled=not st.session_state.is_running):
                st.session_state.is_running = False
                if st.session_state.webcam:
                    st.session_state.webcam.stop_capture()
                    st.session_state.webcam = None
                st.rerun()
    
    # Main content area
    if st.session_state.is_running and st.session_state.webcam and st.session_state.detector:
        # Create columns for video and stats
        col1, col2 = st.columns([3, 1])
        
        with col1:
            st.subheader("๐Ÿ“น Live Video Feed")
            video_placeholder = st.empty()
        
        with col2:
            if show_stats:
                st.subheader("๐Ÿ“Š Statistics")
                stats_placeholder = st.empty()
        
        # Performance metrics placeholders
        if show_fps:
            fps_placeholder = st.empty()
        
        # Detection loop
        while st.session_state.is_running:
            frame = st.session_state.webcam.get_frame()
            
            if frame is not None:
                st.session_state.frame_count += 1
                
                # Process frame for detection
                if st.session_state.frame_count % detection_interval == 0:
                    detections = st.session_state.detector.detect_objects(frame)
                    annotated_frame = st.session_state.detector.draw_detections(frame, detections)
                    
                    # Update detection stats
                    st.session_state.detection_stats = st.session_state.detector.get_detection_stats(detections)
                else:
                    annotated_frame = frame
                    detections = []
                
                # Convert BGR to RGB for display
                display_frame = cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB)
                
                # Display video
                with video_placeholder.container():
                    st.image(display_frame, channels="RGB", use_column_width=True)
                
                # Display FPS
                if show_fps:
                    current_fps = st.session_state.webcam.get_fps()
                    fps_placeholder.metric("๐ŸŽฏ FPS", f"{current_fps:.1f}")
                
                # Display statistics
                if show_stats and st.session_state.detection_stats:
                    with stats_placeholder.container():
                        stats = st.session_state.detection_stats
                        
                        st.metric("๐ŸŽฏ Objects Detected", stats['total_objects'])
                        
                        if stats['total_objects'] > 0:
                            st.metric("๐Ÿ“ˆ Avg Confidence", f"{stats['avg_confidence']:.2f}")
                            st.metric("๐Ÿ” Max Confidence", f"{stats['max_confidence']:.2f}")
                            
                            # Class breakdown
                            st.write("**Object Classes:**")
                            for class_name, count in stats['class_counts'].items():
                                st.write(f"โ€ข {class_name}: {count}")
                
                # Small delay to prevent overwhelming the interface
                time.sleep(0.01)
            
            else:
                st.warning("โš ๏ธ No frame received from camera")
                time.sleep(0.1)
    
    else:
        # Show instructions when not running
        st.info("๐Ÿ‘† Configure your settings in the sidebar and click 'Start Detection' to begin!")
        
        # Show system information
        st.subheader("โ„น๏ธ System Information")
        
        col1, col2, col3 = st.columns(3)
        
        with col1:
            st.write("**Available Cameras:**")
            cameras = get_available_cameras()
            if cameras:
                for cam in cameras:
                    st.write(f"โ€ข Camera {cam}")
            else:
                st.write("โ€ข No cameras detected")
        
        with col2:
            st.write("**Available Models:**")
            for model in app_config.available_models:
                st.write(f"โ€ข {model}")
        
        with col3:
            st.write("**Current Configuration:**")
            st.write(f"โ€ข Model: {app_config.detection.model_name}")
            st.write(f"โ€ข Confidence: {app_config.detection.confidence_threshold}")
            st.write(f"โ€ข Camera: {app_config.camera.camera_index}")

if __name__ == "__main__":
    main()