Update visualization.py

visualization.py

@@ -222,7 +222,7 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, outpu
     
     # Create custom colormap
     cmap = mcolors.LinearSegmentedColormap.from_list("custom", 
-        [(0, 0, 1), (1, 1, 1), (0.5, 0, 0.5)], N=256)
+        [(1, 1, 1), (0, 0, 1), (0.5, 0, 0.5)], N=256)
     
     # Ensure heatmap data covers all frames
     mse_embeddings = np.interp(np.linspace(0, len(mse_embeddings) - 1, total_frames), 
@@ -235,22 +235,19 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, outpu
     mse_posture_norm = (mse_posture - np.min(mse_posture)) / (np.max(mse_posture) - np.min(mse_posture))
     
     # Combine MSEs
-    combined_mse = np.zeros((2, total_frames))
-    combined_mse[0] = mse_embeddings_norm
-    combined_mse[1] = mse_posture_norm
+    combined_mse = np.zeros((2, total_frames, 3))
+    combined_mse[0] = np.array([1 - mse_embeddings_norm, 1 - mse_embeddings_norm, mse_embeddings_norm]).T  # RGB for facial
+    combined_mse[1] = np.array([1 - mse_posture_norm, mse_posture_norm, 1 - mse_posture_norm]).T  # RGB for posture
     
     fig, ax = plt.subplots(figsize=(width/100, 2))
-    im = ax.imshow(combined_mse, aspect='auto', cmap=cmap, extent=[0, total_frames, 0, 2], vmin=0, vmax=1)
+    im = ax.imshow(combined_mse, aspect='auto', extent=[0, total_frames, 0, 2])
     ax.set_yticks([0.5, 1.5])
     ax.set_yticklabels(['Facial', 'Posture'])
     ax.set_xticks([])
-    cbar = plt.colorbar(im, ax=ax, orientation='horizontal', pad=0.1, aspect=20, shrink=0.5)
-    cbar.set_ticks([0, 0.5, 1])
-    cbar.set_ticklabels(['Low MSE', 'Medium MSE', 'High MSE'])
     plt.tight_layout()
     
     line = None
-    frame_interval = max(1, int(original_fps / desired_fps))
+    frame_interval = int(original_fps / desired_fps)
     
     for frame_count in range(0, total_frames, frame_interval):
         cap.set(cv2.CAP_PROP_POS_FRAMES, frame_count)