Updated expression recognition

app/Hackathon_setup/exp_recognition.py

@@ -4,7 +4,7 @@ import torch
 from PIL import Image
 import os
 
-# Remove '.' if running locally; keep it for server (Spaces)
+# Keep the dot when deploying (Spaces); remove locally if needed
 from .exp_recognition_model import facExpRec, processor, device
 
 #############################################################################################################################
@@ -20,10 +20,14 @@ current_path = os.path.dirname(os.path.abspath(__file__))
 # =====================================================
 # LOAD MODEL ONCE (global)
 # =====================================================
-print("🔹 Loading Expression Recognition Model...")
-exp_model = facExpRec().to(device)
-exp_model.eval()
-print("✅ Expression model loaded successfully.")
+print("Loading Expression Recognition Model...")
+try:
+    exp_model = facExpRec().to(device)
+    exp_model.eval()
+    print("Expression model loaded successfully.")
+except Exception as e:
+    print(f"Failed to load Expression Model: {e}")
+    exp_model = None
 
 # =====================================================
 # FACE DETECTION FUNCTION
@@ -43,8 +47,7 @@ def detected_face(image):
     if len(faces) == 0:
         return 0
 
-    face_areas = []
-    images = []
+    face_areas, images = [], []
     for (x, y, w, h) in faces:
         face_cropped = gray[y:y + h, x:x + w]
         face_areas.append(w * h)
@@ -64,6 +67,9 @@ def get_expression(img):
     Takes an OpenCV BGR image as input, detects the face, and returns the
     predicted facial expression as a string.
     """
+    if exp_model is None:
+        raise RuntimeError("Expression model not loaded properly.")
+
     # Detect face
     face = detected_face(img)
     if face == 0:
@@ -80,8 +86,12 @@ def get_expression(img):
         pred_idx = torch.argmax(probs, dim=-1).item()
         confidence = probs[0][pred_idx].item()
 
-    # Map index → expression label
-    expression_label = exp_model.processor.config.id2label.get(pred_idx, "Unknown")
+    # Use model.config.id2label safely
+    id2label = getattr(exp_model.model.config, "id2label", None)
+    if id2label and pred_idx in id2label:
+        expression_label = id2label[pred_idx]
+    else:
+        expression_label = "Unknown"
 
     print(f"Detected Expression: {expression_label} ({confidence:.2f})")
 

app/Hackathon_setup/exp_recognition_model.py

@@ -16,19 +16,6 @@ from PIL import Image
 # =====================================================
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# =====================================================
-# CLASS DEFINITIONS
-# =====================================================
-classes = {
-    0: 'ANGER',
-    1: 'DISGUST',
-    2: 'FEAR',
-    3: 'HAPPINESS',
-    4: 'NEUTRAL',
-    5: 'SADNESS',
-    6: 'SURPRISE'
-}
-
 # =====================================================
 # MODEL NAME (ViT pretrained on facial expressions)
 # =====================================================
@@ -42,9 +29,9 @@ try:
     model = AutoModelForImageClassification.from_pretrained(MODEL_NAME)
     model.to(device)
     model.eval()
-    print(f"✅ Loaded ViT model '{MODEL_NAME}' successfully on {device}.")
+    print(f"Loaded ViT model '{MODEL_NAME}' successfully on {device}.")
 except Exception as e:
-    print(f"❌ Error loading ViT model: {e}")
+    print(f"Error loading ViT model: {e}")
     processor, model = None, None
 
 # =====================================================
@@ -74,6 +61,7 @@ class facExpRec(nn.Module):
             raise RuntimeError("ViT model not loaded correctly.")
         self.model = model
         self.processor = processor
+        self.id2label = getattr(self.model.config, "id2label", None)
 
     def forward(self, x):
         """
@@ -91,8 +79,15 @@ class facExpRec(nn.Module):
             pred_idx = torch.argmax(probs, dim=-1).item()
             confidence = probs[0][pred_idx].item()
 
+        # Prefer model.config.id2label if available
+        label = (
+            self.id2label[pred_idx]
+            if self.id2label and pred_idx in self.id2label
+            else "Unknown"
+        )
+
         return {
-            "expression": classes.get(pred_idx, "Unknown"),
+            "expression": label,
             "confidence": round(confidence, 3)
         }