shape_predictor / main.py
Rudraaaa76's picture
Upload 4 files
f4a2fda verified
Raw
History Blame Contribute Delete
4.06 kB
import cv2
import dlib
import numpy as np
import base64
from fastapi import FastAPI, WebSocket
from scipy.spatial import distance as dist
from imutils import face_utils
import uvicorn
import json
app = FastAPI()
import os
# --- CONFIGURATION FOR HUGGING FACE ---
# Try multiple paths for the model file
possible_paths = [
"shape_predictor_68_face_landmarks.dat",
"data/models/shape_predictor_68_face_landmarks.dat",
"/app/shape_predictor_68_face_landmarks.dat",
"/app/data/models/shape_predictor_68_face_landmarks.dat"
]
predictor_path = None
for path in possible_paths:
if os.path.exists(path):
predictor_path = path
break
if not predictor_path:
print("[ERROR] shape_predictor_68_face_landmarks.dat not found in any of the searched paths.")
print(f"Current working directory: {os.getcwd()}")
print(f"Files in current directory: {os.listdir('.')}")
else:
print(f"[INFO] Using model at: {predictor_path}")
print("[INFO] Loading AI models...")
detector = dlib.get_frontal_face_detector()
try:
if predictor_path:
predictor = dlib.shape_predictor(predictor_path)
print("[SUCCESS] Model loaded!")
else:
print("[ERROR] Predictor path is None, cannot load model.")
except Exception as e:
print(f"[ERROR] Failed to load model: {e}")
# Landmark indexes
(lStart, lEnd) = face_utils.FACIAL_LANDMARKS_IDXS["left_eye"]
(rStart, rEnd) = face_utils.FACIAL_LANDMARKS_IDXS["right_eye"]
(mStart, mEnd) = face_utils.FACIAL_LANDMARKS_IDXS["mouth"]
(nStart, nEnd) = face_utils.FACIAL_LANDMARKS_IDXS["nose"]
def get_ear(eye):
a = dist.euclidean(eye[1], eye[5])
b = dist.euclidean(eye[2], eye[4])
c = dist.euclidean(eye[0], eye[3])
return (a + b) / (2.0 * c)
def get_mar(mouth):
a = dist.euclidean(mouth[13], mouth[19])
b = dist.euclidean(mouth[14], mouth[18])
c = dist.euclidean(mouth[15], mouth[17])
d = dist.euclidean(mouth[12], mouth[16])
return (a + b + c) / (2.0 * d + 1e-6)
@app.websocket("/ws/analyze")
async def websocket_endpoint(websocket: WebSocket):
await websocket.accept()
counter = 0
print("[INFO] Mobile App connected!")
try:
while True:
data = await websocket.receive_text()
header, encoded = data.split(",", 1)
nparr = np.frombuffer(base64.b64decode(encoded), np.uint8)
frame = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
if frame is None:
continue
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
faces = detector(gray, 0)
response = {
"drowsy": False,
"distracted": False,
"yawning": False,
"ear": 0,
"mar": 0
}
for rect in faces:
shape = predictor(gray, rect)
shape = face_utils.shape_to_np(shape)
# Detection Logic
ear = (get_ear(shape[lStart:lEnd]) + get_ear(shape[rStart:rEnd])) / 2.0
mar = get_mar(shape[mStart:mEnd])
nose_center = shape[nStart:nEnd].mean(axis=0)
dist_left = dist.euclidean(shape[lStart:lEnd].mean(axis=0), nose_center)
dist_right = dist.euclidean(shape[rStart:rEnd].mean(axis=0), nose_center)
gaze_ratio = dist_left / (dist_right + 1e-6)
response["ear"] = round(ear, 3)
response["mar"] = round(mar, 3)
if ear < 0.22: counter += 1
else: counter = 0
if counter >= 5: response["drowsy"] = True
if gaze_ratio < 0.7 or gaze_ratio > 1.3: response["distracted"] = True
if mar > 0.6: response["yawning"] = True
await websocket.send_json(response)
except Exception as e:
print(f"[ERROR] Connection closed: {e}")
if __name__ == "__main__":
# Port 7860 is required for Hugging Face Spaces
uvicorn.run(app, host="0.0.0.0", port=7860)