Update visualization.py

visualization.py

@@ -278,7 +278,7 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
     mse_posture = fill_with_previous_values(mse_posture, total_frames)
     mse_voice = fill_with_previous_values(mse_voice, total_frames)
     
-    def create_heatmap(t, mse_embeddings, mse_posture, mse_voice, fps, total_frames, width):
+    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]
@@ -293,17 +293,18 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
             color = plt.cm.hot(normalized_value)
             return (color[0] * 255, color[1] * 255, color[2] * 255)  # Convert to RGB
 
-        heatmap_face = np.full((height // 3, width, 3), get_heatmap_color(mse_face, max_mse_face), dtype=np.uint8)
-        heatmap_body = np.full((height // 3, width, 3), get_heatmap_color(mse_body, max_mse_body), dtype=np.uint8)
-        heatmap_audio = np.full((height // 3, width, 3), get_heatmap_color(mse_audio, max_mse_audio), dtype=np.uint8)
+        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)
 
-        heatmap = np.vstack((heatmap_face, heatmap_body, heatmap_audio))
-        return heatmap
+        # 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)
-        combined_frame = np.vstack((video_frame, heatmap_frame))
+        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)
         return combined_frame
 
     final_clip = video.fl(combine_video_and_heatmap)