Update visualization.py

visualization.py

@@ -231,8 +231,8 @@ 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))
     
     combined_mse = np.zeros((2, total_frames, 3))
-    combined_mse[0] = np.array([1, 1 - mse_embeddings_norm, 1 - mse_embeddings_norm]).T  # More reddish for facial
-    combined_mse[1] = np.array([1, 1 - mse_posture_norm, 1 - mse_posture_norm]).T  # More reddish for posture
+    combined_mse[0] = np.array([1 - mse_embeddings_norm, 1 - mse_embeddings_norm, 1 - mse_embeddings_norm]).T.reshape(-1, 3)  # More reddish for facial
+    combined_mse[1] = np.array([1 - mse_posture_norm, 1 - mse_posture_norm, 1 - mse_posture_norm]).T.reshape(-1, 3)  # More reddish for posture
     
     fig, ax = plt.subplots(figsize=(width/100, 2))
     im = ax.imshow(combined_mse, aspect='auto', extent=[0, total_frames, 0, 2])