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()
|