Update visualization.py

visualization.py

@@ -260,7 +260,7 @@ def pad_or_trim_array(arr, target_length):
     return arr
 
 def create_heatmap(t, mse_embeddings_filtered, mse_posture_filtered, mse_voice_filtered, fps, total_frames, width):
-    frame_index = int(t * fps)
+    frame_index = min(int(t * fps), len(mse_embeddings_filtered) - 1)
     
     # Normalize the MSE values
     mse_embeddings_norm = normalize_mse(mse_embeddings_filtered)
@@ -282,7 +282,7 @@ def create_heatmap(t, mse_embeddings_filtered, mse_posture_filtered, mse_voice_f
     heatmap_frame[2, :, 2] = (mse_voice_norm[frame_index] * 255).astype(np.uint8)      # Blue channel for voice
 
     return heatmap_frame
-
+    
 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}")
     
@@ -307,6 +307,11 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
                             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)
+
+    print(f"Total frames: {total_frames}")
+    print(f"mse_embeddings length: {len(mse_embeddings)}")
+    print(f"mse_posture length: {len(mse_posture)}")
+    print(f"mse_voice length: {len(mse_voice)}")
     
     # Filter MSE arrays for the most frequent person frames
     mse_embeddings_filtered, mse_posture_filtered, mse_voice_filtered = filter_mse_for_most_frequent_person(df, mse_embeddings, mse_posture, mse_voice, most_frequent_person_frames)