evolve-build/visual_emotion_recognition

1import cv2
2import torch
3import torch.nn as nn
4import numpy as np
5import pickle
6from torchvision import transforms
7from PIL import Image
8
9# Emotion labels
10EMOTION_LABELS = {0: "Angry", 1: "Disgust", 2: "Fear", 3: "Happy", 4: "Sad", 5: "Surprise", 6: "Neutral"}
11
12# Load the trained model
13
14class SimpleFERCNN(nn.Module):
15    def __init__(self, num_classes=7):
16        super(SimpleFERCNN, self).__init__()
17        self.features = nn.Sequential(
18            nn.Conv2d(1, 32, kernel_size=3, padding=1),
19            nn.ReLU(),
20            nn.MaxPool2d(2),
21            nn.Conv2d(32, 64, kernel_size=3, padding=1),
22            nn.ReLU(),
23            nn.MaxPool2d(2),
24            nn.Conv2d(64, 128, kernel_size=3, padding=1),
25            nn.ReLU(),
26            nn.MaxPool2d(2)
27        )
28        self.classifier = nn.Sequential(
29            nn.Linear(128 * 6 * 6, 256),
30            nn.ReLU(),
31            nn.Dropout(0.5),
32            nn.Linear(256, num_classes)
33        )
34
35    def forward(self, x):
36        x = self.features(x)
37        x = x.view(x.size(0), -1)
38        return self.classifier(x)
39
40# Load the saved model
41device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
42model = SimpleFERCNN(num_classes=7).to(device)
43with open("trained_fer_model.pkl", "rb") as f:
44    model.load_state_dict(pickle.load(f))
45model.eval()
46
47
48# Haar cascade for face detection
49face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + "haarcascade_frontalface_default.xml")
50
51# Preprocessing function
52transform = transforms.Compose([
53    transforms.ToPILImage(),
54    transforms.Grayscale(),
55    transforms.Resize((48, 48)),
56    transforms.ToTensor(),
57    transforms.Normalize((0.5,), (0.5,))
58])
59
60# Start webcam capture
61cap = cv2.VideoCapture(0)
62
63while True:
64    ret, frame = cap.read()
65    if not ret:
66        break
67
68    gray_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
69    faces = face_cascade.detectMultiScale(gray_frame, scaleFactor=1.3, minNeighbors=5, minSize=(30, 30))
70
71    for (x, y, w, h) in faces:
72        face = gray_frame[y:y+h, x:x+w]
73        face_tensor = transform(face).unsqueeze(0).to(device)  # Add batch dimension
74
75        # Predict emotion
76        with torch.no_grad():
77            outputs = model(face_tensor)
78            _, predicted = torch.max(outputs, 1)
79            emotion = EMOTION_LABELS[predicted.item()]
80
81        # Draw rectangle around face and show emotion label
82        cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2)
83        cv2.putText(frame, emotion, (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (36, 255, 12), 2)
84
85    # Display video feed
86    cv2.imshow("Real-Time Emotion Detection", frame)
87
88    if cv2.waitKey(1) & 0xFF == ord('q'):  # Press 'q' to quit
89        break
90
91cap.release()
92cv2.destroyAllWindows()