new

app.py

@@ -9,6 +9,7 @@ from utils.ear_utils import BlinkDetector
 from gradio_webrtc import WebRTC
 from ultralytics import YOLO
 import torch
+import spaces  # Add spaces import
 
 def smooth_values(history, current_value, window_size=5):
     if current_value is not None:
@@ -38,12 +39,12 @@ GAZE_MODEL_PATH = os.path.join("models", "gaze_estimation_model.pth")
 DISTRACTION_MODEL_PATH = "best.pt"
 
 # --- Global Initializations ---
-gaze_predictor = GazePredictor(GAZE_MODEL_PATH)
+# Load models on CPU initially
+gaze_predictor = GazePredictor(GAZE_MODEL_PATH, device='cpu')  # Assuming GazePredictor accepts device arg
 blink_detector = BlinkDetector()
 
-# Load Distraction Model
+# Load Distraction Model on CPU initially
 distraction_model = YOLO(DISTRACTION_MODEL_PATH)
-distraction_model.to('cpu')
 
 # Distraction Class Names
 distraction_class_names = [
@@ -64,6 +65,9 @@ is_unconscious = False
 frame_count_webcam = 0
 stop_gaze_processing = False
 
+# --- Global State Variables for Distraction Webcam ---
+stop_distraction_processing = False
+
 # Constants
 GAZE_STABILITY_THRESHOLD = 0.5
 TIME_THRESHOLD = 15
@@ -285,10 +289,22 @@ def terminate_gaze_stream():
     frame_count_webcam = 0
     return "Gaze Processing Terminated. State Reset."
 
+def terminate_distraction_stream():
+    global stop_distraction_processing
+    print("Distraction Live Termination signal received. Stopping processing.")
+    stop_distraction_processing = True
+    return "Distraction Live Processing Terminated."
+
+@spaces.GPU  # Add ZeroGPU decorator
 def process_gaze_frame(frame):
     global gaze_history, head_history, ear_history, stable_gaze_time, stable_head_time
     global eye_closed_time, blink_count, start_time, is_unconscious, frame_count_webcam, stop_gaze_processing
 
+    try:
+        gaze_predictor.model.to('cuda')
+    except Exception as e:
+        print(f"Warning: Could not move gaze model to CUDA: {e}")
+
     if stop_gaze_processing:
         return np.zeros((480, 640, 3), dtype=np.uint8)
 
@@ -397,6 +413,66 @@ def process_gaze_frame(frame):
         cv2.putText(error_frame, f"Error: {e}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 0, 0), 2)
         return error_frame
 
+@spaces.GPU  # Add ZeroGPU decorator
+def process_distraction_frame(frame):
+    global stop_distraction_processing
+
+    distraction_model.to('cuda')
+
+    if stop_distraction_processing:
+        return np.zeros((480, 640, 3), dtype=np.uint8)
+
+    if frame is None:
+        return np.zeros((480, 640, 3), dtype=np.uint8)
+
+    try:
+        frame_to_process = frame
+
+        results = distraction_model(frame_to_process, conf=DISTRACTION_CONF_THRESHOLD, verbose=False)
+
+        display_text = "safe driving"
+        alarm_action = None
+
+        for result in results:
+            if result.boxes is not None and len(result.boxes) > 0:
+                boxes = result.boxes.xyxy.cpu().numpy()
+                scores = result.boxes.conf.cpu().numpy()
+                classes = result.boxes.cls.cpu().numpy()
+
+                if len(boxes) > 0:
+                    max_score_idx = scores.argmax()
+                    detected_action_idx = int(classes[max_score_idx])
+                    if 0 <= detected_action_idx < len(distraction_class_names):
+                        detected_action = distraction_class_names[detected_action_idx]
+                        confidence = scores[max_score_idx]
+                        display_text = f"{detected_action}: {confidence:.2f}"
+                        if detected_action != 'safe driving':
+                            alarm_action = detected_action
+                    else:
+                        print(f"Warning: Detected class index {detected_action_idx} out of bounds.")
+                        display_text = "Unknown Detection"
+
+        frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
+        if alarm_action:
+            print(f"ALARM: Unsafe behavior detected - {alarm_action}!")
+            cv2.putText(frame_bgr, f"ALARM: {alarm_action}", (10, 70), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
+
+        text_color = (0, 255, 0) if alarm_action is None else (0, 255, 255)
+        cv2.putText(frame_bgr, display_text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, text_color, 2)
+
+        frame_rgb_processed = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
+        return frame_rgb_processed
+
+    except Exception as e:
+        print(f"Error processing distraction frame: {e}")
+        error_frame = np.zeros((480, 640, 3), dtype=np.uint8)
+        if not error_frame.flags.writeable:
+            error_frame = error_frame.copy()
+        error_frame_bgr = cv2.cvtColor(error_frame, cv2.COLOR_RGB2BGR)
+        cv2.putText(error_frame_bgr, f"Error: {e}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
+        error_frame_rgb = cv2.cvtColor(error_frame_bgr, cv2.COLOR_BGR2RGB)
+        return error_frame_rgb
+
 def create_gaze_interface():
     with gr.Blocks() as gaze_demo:
         gr.Markdown("## Real-time Gaze & Drowsiness Tracking")
@@ -425,6 +501,24 @@ def create_distraction_interface():
     )
     return distraction_demo
 
+def create_distraction_live_interface():
+    with gr.Blocks() as distraction_live_demo:
+        gr.Markdown("## Real-time Distraction Detection (Live)")
+        with gr.Row():
+            webcam_stream = WebRTC(label="Webcam Stream")
+        with gr.Row():
+            terminate_btn = gr.Button("Terminate Process")
+
+        webcam_stream.stream(
+            fn=process_distraction_frame,
+            inputs=[webcam_stream],
+            outputs=[webcam_stream]
+        )
+
+        terminate_btn.click(fn=terminate_distraction_stream, inputs=None, outputs=None)
+
+    return distraction_live_demo
+
 def create_video_interface():
     video_demo = gr.Interface(
         fn=analyze_video,
@@ -436,8 +530,8 @@ def create_video_interface():
     return video_demo
 
 demo = gr.TabbedInterface(
-    [create_video_interface(), create_gaze_interface(), create_distraction_interface()],
-    ["Gaze Video Upload", "Gaze & Drowsiness (Live)", "Distraction Video Upload"],
+    [create_video_interface(), create_gaze_interface(), create_distraction_interface(), create_distraction_live_interface()],
+    ["Gaze Video Upload", "Gaze & Drowsiness (Live)", "Distraction Video Upload", "Distraction Detection (Live)"],
     title="Driver Monitoring System"
 )
 
@@ -453,4 +547,5 @@ if __name__ == "__main__":
     is_unconscious = False
     frame_count_webcam = 0
     stop_gaze_processing = False
+    stop_distraction_processing = False
     demo.launch()

requirements.txt

@@ -11,4 +11,6 @@ tensorflow
 pygame
 twilio
 ultralytics==8.3.93
-torch==2.6.0
+# torch==2.6.0 # Replace with ZeroGPU compatible version, e.g., 2.4.0
+torch==2.4.0 # Example compatible version
+spaces # Add spaces for ZeroGPU