Update visualization.py

visualization.py

@@ -249,32 +249,66 @@ def create_heatmap(t, mse_embeddings, mse_posture, mse_voice, video_fps, total_f
 
 def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_voice, output_folder, desired_fps, largest_cluster):
     print(f"Creating heatmap video. Output folder: {output_folder}")
-    
     os.makedirs(output_folder, exist_ok=True)
-    
     output_filename = os.path.basename(video_path).rsplit('.', 1)[0] + '_heatmap.mp4'
     heatmap_video_path = os.path.join(output_folder, output_filename)
-    
     print(f"Heatmap video will be saved at: {heatmap_video_path}")
 
     # Load the original video
     video = VideoFileClip(video_path)
-    
+
     # Get video properties
     width, height = video.w, video.h
     total_frames = int(video.duration * video.fps)
-    
-    def fill_with_zeros_and_values(mse_array, total_frames):
+
+    def fill_with_previous_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_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)
-    
+    # Fill gaps with previous values
+    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)
+
+    # Find max values for each variable
+    max_embeddings = np.max(mse_embeddings)
+    max_posture = np.max(mse_posture)
+    max_voice = np.max(mse_voice)
+
+    def create_heatmap(t, mse_embeddings, mse_posture, mse_voice, fps, total_frames, width):
+        frame_index = int(t * fps)
+        
+        # Create separate color maps for each variable
+        cmap_embeddings = plt.get_cmap('Reds')
+        cmap_posture = plt.get_cmap('Reds')
+        cmap_voice = plt.get_cmap('Reds')
+
+        # Normalize values within their own scales
+        norm_embeddings = mse_embeddings[frame_index] / max_embeddings
+        norm_posture = mse_posture[frame_index] / max_posture
+        norm_voice = mse_voice[frame_index] / max_voice
+
+        # Create color arrays for each variable
+        color_embeddings = cmap_embeddings(norm_embeddings)
+        color_posture = cmap_posture(norm_posture)
+        color_voice = cmap_voice(norm_voice)
+
+        # Create the heatmap frame
+        heatmap_height = 100  # Adjust as needed
+        heatmap_frame = np.zeros((heatmap_height, width, 4))
+
+        # Fill the heatmap sections
+        heatmap_frame[:heatmap_height//3, :] = color_embeddings
+        heatmap_frame[heatmap_height//3:2*heatmap_height//3, :] = color_posture
+        heatmap_frame[2*heatmap_height//3:, :] = color_voice
+
+        return (heatmap_frame * 255).astype(np.uint8)
+
     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)
@@ -287,18 +321,18 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
 
     # Write the final video
     final_clip.write_videofile(heatmap_video_path, codec='libx264', audio_codec='aac', fps=video.fps)
-    
+
     # Close the video clips
     video.close()
     final_clip.close()
-    
+
     if os.path.exists(heatmap_video_path):
         print(f"Heatmap video created at: {heatmap_video_path}")
         print(f"Heatmap video size: {os.path.getsize(heatmap_video_path)} bytes")
         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