asrcoddeploy commited on
Commit
c5c6603
·
verified ·
1 Parent(s): 6f93339

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +85 -0
app.py ADDED
@@ -0,0 +1,85 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import torch
3
+ import torch.nn as nn
4
+ from torchvision import models, transforms
5
+ from PIL import Image
6
+ import cv2
7
+ import numpy as np
8
+
9
+ # 1. Download OpenCV Haarcascades for eye tracking
10
+ import urllib.request
11
+ import os
12
+ cascade_path = 'haarcascade_eye.xml'
13
+ if not os.path.exists(cascade_path):
14
+ urllib.request.urlretrieve(
15
+ 'https://raw.githubusercontent.com/opencv/opencv/master/data/haarcascades/haarcascade_eye.xml',
16
+ cascade_path
17
+ )
18
+ eye_cascade = cv2.CascadeClassifier(cascade_path)
19
+
20
+ # 2. Re-initialize and load the model (Using your weights)
21
+ model = models.mobilenet_v2(weights=None)
22
+ # Ensure this matches exactly how you defined it in the Masterpiece training step
23
+ model.classifier = nn.Sequential(
24
+ nn.Dropout(p=0.5),
25
+ nn.Linear(model.last_channel, 2)
26
+ )
27
+ model.load_state_dict(torch.load('ddobj_model.pth', map_location=torch.device('cpu')))
28
+ model.eval()
29
+
30
+ # 3. Transforms (Grayscale is key to matching the MRL dataset!)
31
+ transform = transforms.Compose([
32
+ transforms.Grayscale(num_output_channels=3), # Convert to 3-channel grayscale
33
+ transforms.Resize((224, 224)),
34
+ transforms.ToTensor(),
35
+ transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
36
+ ])
37
+
38
+ # 4. The Smart Prediction Function
39
+ def predict_drowsiness(image):
40
+ # Convert Gradio image to OpenCV format
41
+ img_cv = np.array(image)
42
+ gray = cv2.cvtColor(img_cv, cv2.COLOR_RGB2GRAY)
43
+
44
+ # Detect eyes in the image
45
+ eyes = eye_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5, minSize=(30, 30))
46
+
47
+ if len(eyes) == 0:
48
+ return "ERROR: Could not detect any eyes in the image. Please upload a clear face photo.", None
49
+
50
+ # Take the first detected eye (largest/clearest)
51
+ (x, y, w, h) = eyes[0]
52
+
53
+ # Crop the eye from the original image
54
+ eye_crop = img_cv[y:y+h, x:x+w]
55
+
56
+ # Convert the cropped eye back to PIL for PyTorch
57
+ eye_pil = Image.fromarray(eye_crop)
58
+ input_tensor = transform(eye_pil).unsqueeze(0)
59
+
60
+ # Run the model
61
+ with torch.no_grad():
62
+ outputs = model(input_tensor)
63
+ _, predicted = torch.max(outputs, 1)
64
+
65
+ classes = ["DROWSY ALERT! 🚨 (Eyes Closed)", "NOT DROWSY ✅ (Eyes Open)"]
66
+ result = classes[predicted.item()]
67
+
68
+ # Return the prediction AND show the user the exact crop the model looked at
69
+ return result, eye_pil
70
+
71
+ # 5. Build the UI
72
+ interface = gr.Interface(
73
+ fn=predict_drowsiness,
74
+ inputs=gr.Image(label="Upload Full Face Photo"),
75
+ outputs=[
76
+ gr.Textbox(label="DDobj System Status"),
77
+ gr.Image(label="What the AI saw (Eye Crop)")
78
+ ],
79
+ title="DDobj: Driver Drowsiness Detection",
80
+ description="Upload a photo. The system will automatically locate the eyes, isolate them, and analyze them for fatigue.",
81
+ theme="default"
82
+ )
83
+
84
+ if __name__ == "__main__":
85
+ interface.launch()