Update visualization.py

visualization.py

@@ -271,21 +271,23 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
 
     # Get video properties
     width, height = video.w, video.h
-    total_frames = int(video.duration * desired_fps)
+    total_frames = int(video.duration * video.fps)
 
-    # Ensure MSE arrays have the same length as total_frames
-    def pad_mse_array(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)
-        return mse_array[:total_frames]
+    # Ensure MSE arrays align with original video frames
+    def align_mse_array(mse_array, original_fps, desired_fps, total_frames):
+        original_times = np.arange(len(mse_array)) / original_fps
+        desired_times = np.arange(total_frames) / desired_fps
+        interpolated_mse = np.interp(desired_times, original_times, mse_array)
+        return interpolated_mse
 
-    mse_embeddings = pad_mse_array(mse_embeddings, total_frames)
-    mse_posture = pad_mse_array(mse_posture, total_frames)
-    mse_voice = pad_mse_array(mse_voice, total_frames)
+    original_fps = len(mse_embeddings) / video.duration
+    mse_embeddings = align_mse_array(mse_embeddings, original_fps, desired_fps, total_frames)
+    mse_posture = align_mse_array(mse_posture, original_fps, desired_fps, total_frames)
+    mse_voice = align_mse_array(mse_voice, original_fps, desired_fps, total_frames)
 
     def combine_video_and_heatmap(t):
-        original_frame = int(t * video.fps)
-        video_frame = video.get_frame(original_frame / video.fps)
+        frame_index = int(t * desired_fps)
+        video_frame = video.get_frame(t)
         
         heatmap_frame = create_heatmap(t, mse_embeddings, mse_posture, mse_voice, desired_fps, total_frames, width)
         heatmap_frame_resized = cv2.resize(heatmap_frame, (width, int(height * 0.2)))