Update visualization.py

visualization.py

@@ -232,10 +232,14 @@ def filter_mse_for_most_frequent_person(df, mse_embeddings, mse_posture, mse_voi
     # Create a mask for the most frequent person frames
     mask = df['Frame'].isin(most_frequent_person_frames)
 
+    # Pad mask to match the length of the video frames
+    padded_mask = np.zeros(len(mse_embeddings), dtype=bool)
+    padded_mask[:len(mask)] = mask
+
     # Apply the mask to filter the MSE arrays
-    mse_embeddings_filtered = np.where(mask, mse_embeddings, 0)
-    mse_posture_filtered = np.where(mask, mse_posture, 0)
-    mse_voice_filtered = np.where(mask, mse_voice, 0)
+    mse_embeddings_filtered = np.where(padded_mask, mse_embeddings, 0)
+    mse_posture_filtered = np.where(padded_mask, mse_posture, 0)
+    mse_voice_filtered = np.where(padded_mask, mse_voice, 0)
 
     return mse_embeddings_filtered, mse_posture_filtered, mse_voice_filtered
 
@@ -292,7 +296,6 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
         print(f"Failed to create heatmap video at: {heatmap_video_path}")
         return None
 
-# Define the create_heatmap function
 def create_heatmap(t, mse_embeddings, mse_posture, mse_voice, video_fps, total_frames, video_width):
     frame_count = int(t * video_fps)