Spaces:
Runtime error
Runtime error
File size: 3,959 Bytes
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 95 96 97 98 99 100 101 |
# 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
import dlib
from PIL import Image
import os
class CnnProcessor(BaseProcessor):
"""
Drowsiness detection using a pre-trained EfficientNet-B7 model.
"""
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()
# Load the model
self.model = self._load_model()
# Define image transformations
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}")
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.load_state_dict(torch.load(self.model_path, map_location=self.device))
model.to(self.device)
model.eval() # Set the model to evaluation mode
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):
"""
Processes a frame to detect drowsiness using the CNN model.
"""
if self.model is None:
return frame, {"cnn_prediction": False}
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
faces = self.face_detector(gray)
is_drowsy_prediction = False
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
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:
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}
|