Spaces:
Sleeping
Sleeping
| import cv2 | |
| import spaces | |
| import numpy as np | |
| import gradio as gr | |
| import tempfile | |
| import os | |
| from datetime import timedelta | |
| os.environ['CUDA_VISIBLE_DEVICES'] = "0" | |
| def preprocess_frame(frame): | |
| resized_frame = cv2.resize(frame, (224, 224)) | |
| normalized_frame = resized_frame / 255.0 | |
| return np.expand_dims(normalized_frame, axis=0) | |
| def draw_label(frame, label, position=(50, 50), font_scale=1, thickness=2): | |
| if label == 'Drowsy': | |
| color = (0, 0, 255) # Red for Drowsy | |
| bg_color = (0, 0, 100) # Darker background for Drowsy | |
| else: | |
| color = (0, 255, 0) # Green for Alert | |
| bg_color = (0, 100, 0) # Darker background for Alert | |
| font = cv2.FONT_HERSHEY_SIMPLEX | |
| text_size = cv2.getTextSize(label, font, font_scale, thickness)[0] | |
| text_x, text_y = position | |
| rect_start = (text_x - 5, text_y - text_size[1] - 15) | |
| rect_end = (text_x + text_size[0] + 5, text_y + 5) | |
| cv2.rectangle(frame, rect_start, rect_end, bg_color, -1) | |
| cv2.putText(frame, label, (text_x, text_y), font, font_scale, (255, 255, 255), thickness + 2, lineType=cv2.LINE_AA) | |
| cv2.putText(frame, label, (text_x, text_y), font, font_scale, color, thickness, lineType=cv2.LINE_AA) | |
| def add_timestamp(frame, timestamp): | |
| font = cv2.FONT_HERSHEY_SIMPLEX | |
| cv2.putText(frame, timestamp, (10, frame.shape[0] - 10), font, 0.5, (255, 255, 255), 1, cv2.LINE_AA) | |
| def draw_progress_bar(frame, progress): | |
| frame_width = frame.shape[1] | |
| bar_height = 5 | |
| bar_width = int(frame_width * progress) | |
| cv2.rectangle(frame, (0, 0), (bar_width, bar_height), (0, 255, 0), -1) | |
| cv2.rectangle(frame, (0, 0), (frame_width, bar_height), (255, 255, 255), 1) | |
| def predict_drowsiness(video_path): | |
| import tensorflow as tf | |
| print(tf.config.list_physical_devices("GPU")) | |
| model = tf.keras.models.load_model('cnn.keras') | |
| cap = cv2.VideoCapture(video_path) | |
| frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) | |
| frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) | |
| fps = int(cap.get(cv2.CAP_PROP_FPS)) | |
| total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) | |
| skip_interval = int(fps * 0.2) | |
| with tempfile.NamedTemporaryFile(suffix='.mp4', delete=False) as temp_output: | |
| temp_output_path = temp_output.name | |
| out = cv2.VideoWriter(temp_output_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (frame_width, frame_height)) | |
| frame_count = 0 | |
| drowsy_count = 0 | |
| alert_count = 0 | |
| while cap.isOpened(): | |
| ret, frame = cap.read() | |
| if not ret: | |
| break | |
| progress = frame_count / total_frames | |
| draw_progress_bar(frame, progress) | |
| timestamp = str(timedelta(seconds=int(frame_count/fps))) | |
| add_timestamp(frame, timestamp) | |
| if frame_count % skip_interval == 0: | |
| preprocessed_frame = preprocess_frame(frame) | |
| prediction = model.predict(preprocessed_frame) | |
| drowsiness = np.argmax(prediction) | |
| label = 'Drowsy' if drowsiness == 0 else 'Alert' | |
| draw_label(frame, label, position=(50, 50)) | |
| if label == 'Drowsy': | |
| drowsy_count += 1 | |
| else: | |
| alert_count += 1 | |
| # Add drowsiness statistics | |
| stats_text = f"Drowsy: {drowsy_count} | Alert: {alert_count}" | |
| cv2.putText(frame, stats_text, (frame_width - 200, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1, cv2.LINE_AA) | |
| out.write(frame) | |
| frame_count += 1 | |
| cap.release() | |
| out.release() | |
| return temp_output_path | |
| interface = gr.Interface( | |
| fn=predict_drowsiness, | |
| inputs=gr.Video(), | |
| outputs="video", | |
| title="Enhanced Drowsiness Detection in Video", | |
| description="Upload a video or record one to detect drowsiness with improved visuals and statistics.", | |
| examples=["003_nightglasses_mix.mp4"] | |
| ) | |
| if __name__ == "__main__": | |
| interface.launch() |