new

app.py

@@ -69,13 +69,8 @@ GAZE_MODEL_PATH = os.path.join("models", "gaze_estimation_model.pth")
 DISTRACTION_MODEL_PATH = "best.pt"
 
 # --- Global Initializations ---
-gaze_predictor = GazePredictor(GAZE_MODEL_PATH)
 blink_detector = BlinkDetector()
 
-# Load Distraction Model
-distraction_model = YOLO(DISTRACTION_MODEL_PATH)
-distraction_model.to('cpu')
-
 # Distraction Class Names
 distraction_class_names = [
     'safe driving', 'drinking', 'eating', 'hair and makeup',
@@ -106,6 +101,7 @@ EYE_CLOSURE_THRESHOLD = 10
 HEAD_STABILITY_THRESHOLD = 0.05
 DISTRACTION_CONF_THRESHOLD = 0.1
 
+@spaces.GPU(duration=30)  # Set duration to 30 seconds for real-time processing
 def analyze_video(input_video):
     cap = cv2.VideoCapture(input_video)
     local_gaze_predictor = GazePredictor(GAZE_MODEL_PATH)
@@ -247,13 +243,16 @@ def analyze_distraction_video(input_video):
 
     fps = cap.get(cv2.CAP_PROP_FPS) or 30
 
+    local_distraction_model = YOLO(DISTRACTION_MODEL_PATH)
+    local_distraction_model.to('cpu')
+
     while True:
         ret, frame = cap.read()
         if not ret:
             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
@@ -312,9 +311,12 @@ def process_distraction_frame(frame):
     if frame is None:
         return np.zeros((480, 640, 3), dtype=np.uint8)
     
+    local_distraction_model = YOLO(DISTRACTION_MODEL_PATH)
+    local_distraction_model.to('cpu')
+
     try:
         # Run distraction detection model
-        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
@@ -410,8 +412,10 @@ def process_gaze_frame(frame):
     if start_time == 0:
         start_time = current_time
 
+    local_gaze_predictor = GazePredictor(GAZE_MODEL_PATH)
+
     try:
-        head_pose_gaze, gaze_h, gaze_v = gaze_predictor.predict_gaze(frame)
+        head_pose_gaze, gaze_h, gaze_v = local_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)
 

scripts/inference.py

@@ -17,13 +17,29 @@ class GazeEstimationModel(torch.nn.Module):
 
 class GazePredictor:
     def __init__(self, model_path):
+        # Initialize without moving to device - we'll do that during prediction
+        self.model_path = model_path
+        self.model = None
+        self.device = None
+        
+        # Define transform outside of initialization
+        self.transform = transforms.Compose([
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+        ])
+        
+    def _initialize_model(self):
+        # Only initialize model when needed (inside ZeroGPU function)
+        if self.model is not None:
+            return
+            
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         
         # Initialize the custom model
         self.model = GazeEstimationModel()
         
         # Load the state dictionary
-        state_dict = torch.load(model_path, map_location=self.device)
+        state_dict = torch.load(self.model_path, map_location=self.device)
         
         # Check if state_dict has 'backbone.' prefix and strip it if necessary
         new_state_dict = {}
@@ -42,14 +58,11 @@ class GazePredictor:
         # Move to device and set to evaluation mode
         self.model.to(self.device)
         self.model.eval()
-        
-        # Define preprocessing transform
-        self.transform = transforms.Compose([
-            transforms.ToTensor(),
-            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
-        ])
     
     def predict_gaze(self, frame):
+        # Initialize model if not already done
+        self._initialize_model()
+        
         preprocessed = preprocess_frame(frame)
         preprocessed = preprocessed[0]
         preprocessed = self.transform(preprocessed).float().unsqueeze(0)