Update visualization.py

visualization.py

@@ -210,7 +210,7 @@ def plot_posture(df, posture_scores, color='blue', anomaly_threshold=3):
 
 def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_voice, output_path, desired_fps, largest_cluster):
     # Filter the DataFrame to only include frames from the largest cluster
-        df_largest_cluster = df[df['Cluster'] == largest_cluster]
+    df_largest_cluster = df[df['Cluster'] == largest_cluster]
     
     # Interpolate mse_voice to match the length of df
     x_voice = np.linspace(0, len(mse_voice)-1, len(mse_voice))
@@ -220,12 +220,8 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
     
     mse_embeddings = mse_embeddings[df['Cluster'] == largest_cluster]
     mse_posture = mse_posture[df['Cluster'] == largest_cluster]
-    mse_voice = mse_voice_interpolated[df['Cluster'] == largest_cluster]
+    mse_voice_interpolated = mse_voice_interpolated[df['Cluster'] == largest_cluster]
 
-    mse_embeddings = mse_embeddings[df['Cluster'] == largest_cluster]
-    mse_posture = mse_posture[df['Cluster'] == largest_cluster]
-    mse_voice = mse_voice[df['Cluster'] == largest_cluster]
-    
     cap = cv2.VideoCapture(video_path)
     original_fps = cap.get(cv2.CAP_PROP_FPS)
     width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
@@ -235,21 +231,21 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
     fourcc = cv2.VideoWriter_fourcc(*'mp4v')
     out = cv2.VideoWriter(output_path, fourcc, original_fps, (width, height + 200))
     
+    # Ensure all MSE arrays have the same length as total_frames
     mse_embeddings = np.interp(np.linspace(0, len(mse_embeddings) - 1, total_frames), 
                                np.arange(len(mse_embeddings)), mse_embeddings)
     mse_posture = np.interp(np.linspace(0, len(mse_posture) - 1, total_frames), 
                             np.arange(len(mse_posture)), mse_posture)
-    mse_voice = np.interp(np.linspace(0, len(mse_voice) - 1, total_frames), 
-                          np.arange(len(mse_voice)), mse_voice)
-
+    mse_voice_interpolated = np.interp(np.linspace(0, len(mse_voice_interpolated) - 1, total_frames), 
+                                       np.arange(len(mse_voice_interpolated)), mse_voice_interpolated)
     
     mse_embeddings_norm = (mse_embeddings - np.min(mse_embeddings)) / (np.max(mse_embeddings) - np.min(mse_embeddings))
     mse_posture_norm = (mse_posture - np.min(mse_posture)) / (np.max(mse_posture) - np.min(mse_posture))
-    mse_voice_norm = (mse_voice - np.min(mse_voice)) / (np.max(mse_voice) - np.min(mse_voice))
+    mse_voice_norm = (mse_voice_interpolated - np.min(mse_voice_interpolated)) / (np.max(mse_voice_interpolated) - np.min(mse_voice_interpolated))
     
-    combined_mse = np.zeros((2, total_frames))
-    combined_mse[0] = mse_embeddings_norm  # Use normalized MSE values for facial
-    combined_mse[1] = mse_posture_norm  # Use normalized MSE values for posture
+    combined_mse = np.zeros((3, total_frames))
+    combined_mse[0] = mse_embeddings_norm
+    combined_mse[1] = mse_posture_norm
     combined_mse[2] = mse_voice_norm
     
     # Custom colormap definition