Update src/detection/strategies/cnn_model.py

src/detection/strategies/cnn_model.py

@@ -5,26 +5,24 @@ import torch
 import torchvision.transforms as transforms
 from torchvision.models import efficientnet_b7
 import cv2
-import dlib
 from PIL import Image
 import os
 
 class CnnProcessor(BaseProcessor):
     """
     Drowsiness detection using a pre-trained EfficientNet-B7 model.
+    This version receives face landmarks from another processor instead of using dlib.
     """
     def __init__(self, config):
         self.settings = config['cnn_model_settings']
         self.model_path = self.settings['model_path']
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         
-        # Initialize dlib for face detection
-        self.face_detector = dlib.get_frontal_face_detector()
+        # dlib is no longer needed.
+        # self.face_detector = dlib.get_frontal_face_detector()
         
-        # Load the model
         self.model = self._load_model()
         
-        # Define image transformations
         self.transform = transforms.Compose([
             transforms.Resize((224, 224)),
             transforms.ToTensor(),
@@ -35,66 +33,61 @@ class CnnProcessor(BaseProcessor):
         """Loads the EfficientNet-B7 model and custom weights."""
         if not os.path.exists(self.model_path):
             print(f"Error: Model file not found at {self.model_path}")
-            print("Please run 'python download_model.py' first.")
             return None
             
         try:
-            # Initialize the model structure
             model = efficientnet_b7()
-            # Modify the final classifier layer to match the number of output classes (e.g., 2: drowsy, not_drowsy)
             num_ftrs = model.classifier[1].in_features
-            model.classifier[1] = torch.nn.Linear(num_ftrs, 2) # Assuming 2 output classes
-
-            # Load the saved weights
+            model.classifier[1] = torch.nn.Linear(num_ftrs, 2)
             model.load_state_dict(torch.load(self.model_path, map_location=self.device))
             model.to(self.device)
-            model.eval() # Set the model to evaluation mode
+            model.eval()
             print(f"CNN Model '{self.model_path}' loaded successfully on {self.device}.")
             return model
         except Exception as e:
             print(f"Error loading CNN model: {e}")
             return None
 
-    def process_frame(self, frame):
+    def process_frame(self, frame, face_landmarks=None):
         """
-        Processes a frame to detect drowsiness using the CNN model.
+        Processes a frame using the CNN model with pre-supplied landmarks.
         """
-        if self.model is None:
+        if self.model is None or face_landmarks is None:
             return frame, {"cnn_prediction": False}
 
-        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
-        faces = self.face_detector(gray)
         is_drowsy_prediction = False
+        h, w, _ = frame.shape
+        
+        landmarks = face_landmarks[0].landmark
+        
+        # Calculate bounding box from landmarks
+        x_coords = [lm.x * w for lm in landmarks]
+        y_coords = [lm.y * h for lm in landmarks]
+        x1, y1 = int(min(x_coords)), int(min(y_coords))
+        x2, y2 = int(max(x_coords)), int(max(y_coords))
 
-        for face in faces:
-            x1, y1, x2, y2 = face.left(), face.top(), face.right(), face.bottom()
-            
-            # Crop the face from the frame
-            face_crop = frame[y1:y2, x1:x2]
-            
-            # Ensure the crop is valid before processing
-            if face_crop.size == 0:
-                continue
-                
-            # Convert to PIL Image and apply transformations
+        # Add some padding to the bounding box
+        padding = 10
+        x1 = max(0, x1 - padding)
+        y1 = max(0, y1 - padding)
+        x2 = min(w, x2 + padding)
+        y2 = min(h, y2 + padding)
+
+        # Crop the face
+        face_crop = frame[y1:y2, x1:x2]
+        
+        if face_crop.size > 0:
             pil_image = Image.fromarray(cv2.cvtColor(face_crop, cv2.COLOR_BGR2RGB))
             image_tensor = self.transform(pil_image).unsqueeze(0).to(self.device)
             
-            # Perform inference
             with torch.no_grad():
                 outputs = self.model(image_tensor)
                 _, preds = torch.max(outputs, 1)
-                # Assuming class 1 is 'drowsy' and class 0 is 'not_drowsy'
-                print(preds)
-                if preds.item() == 1:
+                if preds.item() == 1: # Assuming class 1 is 'drowsy'
                     is_drowsy_prediction = True
 
-            # Draw bounding box for visualization
             cv2.rectangle(frame, (x1, y1), (x2, y2), (255, 255, 0), 2)
             label = "Drowsy" if is_drowsy_prediction else "Awake"
             cv2.putText(frame, f"CNN: {label}", (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 0), 2)
             
-            # Process only the first detected face
-            break
-            
         return frame, {"cnn_prediction": is_drowsy_prediction}