Spaces:
Sleeping
Sleeping
File size: 3,602 Bytes
19f420a 42e8aa3 19f420a 42e8aa3 19f420a 42e8aa3 19f420a 42e8aa3 19f420a 42e8aa3 19f420a 42e8aa3 19f420a 42e8aa3 19f420a 42e8aa3 19f420a 42e8aa3 19f420a 42e8aa3 19f420a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 |
# drive_paddy/detection/strategies/cnn_model.py
from src.detection.base_processor import BaseProcessor
import numpy as np
import torch
import torchvision.transforms as transforms
from torchvision.models import efficientnet_b7
import cv2
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")
# dlib is no longer needed.
# self.face_detector = dlib.get_frontal_face_detector()
self.model = self._load_model()
self.transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
def _load_model(self):
"""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}")
return None
try:
model = efficientnet_b7()
num_ftrs = model.classifier[1].in_features
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()
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, face_landmarks=None):
"""
Processes a frame using the CNN model with pre-supplied landmarks.
"""
if self.model is None or face_landmarks is None:
return frame, {"cnn_prediction": False}
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))
# 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)
with torch.no_grad():
outputs = self.model(image_tensor)
_, preds = torch.max(outputs, 1)
if preds.item() == 1: # Assuming class 1 is 'drowsy'
is_drowsy_prediction = True
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)
return frame, {"cnn_prediction": is_drowsy_prediction}
|