Update visualization.py

visualization.py

@@ -264,50 +264,25 @@ 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)
     
-    def fill_with_previous_values(mse_array, total_frames):
+    def fill_with_zeros_and_values(mse_array, total_frames):
         result = np.zeros(total_frames)
         indices = np.linspace(0, total_frames - 1, len(mse_array)).astype(int)
         result[indices] = mse_array
-        for i in range(1, total_frames):
-            if result[i] == 0:
-                result[i] = result[i - 1]
         return result
 
     # Ensure all MSE arrays have the same length as total_frames
-    mse_embeddings = fill_with_previous_values(mse_embeddings, total_frames)
-    mse_posture = fill_with_previous_values(mse_posture, total_frames)
-    mse_voice = fill_with_previous_values(mse_voice, total_frames)
+    mse_embeddings = fill_with_zeros_and_values(mse_embeddings, total_frames)
+    mse_posture = fill_with_zeros_and_values(mse_posture, total_frames)
+    mse_voice = fill_with_zeros_and_values(mse_voice, total_frames)
     
-    def create_heatmap(t, mse_embeddings, mse_posture, mse_voice, fps, total_frames, width, height):
-        frame_index = int(t * fps)
-        mse_face = mse_embeddings[frame_index]
-        mse_body = mse_posture[frame_index]
-        mse_audio = mse_voice[frame_index]
-
-        max_mse_face = np.max(mse_embeddings)
-        max_mse_body = np.max(mse_posture)
-        max_mse_audio = np.max(mse_voice)
-
-        def get_heatmap_color(mse_value, max_mse_value):
-            normalized_value = mse_value / max_mse_value if max_mse_value != 0 else 0
-            color = plt.cm.hot(normalized_value)
-            return (color[0] * 255, color[1] * 255, color[2] * 255)  # Convert to RGB
-
-        heatmap_face = np.full((height, width, 3), get_heatmap_color(mse_face, max_mse_face), dtype=np.uint8)
-        heatmap_body = np.full((height, width, 3), get_heatmap_color(mse_body, max_mse_body), dtype=np.uint8)
-        heatmap_audio = np.full((height, width, 3), get_heatmap_color(mse_audio, max_mse_audio), dtype=np.uint8)
-
-        # Combine heatmaps
-        combined_heatmap = (heatmap_face + heatmap_body + heatmap_audio) / 3
-        return combined_heatmap.astype(np.uint8)
-
-    def combine_video_and_heatmap(get_frame, t):
-        video_frame = get_frame(t)
-        heatmap_frame = create_heatmap(t, mse_embeddings, mse_posture, mse_voice, desired_fps, total_frames, width, height)
-        combined_frame = cv2.addWeighted(video_frame, 0.7, heatmap_frame, 0.3, 0)
+    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)
+        heatmap_frame_resized = cv2.resize(heatmap_frame, (width, heatmap_frame.shape[0]))
+        combined_frame = np.vstack((video_frame, heatmap_frame_resized))
         return combined_frame
 
-    final_clip = video.fl(combine_video_and_heatmap)
+    final_clip = VideoClip(combine_video_and_heatmap, duration=video.duration)
     final_clip = final_clip.set_audio(video.audio)
 
     # Write the final video
@@ -323,7 +298,7 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
         return heatmap_video_path
     else:
         print(f"Failed to create heatmap video at: {heatmap_video_path}")
-        return None
+        return None] 
 
 
 # Function to create the correlation heatmap