Update app.py

app.py

@@ -6,6 +6,7 @@ import torch.nn as nn
 import torch.optim as optim
 from facenet_pytorch import InceptionResnetV1, MTCNN
 import mediapipe as mp
+from fer import FER
 from sklearn.cluster import KMeans
 from sklearn.preprocessing import StandardScaler, MinMaxScaler
 from sklearn.metrics import silhouette_score
@@ -24,6 +25,7 @@ mtcnn = MTCNN(keep_all=False, device=device, thresholds=[0.999, 0.999, 0.999], m
 model = InceptionResnetV1(pretrained='vggface2').eval().to(device)
 mp_face_mesh = mp.solutions.face_mesh
 face_mesh = mp_face_mesh.FaceMesh(static_image_mode=False, max_num_faces=1, min_detection_confidence=0.5)
+emotion_detector = FER(mtcnn=False)
 
 def frame_to_timecode(frame_num, original_fps, desired_fps):
     total_seconds = frame_num / original_fps
@@ -40,9 +42,12 @@ def get_face_embedding_and_emotion(face_img):
     with torch.no_grad():
         embedding = model(face_tensor)
 
-    # Placeholder for emotion detection
-    emotion_dict = {e: np.random.random() for e in ['angry', 'disgust', 'fear', 'happy', 'sad', 'surprise', 'neutral']}
-    
+    emotions = emotion_detector.detect_emotions(face_img)
+    if emotions:
+        emotion_dict = emotions[0]['emotions']
+    else:
+        emotion_dict = {e: 0 for e in ['angry', 'disgust', 'fear', 'happy', 'sad', 'surprise', 'neutral']}
+
     return embedding.cpu().numpy().flatten(), emotion_dict
 
 def alignFace(img):
@@ -278,7 +283,7 @@ def process_video(video_path, num_anomalies, num_components, desired_fps, batch_
         progress(0.4, "Organizing faces")
         organize_faces_by_person(embeddings_by_frame, clusters, aligned_faces_folder, organized_faces_folder)
 
-        progress(0.5, "Saving person data")
+progress(0.5, "Saving person data")
         df, largest_cluster = save_person_data_to_csv(embeddings_by_frame, emotions_by_frame, clusters, desired_fps, original_fps, temp_dir, num_components)
 
         progress(0.6, "Performing anomaly detection")