Spaces:
Starting
Starting
File size: 3,247 Bytes
c20196f |
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 102 103 104 |
import cv2
from ultralytics import YOLO
import supervision as sv
import numpy as np
import os
from typing import Union, Tuple
class FacialEmotionDetector:
"""
Detect facial emotions from an image or video frame using a YOLO model.
"""
def __init__(self, model_path: str = "best.pt"):
"""
Initialize the detector.
Args:
model_path (str): Path to YOLO model weights (.pt).
"""
if not os.path.exists(model_path):
raise FileNotFoundError(
f"❌ Model file not found at '{model_path}'. "
f"Please ensure 'best.pt' is available in the project directory."
)
# Load YOLO model
self.model = YOLO(model_path)
# Supervision annotators for boxes + labels
self.box_annotator = sv.BoxAnnotator(thickness=2)
self.label_annotator = sv.LabelAnnotator(text_scale=0.5, text_thickness=1)
print("✅ FacialEmotionDetector initialized successfully.")
def detect_emotion(self, frame: np.ndarray) -> Tuple[np.ndarray, Union[str, None]]:
"""
Detect emotions in a single frame.
Args:
frame (np.ndarray): BGR image (OpenCV).
Returns:
Tuple[np.ndarray, str|None]:
- Annotated frame (with boxes + labels).
- Most confident emotion label (or None if no detection).
"""
# YOLO inference
result = self.model(frame, agnostic_nms=True)[0]
# Convert YOLO results → Supervision detections
detections = sv.Detections.from_ultralytics(result)
# Find dominant (highest confidence) detection
dominant_emotion = None
if len(detections) > 0:
most_confident_idx = np.argmax(detections.confidence)
dominant_emotion = detections.data["class_name"][most_confident_idx]
# Build label strings
labels = [
f"{self.model.model.names[class_id]} {confidence:.2f}"
for _, _, confidence, class_id, _, _ in detections
]
# Annotate boxes
annotated = self.box_annotator.annotate(scene=frame.copy(), detections=detections)
# Annotate labels
annotated = self.label_annotator.annotate(scene=annotated, detections=detections, labels=labels)
return annotated, dominant_emotion
if __name__ == "__main__":
# Quick webcam test
try:
detector = FacialEmotionDetector(model_path="best.pt")
cap = cv2.VideoCapture(0)
if not cap.isOpened():
print("❌ Could not open webcam.")
else:
while True:
ret, frame = cap.read()
if not ret:
break
annotated_frame, emotion = detector.detect_emotion(frame)
cv2.imshow("Facial Emotion Detection", annotated_frame)
if emotion:
print(f"Detected Emotion: {emotion}")
if cv2.waitKey(1) & 0xFF == ord("q"):
break
cap.release()
cv2.destroyAllWindows()
except FileNotFoundError as e:
print(e)
except Exception as e:
print(f"⚠️ Unexpected error: {e}")
|