Update visualization.py

visualization.py

@@ -219,6 +219,10 @@ def plot_posture(df, posture_scores, color='blue', anomaly_threshold=3):
 
 
 def filter_mse_for_most_frequent_person(df, mse_embeddings, mse_posture, mse_voice, most_frequent_person_frames):
+    # Ensure most_frequent_person_frames is a list
+    if not isinstance(most_frequent_person_frames, (list, np.ndarray)):
+        most_frequent_person_frames = [most_frequent_person_frames]
+
     # Create a mask for the most frequent person frames
     mask = df['Frame'].isin(most_frequent_person_frames)
 
@@ -229,6 +233,7 @@ def filter_mse_for_most_frequent_person(df, mse_embeddings, mse_posture, mse_voi
 
     return mse_embeddings_filtered, mse_posture_filtered, mse_voice_filtered
 
+
 def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_voice, output_folder, desired_fps, most_frequent_person_frames):
     print(f"Creating heatmap video. Output folder: {output_folder}")
     
@@ -282,7 +287,7 @@ 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)