Update app.py

app.py

@@ -2,157 +2,136 @@
 import gradio as gr
 import numpy as np
 import torch
-import soundfile as sf
-import os
-import yaml
+import os, yaml, soundfile as sf
 from dotenv import load_dotenv
 from threading import Thread
 
 # --- 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 file
 
 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
+# ──────────────────────────────────────────────────────────
+# CONFIG & BACKEND SET-UP
+# ──────────────────────────────────────────────────────────
 load_dotenv()
-config_path = 'config.yaml'
-with open(config_path, 'r') as f:
+
+with open("config.yaml", "r") as f:
     config = yaml.safe_load(f)
-secrets = {
-    "gemini_api_key": os.getenv("GEMINI_API_KEY"),
-}
 
-# --- Initialize Backend Components ---
-print("Initializing detector and alerter...")
+secrets = {"gemini_api_key": os.getenv("GEMINI_API_KEY")}
+
+print("Initializing detector and alerter …")
 detector = get_detector(config)
 alerter = get_alerter(config, secrets["gemini_api_key"])
-print("Initialization complete. Launching UI...")
+print("Backend ready.")
 
-# --- Parler-TTS Model Setup (Requires GPU) ---
-print("Loading Parler-TTS model. This may take a moment...")
+# ──────────────────────────────────────────────────────────
+# TTS MODEL (Parler-TTS mini)
+# ──────────────────────────────────────────────────────────
 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")
-torch_dtype = torch.float16 if device != "cpu" else torch.float32
-
+    print("\n⚠️  Running TTS on CPU will be slow; only ‘Very Drowsy’ alerts will use it.\n")
 
-# Using a smaller, faster model suitable for real-time alerts
+model_dtype = torch.float16 if device != "cpu" else torch.float32
 repo_id = "parler-tts/parler_tts_mini_v0.1"
-model = ParlerTTSForConditionalGeneration.from_pretrained(repo_id, torch_dtype=torch_dtype).to(device)
+
+print("Loading Parler-TTS …")
+model = ParlerTTSForConditionalGeneration.from_pretrained(repo_id,
+                                                         torch_dtype=model_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 ---
-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.
-    """
+print("TTS loaded.")
+
+# ──────────────────────────────────────────────────────────
+# AUDIO STREAMER
+# ──────────────────────────────────────────────────────────
+def stream_alert_audio(text_prompt: str):
+    """Yields (sampling_rate, np.ndarray) chunks for Gradio streaming."""
     sampling_rate = model.config.sampling_rate
-    description = "Jenny is A female speaker with a clear and urgent voice." # Voice prompt for TTS
-    
+    voice_desc = "Jenny is 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)
+    desc_ids   = tokenizer(voice_desc, 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,
+
+    gen_kwargs = dict(
+        input_ids=desc_ids,
         prompt_input_ids=prompt_ids,
         streamer=streamer,
         do_sample=True,
-        temperature=1.0, # Increase for more vocal variety
+        temperature=1.0,
         repetition_penalty=1.2,
     )
-    
-    # Run generation in a separate thread to not block the UI
-    thread = Thread(target=model.generate, kwargs=generation_kwargs)
-    
+
+    thread = Thread(target=model.generate, kwargs=gen_kwargs, daemon=True)
+
     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 new_audio_chunk
+        for chunk in streamer:
+            yield (sampling_rate, 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.")
+        thread.join(timeout=0.1)
         alerter.reset_alert()
-        
-# --- Main Webcam Processing Function ---
+
+# ──────────────────────────────────────────────────────────
+# FRAME PROCESSOR
+# ──────────────────────────────────────────────────────────
 def process_live_frame(frame):
-    """
-    Processes each webcam frame, performs drowsiness detection, and
-    returns a generator for audio streaming when an alert is triggered.
-    """
     if frame is None:
-        return np.zeros((480, 640, 3), dtype=np.uint8), "Status: Inactive", None
-
-    processed_frame, indicators, _ = detector.process_frame(frame)
-    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"
-    if lighting == "Low":
-        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)
-
-        else:
-            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 ---
+        return np.zeros((480, 640, 3), np.uint8), "Status: Inactive", None
+
+    processed, indicators, _ = detector.process_frame(frame)
+    level      = indicators.get("drowsiness_level", "Awake")
+    lighting   = indicators.get("lighting", "Good")
+    score      = indicators.get("details", {}).get("Score", 0)
+
+    status_txt = f"Lighting: {lighting}\n"
+    status_txt += ("Detection paused due to low light."
+                   if lighting == "Low"
+                   else f"Status: {level}\nScore: {score:.2f}")
+
+    audio_out = None
+    if level != "Awake" and lighting != "Low":
+        payload = alerter.trigger_alert(level=level)
+        if payload:
+            # Static file path → bytes, Dynamic Gemini path → str
+            if isinstance(payload, bytes):
+                # Return raw bytes (Gradio accepts bytes for .wav / .mp3)
+                audio_out = payload
+            elif isinstance(payload, str):
+                audio_out = stream_alert_audio(payload)
+
+    return processed, status_txt, audio_out
+
+# ──────────────────────────────────────────────────────────
+# GRADIO UI
+# ──────────────────────────────────────────────────────────
 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")
+    gr.Markdown("Live detection with real-time voice alerts.")
 
     with gr.Row():
         with gr.Column(scale=2):
-            webcam_input = gr.Image(sources=["webcam"], streaming=True, label="Live Camera Feed")
+            webcam = 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)
-            
-            # --- 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(
+            processed_img = gr.Image(label="Processed Feed")
+            status_box    = gr.Textbox(label="Live Status", lines=3, interactive=False)
+            alert_audio   = gr.Audio(label="Alert",
+                                     autoplay=True,
+                                     streaming=True,
+                                     height=40)
+
+    webcam.stream(
         fn=process_live_frame,
-        inputs=[webcam_input],
-        outputs=[processed_output, status_output, audio_alert_output]
+        inputs=webcam,
+        outputs=[processed_img, status_box, alert_audio],
     )
 
-
-# --- Launch the App ---
 if __name__ == "__main__":
-    app.launch(debug=True)
+    app.launch(debug=True)