1

app.py

@@ -14,7 +14,7 @@ TAB_STRUCTURE = [
 def create_demo():
     # Import model-related functions here to ensure spaces is imported first
     from app.model import load_models
-    
+
     # Load models outside of the Gradio blocks
     pth_model_static, pth_model_dynamic, cam = load_models()
 
@@ -28,11 +28,11 @@ def create_demo():
                             with gr.Tab(sub_tab):
                                 create_fn(pth_model_static, pth_model_dynamic, cam)
         gr.HTML(DISCLAIMER_HTML)
-    
+
     return demo
 
 # Create the demo instance
 demo = create_demo()
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()

app/model.py

@@ -20,59 +20,21 @@ STATIC_MODEL_PATH = 'assets/models/FER_static_ResNet50_AffectNet.pt'
 DYNAMIC_MODEL_PATH = 'assets/models/FER_dynamic_LSTM.pt'
 
 def load_model(model_class, model_path, *args, **kwargs):
-    model = model_class(*args, **kwargs).to(device)
+    model = model_class(*args, **kwargs)
     if os.path.exists(model_path):
-        try:
-            model.load_state_dict(torch.load(model_path, map_location=device))
-            model.eval()
-            logger.info(f"Model loaded successfully from {model_path}")
-        except Exception as e:
-            logger.error(f"Error loading model from {model_path}: {str(e)}")
-            logger.info("Initializing with random weights.")
+        model.load_state_dict(torch.load(model_path, map_location=device), strict=False)
+        model.to(device)
+        model.eval()
+        logger.info(f"Loaded model from {model_path}")
     else:
-        logger.warning(f"Model file not found at {model_path}. Initializing with random weights.")
+        logger.error(f"Model file not found: {model_path}")
+        model = model.to(device)
     return model
 
-# Load the static model
-pth_model_static = load_model(ResNet50, STATIC_MODEL_PATH, num_classes=7, channels=3)
-
-# Load the dynamic model
-pth_model_dynamic = load_model(LSTMPyTorch, DYNAMIC_MODEL_PATH, input_size=2048, hidden_size=256, num_layers=2, num_classes=7)
-
-# Set up GradCAM
-target_layers = [pth_model_static.resnet.layer4[-1]]
-cam = GradCAM(model=pth_model_static, target_layers=target_layers)
-
-# Define image preprocessing
-pth_transform = transforms.Compose([
-    transforms.Resize((224, 224)),
-    transforms.ToTensor(),
-    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
-])
-
-def pth_processing(img):
-    img = pth_transform(img).unsqueeze(0).to(device)
-    return img
-
-def predict_emotion(img):
-    with torch.no_grad():
-        output = pth_model_static(pth_processing(img))
-        _, predicted = torch.max(output, 1)
-    return predicted.item()
-
-def get_emotion_probabilities(img):
-    with torch.no_grad():
-        output = nn.functional.softmax(pth_model_static(pth_processing(img)), dim=1)
-    return output.squeeze().cpu().numpy()
-
-def generate_cam(img):
-    input_tensor = pth_processing(img)
-    targets = [ClassifierOutputTarget(predict_emotion(img))]
-    grayscale_cam = cam(input_tensor=input_tensor, targets=targets)
-    return grayscale_cam[0, :]
-
-# Add any other necessary functions or variables here
-
-if __name__ == "__main__":
-    logger.info("Model initialization complete.")
-    # You can add some test code here to verify everything is working correctly
+def load_models():
+    pth_model_static = load_model(ResNet50, STATIC_MODEL_PATH)
+    pth_model_dynamic = load_model(LSTMPyTorch, DYNAMIC_MODEL_PATH)
+    
+    cam = GradCAM(model=pth_model_static, target_layers=[pth_model_static.layer4], use_cuda=device == 'cuda')
+    
+    return pth_model_static, pth_model_dynamic, cam

tabs/FACS_analysis.py

@@ -1,55 +1,25 @@
+from gradio import Interface
+from app.app_utils import preprocess_frame_and_predict_aus
 import gradio as gr
-import cv2
-import numpy as np
 import matplotlib.pyplot as plt
-from app.app_utils import preprocess_frame_and_predict_aus
-
-# Define the AUs associated with stress, anxiety, and depression
-STRESS_AUS = [4, 7, 17, 23, 24]
-ANXIETY_AUS = [1, 2, 4, 5, 20]
-DEPRESSION_AUS = [1, 4, 15, 17]
 
+# Define stress, anxiety, and depression AU mappings
+STRESS_AUS = [1, 2, 4]
+ANXIETY_AUS = [5, 9, 14]
+DEPRESSION_AUS = [15, 17, 20]
 AU_DESCRIPTIONS = {
-    1: "Inner Brow Raiser",
-    2: "Outer Brow Raiser",
-    4: "Brow Lowerer",
-    5: "Upper Lid Raiser",
-    7: "Lid Tightener",
-    15: "Lip Corner Depressor",
-    17: "Chin Raiser",
-    20: "Lip Stretcher",
-    23: "Lip Tightener",
-    24: "Lip Pressor"
+    1: "Inner Brow Raiser", 2: "Outer Brow Raiser", 4: "Brow Lowerer",
+    5: "Upper Lid Raiser", 9: "Nose Wrinkler", 14: "Dimpler",
+    15: "Lip Corner Depressor", 17: "Chin Raiser", 20: "Lip Stretcher"
 }
 
-def normalize_score(score):
-    return max(0, min(1, (score + 1.5) / 3))  # Adjust the range as needed
+def process_video_for_facs(video):
+    frames, avg_au_intensities = preprocess_frame_and_predict_aus(video)
+
+    # Calculate emotional state scores
+    def normalize_score(score):
+        return max(0, min(1, score))
 
-def process_video_for_facs(video_path):
-    cap = cv2.VideoCapture(video_path)
-    frames = []
-    au_intensities_list = []
-    
-    while True:
-        ret, frame = cap.read()
-        if not ret:
-            break
-        
-        processed_frame, au_intensities, _ = preprocess_frame_and_predict_aus(frame)
-        
-        if processed_frame is not None and au_intensities is not None:
-            frames.append(processed_frame)
-            au_intensities_list.append(au_intensities)
-    
-    cap.release()
-    
-    if not frames:
-        return None, None
-    
-    # Calculate average AU intensities
-    avg_au_intensities = np.mean(au_intensities_list, axis=0)
-    
-    # Calculate and normalize emotional state scores
     stress_score = normalize_score(np.mean([avg_au_intensities[au-1] for au in STRESS_AUS if au <= len(avg_au_intensities)]))
     anxiety_score = normalize_score(np.mean([avg_au_intensities[au-1] for au in ANXIETY_AUS if au <= len(avg_au_intensities)]))
     depression_score = normalize_score(np.mean([avg_au_intensities[au-1] for au in DEPRESSION_AUS if au <= len(avg_au_intensities)]))
@@ -82,7 +52,7 @@ def process_video_for_facs(video_path):
     
     return frames[-1], fig  # Return the last processed frame and the plot
 
-def create_facs_analysis_tab():
+def create_facs_analysis_tab(pth_model_static, pth_model_dynamic, cam):
     with gr.Row():
         with gr.Column(scale=1):
             input_video = gr.Video()