Update app.py

app.py

@@ -2,10 +2,8 @@
 import gradio as gr
 import numpy as np
 import torch
-import soundfile as sf
 import os
 import yaml
-import spaces
 from dotenv import load_dotenv
 from threading import Thread
 from gradio_webrtc import WebRTC
@@ -13,168 +11,145 @@ from twilio.rest import Client
 
 # --- TTS & AI Imports ---
 from parler_tts import ParlerTTSForConditionalGeneration
-from transformers import AutoTokenizer, AutoFeatureExtractor, set_seed
-from streamer import ParlerTTSStreamer # Make sure streamer.py is available
+from transformers import AutoTokenizer, AutoFeatureExtractor
+from streamer import ParlerTTSStreamer 
 
+# --- 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. Launching UI...")
+print("Initialization complete.")
 
+# --- 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:
-    client = Client(account_sid, auth_token)
-
-    token = client.tokens.create()
-
-    rtc_configuration = {
-        "iceServers": token.ice_servers,
-        "iceTransportPolicy": "relay",
-    }
-else:
-    rtc_configuration = None
-
-
-# --- Parler-TTS Model Setup (Requires GPU) ---
-print("Loading Parler-TTS model. This may take a moment...")
+    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...")
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
 if device == "cpu":
-    print("\nWARNING: Running Parler-TTS on a CPU will be extremely slow. A GPU is highly recommended.\n")
+    print("\nWARNING: Running Parler-TTS on a CPU is 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 Function ---
-@spaces.GPU
+# --- Audio Streaming Generator ---
 def stream_alert_audio(text_prompt):
-    """
-    A generator function that yields audio chunks for a given text prompt.
-    This is the core of the streaming implementation.
-    """
+    """A generator that streams audio chunks for a given text prompt."""
     sampling_rate = model.config.sampling_rate
-    description = "Jenny is A female speaker with a clear and urgent voice." # Voice prompt for TTS
-    
+    description = "A female speaker with a clear and urgent voice."
     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, # Increase for more vocal variety
-        repetition_penalty=1.2,
-    )
-    
-    # Run generation in a separate thread to not block the UI
+    generation_kwargs = dict(input_ids=description_ids, prompt_input_ids=prompt_ids, streamer=streamer, do_sample=True, temperature=1.0, repetition_penalty=1.2)
     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()
-        
-# --- Main Webcam Processing Function ---
-@spaces.GPU
-def process_live_frame(frame):
+
+# --- Decoupled Processing Functions ---
+
+def process_video_and_update_state(frame_dict: dict, state: dict):
     """
-    Processes each webcam frame, performs drowsiness detection, and
-    returns a generator for audio streaming when an alert is triggered.
+    HIGH-FREQUENCY LOOP: Processes video, updates shared state, and returns the processed frame.
+    This function's speed directly impacts video latency.
     """
-    if frame is None:
-        return np.zeros((480, 640, 3), dtype=np.uint8), "Status: Inactive", None
+    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
 
+    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)
 
-    # Build status text
-    status_text = f"Lighting: {lighting}\n"
+    status_text = f"Lighting: {lighting}\nStatus: {drowsiness_level}\nScore: {score:.2f}"
     if lighting == "Low":
-        status_text += "Detection paused due to low light."
-    else:
-        status_text += f"Status: {drowsiness_level}\nScore: {score:.2f}"
+        status_text = "Lighting: Low\nDetection paused due to low light."
 
-    # --- 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)
-        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
-
+    return 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 (Streaming)")
-    gr.Markdown("Live drowsiness detection with real-time, streaming voice alerts.")
+    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': {}})
 
     with gr.Row():
         with gr.Column(scale=2):
-            webcam_input = WebRTC(label="Stream", rtc_configuration=rtc_configuration)
+            # 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)
         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)
             
-            # --- KEY CHANGE: The Audio component now uses streaming=True ---
-            audio_alert_output = gr.Audio(
-                label="Alert System", 
-                autoplay=True, 
-                visible=True, # Hide the player controls
-                streaming=True
-            )
-
-    webcam_input.stream(
-        fn=process_live_frame,
-        inputs=[webcam_input],
-        outputs=[status_output, audio_alert_output],
-        time_limit=10
+    # 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],
     )
 
+    # 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],
+        every=0.25  # Run this function every 250 milliseconds
+    )
 
-# --- Launch the App ---
 if __name__ == "__main__":
-    app.launch(debug=True)
+    print("Starting Drive Paddy WebRTC Application...")
+    app.launch(debug=True)