Update visualization.py

visualization.py

@@ -264,14 +264,17 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
     width, height = video.w, video.h
     total_frames = int(video.duration * video.fps)
     
-    # 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)
-    
+    # Fill missing MSE values with 0
+    def pad_with_zeros(mse_array, total_frames):
+        if len(mse_array) < total_frames:
+            return np.pad(mse_array, (0, total_frames - len(mse_array)), 'constant', constant_values=0)
+        else:
+            return mse_array[:total_frames]
+
+    mse_embeddings = pad_with_zeros(mse_embeddings, total_frames)
+    mse_posture = pad_with_zeros(mse_posture, total_frames)
+    mse_voice = pad_with_zeros(mse_voice, total_frames)
+
     def combine_video_and_heatmap(t):
         video_frame = video.get_frame(t)
         heatmap_frame = create_heatmap(t, mse_embeddings, mse_posture, mse_voice, video.fps, total_frames, width)