new

app.py

@@ -39,14 +39,7 @@ GAZE_MODEL_PATH = os.path.join("models", "gaze_estimation_model.pth")
 DISTRACTION_MODEL_PATH = "best.pt"
 
 # --- Global Initializations ---
-# Load models on CPU initially
-gaze_predictor = GazePredictor(GAZE_MODEL_PATH)
-blink_detector = BlinkDetector()
-
-# Load Distraction Model on CPU initially
-distraction_model = YOLO(DISTRACTION_MODEL_PATH)
-distraction_model.to('cpu')
-
+blink_detector = BlinkDetector()  # Keep BlinkDetector global as it is CPU-only
 
 # Distraction Class Names
 distraction_class_names = [
@@ -78,10 +71,12 @@ EYE_CLOSURE_THRESHOLD = 10
 HEAD_STABILITY_THRESHOLD = 0.05
 DISTRACTION_CONF_THRESHOLD = 0.1
 
+@spaces.GPU
 def analyze_video(input_video):
+    local_gaze_predictor = GazePredictor(GAZE_MODEL_PATH, device='cuda')  # Load directly to CUDA
+    local_blink_detector = blink_detector  # Use global CPU instance
+
     cap = cv2.VideoCapture(input_video)
-    local_gaze_predictor = GazePredictor(GAZE_MODEL_PATH)
-    local_blink_detector = BlinkDetector()
     fourcc = cv2.VideoWriter_fourcc(*'mp4v')
     temp_fd, temp_path = tempfile.mkstemp(suffix='.mp4')
     os.close(temp_fd)
@@ -205,7 +200,11 @@ def analyze_video(input_video):
         out.release()
     return temp_path
 
+@spaces.GPU
 def analyze_distraction_video(input_video):
+    local_distraction_model = YOLO(DISTRACTION_MODEL_PATH)
+    local_distraction_model.to('cuda')  # Move to GPU
+
     cap = cv2.VideoCapture(input_video)
     if not cap.isOpened():
         print("Error: Could not open video file.")
@@ -224,7 +223,7 @@ def analyze_distraction_video(input_video):
             break
 
         try:
-            results = distraction_model(frame, conf=DISTRACTION_CONF_THRESHOLD, verbose=False)
+            results = local_distraction_model(frame, conf=DISTRACTION_CONF_THRESHOLD, verbose=False)
 
             display_text = "safe driving"
             alarm_action = None
@@ -297,16 +296,14 @@ def terminate_distraction_stream():
     stop_distraction_processing = True
     return "Distraction Live Processing Terminated."
 
-@spaces.GPU  # Add ZeroGPU decorator
+@spaces.GPU
 def process_gaze_frame(frame):
+    gaze_predictor_live = GazePredictor(GAZE_MODEL_PATH, device='cuda')  # Load directly to CUDA
+    local_blink_detector = blink_detector  # Use global CPU instance
+
     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)
 
@@ -319,11 +316,8 @@ def process_gaze_frame(frame):
         start_time = current_time
 
     try:
-        head_pose_gaze, gaze_h, gaze_v = gaze_predictor.predict_gaze(frame)
-        current_gaze = np.array([gaze_h, gaze_v]) if gaze_h is not None and gaze_v is not None else None
-        smoothed_gaze = smooth_values(gaze_history, current_gaze)
-
-        ear, left_eye, right_eye, head_pose, left_iris, right_iris = blink_detector.detect_blinks(frame)
+        head_pose_gaze, gaze_h, gaze_v = gaze_predictor_live.predict_gaze(frame)
+        ear, left_eye, right_eye, head_pose, left_iris, right_iris = local_blink_detector.detect_blinks(frame)
 
         if ear is None:
             cv2.putText(frame, "No face detected", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
@@ -415,11 +409,12 @@ 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
+@spaces.GPU
 def process_distraction_frame(frame):
-    global stop_distraction_processing
+    distraction_model_live = YOLO(DISTRACTION_MODEL_PATH)
+    distraction_model_live.to('cuda')
 
-    distraction_model.to('cuda')
+    global stop_distraction_processing
 
     if stop_distraction_processing:
         return np.zeros((480, 640, 3), dtype=np.uint8)
@@ -430,7 +425,7 @@ def process_distraction_frame(frame):
     try:
         frame_to_process = frame
 
-        results = distraction_model(frame_to_process, conf=DISTRACTION_CONF_THRESHOLD, verbose=False)
+        results = distraction_model_live(frame_to_process, conf=DISTRACTION_CONF_THRESHOLD, verbose=False)
 
         display_text = "safe driving"
         alarm_action = None