Upload 2 files

main.py

@@ -0,0 +1,61 @@
+import cv2
+from model import EmotionPredictor
+
+cap = cv2.VideoCapture(0)
+predictor = EmotionPredictor()
+
+face_cascade = cv2.CascadeClassifier(
+    cv2.data.haarcascades + "haarcascade_frontalface_default.xml"
+)
+
+if face_cascade.empty():
+    raise RuntimeError("Failed to load Haar Cascade")
+
+FRAME_SKIP = 2
+frame_count = 0
+current_faces = []
+
+
+while True:
+    ret, frame = cap.read()
+    if not ret:
+        break
+
+    frame_count += 1
+
+    if frame_count % FRAME_SKIP == 0:
+        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
+        detected = face_cascade.detectMultiScale(
+            gray, scaleFactor=1.1, minNeighbors=5, minSize=(30, 30)
+        )
+        current_faces = []
+        if len(detected) > 0:
+            x, y, w, h = max(detected, key=lambda r: r[2]*r[3])
+            
+            y1, y2 = max(0, y), min(frame.shape[0], y + h)
+            x1, x2 = max(0, x), min(frame.shape[1], x + w)
+            
+            if y2 > y1 and x2 > x1:
+                face_rgb = cv2.cvtColor(frame[y1:y2, x1:x2], cv2.COLOR_BGR2RGB)
+                label = predictor.predict(face_rgb)
+                current_faces.append((x, y, w, h, label))
+
+    for (x, y, w, h, label) in current_faces:
+        cv2.rectangle(frame, (x, y), (x+w, y+h), (255, 0, 0), 2)
+        cv2.putText(
+            frame, label, (x, y - 10),
+            cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 255, 255), 2
+        )
+
+    cv2.imshow("Emotion Detection", frame)
+
+    if cv2.waitKey(1) & 0xFF == ord("q"):
+        break
+    if cv2.getWindowProperty("Emotion Detection", cv2.WND_PROP_VISIBLE) < 1:
+        break
+
+cap.release()
+cv2.destroyAllWindows()
+
+
+

model.py

@@ -0,0 +1,60 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import pickle
+from torchvision import transforms
+
+import numpy as np
+from PIL import Image
+
+
+class FaceClassifier(nn.Module):
+    def __init__(self, num_classes):
+        super().__init__()
+        self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
+        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
+        self.conv3 = nn.Conv2d(64, 128, 3, padding=1)
+        self.pool = nn.MaxPool2d(2, 2)
+        
+        self.dropout = nn.Dropout(0.1)
+        self.fc1 = nn.Linear(128 * 16 * 16, 512)
+        self.fc2 = nn.Linear(512, num_classes) 
+
+    def forward(self, x):
+        x = self.pool(F.relu(self.conv1(x)))
+        x = self.pool(F.relu(self.conv2(x)))
+        x = self.pool(F.relu(self.conv3(x)))
+        x = x.view(-1, 128 * 16 * 16)
+        x = self.dropout(F.relu(self.fc1(x)))
+        x = self.fc2(x)
+        return x
+
+
+
+class EmotionPredictor:
+    def __init__(self):
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+        with open("classes.pkl", "rb") as f:
+            self.classes = pickle.load(f)
+
+        self.model = FaceClassifier(len(self.classes))
+        self.model.load_state_dict(
+            torch.load("face_classifier.pth", map_location=self.device)
+        )
+        self.model.to(self.device).eval()
+
+        self.transform = transforms.Compose([
+            transforms.Resize((128, 128)),
+            transforms.ToTensor(),
+            transforms.Normalize((0.5,), (0.5,))
+        ])
+    @torch.inference_mode()
+    def predict(self, image_np: np.ndarray) -> str:
+        img = Image.fromarray(image_np)
+        tensor = self.transform(img).unsqueeze(0).to(self.device)
+        output = self.model(tensor)
+        return self.classes[output.argmax(1).item()]
+
+
+