Update app.py

app.py

@@ -2,153 +2,158 @@
 import gradio as gr
 import numpy as np
 import torch
+import soundfile as sf
 import os
 import yaml
 from dotenv import load_dotenv
 from threading import Thread
-from gradio_webrtc import WebRTC
-from twilio.rest import Client
 
 # --- TTS & AI Imports ---
 from parler_tts import ParlerTTSForConditionalGeneration
-from transformers import AutoTokenizer, AutoFeatureExtractor
-from streamer import ParlerTTSStreamer 
+from transformers import AutoTokenizer, AutoFeatureExtractor, set_seed
+from streamer import ParlerTTSStreamer # Make sure streamer.py is available
 
-# --- Local Project Imports ---
 from src.detection.factory import get_detector
 from src.alerting.alert_system import get_alerter
 
 # --- Load Configuration and Environment Variables ---
+# This part is the same as our Streamlit app
 load_dotenv()
 config_path = 'config.yaml'
 with open(config_path, 'r') as f:
     config = yaml.safe_load(f)
-secrets = {"gemini_api_key": os.getenv("GEMINI_API_KEY")}
+secrets = {
+    "gemini_api_key": os.getenv("GEMINI_API_KEY"),
+}
 
 # --- Initialize Backend Components ---
 print("Initializing detector and alerter...")
 detector = get_detector(config)
 alerter = get_alerter(config, secrets["gemini_api_key"])
-print("Initialization complete.")
+print("Initialization complete. Launching UI...")
 
-# --- Twilio TURN Server Setup ---
-account_sid = os.environ.get("TURN_USERNAME")
-auth_token = os.environ.get("TURN_CREDENTIAL")
-rtc_configuration = None
-if account_sid and auth_token:
-    try:
-        client = Client(account_sid, auth_token)
-        token = client.tokens.create()
-        rtc_configuration = {"iceServers": token.ice_servers}
-        print("Twilio TURN server configured successfully.")
-    except Exception as e:
-        print(f"Warning: Failed to create Twilio token. Using public STUN server. Error: {e}")
-# Fallback to a public STUN server if Twilio fails or is not configured
-if rtc_configuration is None:
-    print("Using public STun server.")
-    rtc_configuration = {"iceServers": [{"urls": ["stun:stun.l.google.com:19302"]}]}
-
-
-# --- Parler-TTS Model Setup ---
-print("Loading Parler-TTS model...")
+# --- Parler-TTS Model Setup (Requires GPU) ---
+print("Loading Parler-TTS model. This may take a moment...")
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
 if device == "cpu":
-    print("\nWARNING: Running Parler-TTS on a CPU is slow. A GPU is highly recommended.\n")
+    print("\nWARNING: Running Parler-TTS on a CPU will be extremely slow. A GPU is highly recommended.\n")
 torch_dtype = torch.float16 if device != "cpu" else torch.float32
 
+
+# Using a smaller, faster model suitable for real-time alerts
 repo_id = "parler-tts/parler_tts_mini_v0.1"
 model = ParlerTTSForConditionalGeneration.from_pretrained(repo_id, torch_dtype=torch_dtype).to(device)
 tokenizer = AutoTokenizer.from_pretrained(repo_id)
 feature_extractor = AutoFeatureExtractor.from_pretrained(repo_id)
 print("Parler-TTS model loaded.")
 
-# --- Audio Streaming Generator ---
+# --- Audio Streaming Generator Function ---
 def stream_alert_audio(text_prompt):
-    """A generator that streams audio chunks for a given text prompt."""
+    """
+    A generator function that yields audio chunks for a given text prompt.
+    This is the core of the streaming implementation.
+    """
     sampling_rate = model.config.sampling_rate
-    description = "A female speaker with a clear and urgent voice."
+    description = "Jenny is A female speaker with a clear and urgent voice." # Voice prompt for TTS
+    
     prompt_ids = tokenizer(text_prompt, return_tensors="pt").input_ids.to(device)
     description_ids = tokenizer(description, return_tensors="pt").input_ids.to(device)
+
+    # Setup the streamer
     streamer = ParlerTTSStreamer(model, device, play_steps=int(sampling_rate * 2.0))
-    generation_kwargs = dict(input_ids=description_ids, prompt_input_ids=prompt_ids, streamer=streamer, do_sample=True, temperature=1.0, repetition_penalty=1.2)
+    
+    generation_kwargs = dict(
+        input_ids=description_ids,
+        prompt_input_ids=prompt_ids,
+        streamer=streamer,
+        do_sample=True,
+        temperature=1.0, # Increase for more vocal variety
+        repetition_penalty=1.2,
+    )
+    
+    # Run generation in a separate thread to not block the UI
     thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    
     try:
         thread.start()
         print(f"Audio stream started for: '{text_prompt}'")
+        # Yield audio chunks as they become available
         for new_audio_chunk in streamer:
             yield (sampling_rate, new_audio_chunk)
     finally:
+        # CRITICAL: This block runs after the generator is exhausted (audio finishes)
+        # We reset the alerter state so that a new alert can be triggered later.
         print("Audio stream finished. Resetting alerter state.")
         alerter.reset_alert()
-
-# --- Decoupled Processing Functions ---
-
-def process_video_and_update_state(frame_dict: dict, state: dict):
+        
+# --- Main Webcam Processing Function ---
+def process_live_frame(frame):
     """
-    HIGH-FREQUENCY LOOP: Processes video, updates shared state, and returns the processed frame.
-    This function's speed directly impacts video latency.
+    Processes each webcam frame, performs drowsiness detection, and
+    returns a generator for audio streaming when an alert is triggered.
     """
-    if not frame_dict or "video" not in frame_dict or frame_dict["video"] is None:
-        return np.zeros((480, 640, 3), dtype=np.uint8), state
+    if frame is None:
+        return np.zeros((480, 640, 3), dtype=np.uint8), "Status: Inactive", None
 
-    frame = frame_dict["video"]
     processed_frame, indicators, _ = detector.process_frame(frame)
-    state['indicators'] = indicators
-    return processed_frame, state
-
-def update_ui_from_state(state: dict):
-    """
-    LOW-FREQUENCY LOOP: Reads from state to update status text and trigger audio.
-    This runs independently of the video loop.
-    """
-    indicators = state.get('indicators', {})
     drowsiness_level = indicators.get("drowsiness_level", "Awake")
     lighting = indicators.get("lighting", "Good")
     score = indicators.get("details", {}).get("Score", 0)
 
-    status_text = f"Lighting: {lighting}\nStatus: {drowsiness_level}\nScore: {score:.2f}"
+    # Build status text
+    status_text = f"Lighting: {lighting}\n"
     if lighting == "Low":
-        status_text = "Lighting: Low\nDetection paused due to low light."
+        status_text += "Detection paused due to low light."
+    else:
+        status_text += f"Status: {drowsiness_level}\nScore: {score:.2f}"
 
+    # --- Alert Trigger Logic ---
     audio_output = None
     if drowsiness_level != "Awake":
+        # alerter.trigger_alert() returns the alert TEXT if not on cooldown, otherwise None.
         alert_text = alerter.trigger_alert(level=drowsiness_level)
         if alert_text:
+            # If we got text, it means we can start an alert.
+            # We return the generator function itself. Gradio will handle it.
             audio_output = stream_alert_audio(alert_text)
-    return status_text, audio_output
+
+        else;
+            alert_text = "WAKE UP"
+            audio_output = stream_alert_audio(alert_text)
+            
+    # On subsequent frames where the user is drowsy, trigger_alert() will return None
+    # due to the cooldown, preventing a new stream from starting, which is what we want.
+
+    return processed_frame, status_text, audio_output
+
 
 # --- Gradio UI Definition ---
 with gr.Blocks(theme=gr.themes.Default(primary_hue="blue")) as app:
-    gr.Markdown("# 🚗 Drive Paddy - Drowsiness Detection (WebRTC)")
-    gr.Markdown("Low-latency video processing via WebRTC, with decoupled UI updates for smooth performance.")
-    
-    # Shared state object to pass data between the two processing loops
-    shared_state = gr.State(value={'indicators': {}})
+    gr.Markdown("# 🚗 Drive Paddy - Drowsiness Detection (Streaming)")
+    gr.Markdown("Live drowsiness detection with real-time, streaming voice alerts.")
 
     with gr.Row():
         with gr.Column(scale=2):
-            # This WebRTC component is now correctly used for both input and output of the video stream.
-            webcam = WebRTC(label="Live Camera Feed", rtc_configuration=rtc_configuration)
+            webcam_input = gr.Image(sources=["webcam"], streaming=True, label="Live Camera Feed")
         with gr.Column(scale=1):
+            processed_output = gr.Image(label="Processed Feed")
             status_output = gr.Textbox(label="Live Status", lines=3, interactive=False)
-            audio_alert_output = gr.Audio(label="Alert System", autoplay=True, visible=False, streaming=True)
             
-    # LOOP 1: High-Frequency Video Stream (as fast as possible)
-    # This takes video from the webcam, processes it, and sends it right back.
-    webcam.stream(
-        fn=process_video_and_update_state,
-        inputs=[webcam, shared_state],
-        outputs=[webcam, shared_state],
+            # --- KEY CHANGE: The Audio component now uses streaming=True ---
+            audio_alert_output = gr.Audio(
+                label="Alert System", 
+                autoplay=True, 
+                visible=False, # Hide the player controls
+                streaming=True
+            )
+
+    webcam_input.stream(
+        fn=process_live_frame,
+        inputs=[webcam_input],
+        outputs=[processed_output, status_output, audio_alert_output]
     )
 
-    # LOOP 2: Low-Frequency UI Updates (4 times per second)
-    # This runs on a timer, reads the shared state, and updates the other UI elements.
-    app.load(
-        fn=update_ui_from_state,
-        inputs=[shared_state],
-        outputs=[status_output, audio_alert_output],
-    )
 
+# --- Launch the App ---
 if __name__ == "__main__":
-    print("Starting Drive Paddy WebRTC Application...")
-    app.launch(debug=True)
+    app.launch(debug=True)