Update app.py

app.py

@@ -1,82 +1,137 @@
-import cv2
+import os
 import streamlit as st
 from ultralytics import YOLO
+import cv2
+import random
+import time
 from gtts import gTTS
+import pygame
 import tempfile
-import os
-
-# Load YOLOv8 model
-@st.cache_resource
-def load_model():
-    return YOLO("yolov8n.pt")  # Pre-trained YOLOv8 model
-
-model = load_model()
+import threading
+from datetime import datetime, timedelta
 
-# Streamlit App
-st.title("Real-Time Object Detection for Blind Assistance")
-st.write(
-    "This app detects objects in real-time via an IP webcam and provides audio feedback."
-)
-
-# Input for IP Webcam URL
-ip_webcam_url = st.text_input(
-    "Enter your IP Webcam URL (e.g., http://<ip-address>:8080/video):",
-    value="http://<your-ip>:8080/video",
-)
-
-# Start detection
-start_detection = st.button("Start Real-Time Detection")
+# Initialize pygame mixer for audio alerts
+pygame.mixer.quit()
+pygame.mixer.init()
 
+# Load YOLOv8 model
+yolo = YOLO("yolov8n.pt")
+
+# Streamlit app layout
+st.set_page_config(page_title="Assistive Vision App", layout="wide")
+st.title("Real-Time Object Detection with Assistive Vision")
+st.write("This app detects objects in real-time using a webcam feed and provides audio feedback for visually impaired people.")
+
+# Placeholder for video feed
+video_placeholder = st.empty()
+
+# User controls
+start_detection = st.button("Start Detection")
+stop_detection = st.button("Stop Detection")
+enable_audio = st.checkbox("Enable Audio Alerts", value=True)
+
+# Audio alerts settings
+alert_categories = {"person", "cat", "dog", "knife", "fire", "gun"}
+last_alert_time = {}
+alert_cooldown = timedelta(seconds=10)  # Cooldown for audio alerts
+
+# Create a directory for temporary audio files
+audio_temp_dir = "audio_temp_files"
+if not os.path.exists(audio_temp_dir):
+    os.makedirs(audio_temp_dir)
+
+def play_audio_alert(label, position):
+    """Generate and play an audio alert."""
+    phrases = [
+        f"Be careful, there's a {label} on your {position}.",
+        f"Watch out! {label} detected on your {position}.",
+        f"Alert! A {label} is on your {position}.",
+    ]
+    alert_text = random.choice(phrases)
+
+    temp_audio_path = os.path.join(audio_temp_dir, f"alert_{datetime.now().strftime('%Y%m%d_%H%M%S_%f')}.mp3")
+    tts = gTTS(alert_text)
+    tts.save(temp_audio_path)
+
+    try:
+        pygame.mixer.music.load(temp_audio_path)
+        pygame.mixer.music.play()
+
+        def cleanup_audio():
+            while pygame.mixer.music.get_busy():
+                time.sleep(0.1)
+            pygame.mixer.music.stop()
+            if os.path.exists(temp_audio_path):
+                os.remove(temp_audio_path)
+
+        threading.Thread(target=cleanup_audio, daemon=True).start()
+
+    except Exception as e:
+        print(f"Audio playback error: {e}")
+
+def process_frame(frame, enable_audio):
+    """Process a single video frame."""
+    results = yolo(frame)
+    result = results[0]
+
+    detected_objects = {}
+    for box in result.boxes:
+        x1, y1, x2, y2 = map(int, box.xyxy[0])
+        label = result.names[int(box.cls[0])]
+
+        # Only alert for specific categories
+        if enable_audio and label not in alert_categories:
+            continue
+
+        # Determine object position
+        frame_center_x = frame.shape[1] // 2
+        object_center_x = (x1 + x2) // 2
+        position = "left" if object_center_x < frame_center_x else "right"
+
+        detected_objects[label] = position
+
+        # Draw bounding box and label on the frame
+        cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
+        cv2.putText(frame, f"{label}", (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
+
+    return detected_objects, frame
+
+# Real-time detection logic
 if start_detection:
-    # Debug: Print the URL being used
-    st.write(f"Connecting to: {ip_webcam_url}")
-
-    # Check if the URL starts with "http://"
-    if not ip_webcam_url.startswith("http://"):
-        st.error("Error: Invalid URL. Ensure the URL starts with 'http://'.")
-    else:
-        # Open video stream
-        st.write("Connecting to the webcam stream...")
-        cap = cv2.VideoCapture(ip_webcam_url)
-
-        if not cap.isOpened():
-            st.error(
-                "Error: Could not connect to the webcam stream. Ensure the IP Webcam URL is correct and accessible."
-            )
+    st.success("Starting real-time object detection...")
+    try:
+        video_capture = cv2.VideoCapture(0)
+        if not video_capture.isOpened():
+            st.error("Error: Could not access the webcam. Please check your webcam settings.")
         else:
-            stframe = st.empty()  # Placeholder for video frames
-
-            while cap.isOpened():
-                ret, frame = cap.read()
+            while not stop_detection:
+                ret, frame = video_capture.read()
                 if not ret:
-                    st.error("Error: Failed to capture frame from the webcam.")
+                    st.error("Error: Unable to read from webcam. Please check your webcam.")
                     break
 
-                # Perform object detection
-                results = model(frame)
-                detected_objects = [model.names[int(box.cls)] for box in results[0].boxes]
-
-                # Generate audio feedback for detected objects
-                if detected_objects:
-                    objects_text = ", ".join(set(detected_objects))
-                    feedback_text = f"Detected: {objects_text}."
-                    st.write(feedback_text)
+                detected_objects, processed_frame = process_frame(frame, enable_audio)
 
-                    # Convert feedback to speech
-                    tts = gTTS(text=feedback_text, lang="en")
-                    audio_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp3")
-                    tts.save(audio_file.name)
-                    os.system(f"mpg123 {audio_file.name}")  # Play the audio file
+                # Convert frame to RGB for Streamlit display
+                frame_rgb = cv2.cvtColor(processed_frame, cv2.COLOR_BGR2RGB)
+                video_placeholder.image(frame_rgb, channels="RGB", use_column_width=True)
 
-                # Annotate frame
-                annotated_frame = results[0].plot()
+                # Play audio alerts
+                if enable_audio:
+                    current_time = datetime.now()
+                    for label, position in detected_objects.items():
+                        if label not in last_alert_time or current_time - last_alert_time[label] > alert_cooldown:
+                            play_audio_alert(label, position)
+                            last_alert_time[label] = current_time
 
-                # Display the annotated frame
-                stframe.image(annotated_frame, channels="BGR", use_column_width=True)
+                time.sleep(0.1)
 
-                # Break loop if the 'q' key is pressed
-                if cv2.waitKey(1) & 0xFF == ord("q"):
-                    break
+    except Exception as e:
+        st.error(f"An error occurred: {e}")
+    finally:
+        video_capture.release()
+        cv2.destroyAllWindows()
+        pygame.mixer.quit()
 
-            cap.release()
-            st.write("Real-Time Detection Stopped.")
+elif stop_detection:
+    st.warning("Real-time object detection stopped.")