1

app/app_utils.py

@@ -6,25 +6,22 @@ import cv2
 from pytorch_grad_cam.utils.image import show_cam_on_image
 import matplotlib.pyplot as plt
 
-# Importing necessary components for the Gradio app
-from app.model import pth_model_static, pth_model_dynamic, cam, pth_processing
+# Importing necessary components
+from app.model import load_models
 from app.face_utils import get_box, display_info
 from app.config import DICT_EMO, config_data
 from app.plot import statistics_plot
 
+# Initialize the face mesh detector
 mp_face_mesh = mp.solutions.face_mesh
 
-def get_device():
-    if torch.backends.mps.is_available():
-        return torch.device("mps")
-    elif torch.cuda.is_available():
-        return torch.device("cuda")
-    else:
-        return torch.device("cpu")
-
-device = get_device()
+# Set the device directly
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 print(f"Using device: {device}")
 
+# Load models and GradCAM
+pth_model_static, pth_model_dynamic, cam = load_models()
+
 # Move models to the selected device
 pth_model_static = pth_model_static.to(device)
 pth_model_dynamic = pth_model_dynamic.to(device)
@@ -152,7 +149,7 @@ def preprocess_video_and_predict(video):
                     startX, startY, endX, endY  = get_box(fl, w, h)
                     cur_face = frame_copy[startY:endY, startX: endX]
 
-                    if count_face%config_data.FRAME_DOWNSAMPLING == 0:
+                    if count_face % config_data.FRAME_DOWNSAMPLING == 0:
                         cur_face_copy = pth_processing(Image.fromarray(cur_face)).to(device)
                         with torch.no_grad():
                             features = torch.nn.functional.relu(pth_model_static.extract_features(cur_face_copy)).detach().cpu().numpy()
@@ -228,8 +225,10 @@ def preprocess_video_and_predict(video):
         
     return video, path_save_video_face, path_save_video_hm, stat, au_stat
 
-# The rest of the functions remain the same
-# ...
+def features_to_au_intensities(features):
+    features_np = features.detach().cpu().numpy()[0]
+    au_intensities = (features_np - features_np.min()) / (features_np.max() - features_np.min())
+    return au_intensities[:24]  # Assuming we want 24 AUs
 
 def au_statistics_plot(frames, au_intensities_list):
     fig, ax = plt.subplots(figsize=(12, 6))
@@ -244,90 +243,3 @@ def au_statistics_plot(frames, au_intensities_list):
     ax.legend(bbox_to_anchor=(1.05, 1), loc='upper left')
     plt.tight_layout()
     return fig
-
-def preprocess_video_and_predict_sleep_quality(video):
-    cap = cv2.VideoCapture(video)
-    w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-    h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-    fps = np.round(cap.get(cv2.CAP_PROP_FPS))
-
-    path_save_video_original = 'result_original.mp4'
-    path_save_video_face = 'result_face.mp4'
-    path_save_video_sleep = 'result_sleep.mp4'
-    
-    vid_writer_original = cv2.VideoWriter(path_save_video_original, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
-    vid_writer_face = cv2.VideoWriter(path_save_video_face, cv2.VideoWriter_fourcc(*'mp4v'), fps, (224, 224))
-    vid_writer_sleep = cv2.VideoWriter(path_save_video_sleep, cv2.VideoWriter_fourcc(*'mp4v'), fps, (224, 224))
-
-    frames = []
-    sleep_quality_scores = []
-    eye_bags_images = []
-
-    with mp_face_mesh.FaceMesh(
-    max_num_faces=1,
-    refine_landmarks=False,
-    min_detection_confidence=0.5,
-    min_tracking_confidence=0.5) as face_mesh:
-
-        while cap.isOpened():
-            ret, frame = cap.read()
-            if not ret:
-                break
-
-            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-            results = face_mesh.process(frame_rgb)
-
-            if results.multi_face_landmarks:
-                for fl in results.multi_face_landmarks:
-                    startX, startY, endX, endY = get_box(fl, w, h)
-                    cur_face = frame_rgb[startY:endY, startX:endX]
-                    
-                    sleep_quality_score, eye_bags_image = analyze_sleep_quality(cur_face)
-                    sleep_quality_scores.append(sleep_quality_score)
-                    eye_bags_images.append(cv2.resize(eye_bags_image, (224, 224)))
-
-                    sleep_quality_viz = create_sleep_quality_visualization(cur_face, sleep_quality_score)
-                    
-                    cur_face = cv2.resize(cur_face, (224, 224))
-                    
-                    vid_writer_face.write(cv2.cvtColor(cur_face, cv2.COLOR_RGB2BGR))
-                    vid_writer_sleep.write(sleep_quality_viz)
-
-            vid_writer_original.write(frame)
-            frames.append(len(frames) + 1)
-
-    cap.release()
-    vid_writer_original.release()
-    vid_writer_face.release()
-    vid_writer_sleep.release()
-
-    sleep_stat = sleep_quality_statistics_plot(frames, sleep_quality_scores)
-    
-    if eye_bags_images:
-        average_eye_bags_image = np.mean(np.array(eye_bags_images), axis=0).astype(np.uint8)
-    else:
-        average_eye_bags_image = np.zeros((224, 224, 3), dtype=np.uint8)
-
-    return (path_save_video_original, path_save_video_face, path_save_video_sleep, 
-            average_eye_bags_image, sleep_stat)
-
-def analyze_sleep_quality(face_image):
-    # Placeholder function - implement your sleep quality analysis here
-    sleep_quality_score = np.random.random()
-    eye_bags_image = cv2.resize(face_image, (224, 224))
-    return sleep_quality_score, eye_bags_image
-
-def create_sleep_quality_visualization(face_image, sleep_quality_score):
-    viz = face_image.copy()
-    cv2.putText(viz, f"Sleep Quality: {sleep_quality_score:.2f}", (10, 30), 
-                cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
-    return cv2.cvtColor(viz, cv2.COLOR_RGB2BGR)
-
-def sleep_quality_statistics_plot(frames, sleep_quality_scores):
-    # Placeholder function - implement your statistics plotting here
-    fig, ax = plt.subplots()
-    ax.plot(frames, sleep_quality_scores)
-    ax.set_xlabel('Frame')
-    ax.set_ylabel('Sleep Quality Score')
-    ax.set_title('Sleep Quality Over Time')
-    return fig