Update app.py

app.py

@@ -4,100 +4,111 @@ import numpy as np
 import gradio as gr
 import tempfile
 import os
+from datetime import timedelta
+
 os.environ['CUDA_VISIBLE_DEVICES'] = "0"
 
-# Function to preprocess each frame
 def preprocess_frame(frame):
-    resized_frame = cv2.resize(frame, (224, 224))  # Adjust size based on your model's input shape
+    resized_frame = cv2.resize(frame, (224, 224))
     normalized_frame = resized_frame / 255.0
-    return np.expand_dims(normalized_frame, axis=0)  # Add batch dimension
+    return np.expand_dims(normalized_frame, axis=0)
 
-# Function to draw pretty label on the frame
 def draw_label(frame, label, position=(50, 50), font_scale=1, thickness=2):
-    # Define label properties
     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]
     
-    # Define rectangle background dimensions
     text_x, text_y = position
-    rect_start = (text_x, text_y - text_size[1] - 10)  # Adjust y to account for text height
-    rect_end = (text_x + text_size[0] + 10, text_y + 10)
+    rect_start = (text_x - 5, text_y - text_size[1] - 15)
+    rect_end = (text_x + text_size[0] + 5, text_y + 5)
     
-    # Draw rectangle background
     cv2.rectangle(frame, rect_start, rect_end, bg_color, -1)
-    
-    # Add border around text
-    cv2.putText(frame, label, position, font, font_scale, (255, 255, 255), thickness + 2, lineType=cv2.LINE_AA)
-    
-    # Add the main colored text
-    cv2.putText(frame, label, position, font, font_scale, color, thickness, lineType=cv2.LINE_AA)
+    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)
 
 @spaces.GPU(duration=120)
 def predict_drowsiness(video_path):
-    # Open the video file
     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))
     
-    # Calculate frame skipping interval based on 0.5 seconds
-    skip_interval = int(fps * 0.5)  # Skip frames to achieve 1 frame every 0.5 seconds
+    skip_interval = int(fps * 0.5)
     
-    # Create a temporary file for the output video
     with tempfile.NamedTemporaryFile(suffix='.mp4', delete=False) as temp_output:
         temp_output_path = temp_output.name
     
-    # Output video settings
     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
         
-        # Only process frames at the specified interval
+        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:
-            # Preprocess frame
             preprocessed_frame = preprocess_frame(frame)
-            
-            # Use the model to predict drowsiness
             prediction = model.predict(preprocessed_frame)
             drowsiness = np.argmax(prediction)
             
-            # Add label to frame with improved visibility
             label = 'Drowsy' if drowsiness == 0 else 'Alert'
-            draw_label(frame, label, position=(50, 50))  # Use the draw_label function
+            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)
         
-        # Write the frame (whether labeled or not) to the output video
         out.write(frame)
         frame_count += 1
     
-    # Release resources
     cap.release()
     out.release()
     
-    return temp_output_path  # Return the path to the temporary output video
+    return temp_output_path
 
-# Gradio interface
 interface = gr.Interface(
     fn=predict_drowsiness,
-    inputs=gr.Video(),  # Video input from webcam or upload
-    outputs="video",  # Return a playable video with predictions
-    title="Drowsiness Detection in Video",
-    description="Upload a video or record one, and this tool will detect if the person is drowsy.",
+    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.",
 )
 
-# Launch the app
 if __name__ == "__main__":
     interface.launch()