Update visualization.py

visualization.py

@@ -233,18 +233,20 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, outpu
     combined_mse[0] = mse_embeddings_norm  # Use normalized MSE values for facial
     combined_mse[1] = mse_posture_norm  # Use normalized MSE values for posture
 
+    # Custom colormap definition
     cdict = {
-    'red':   [(0.0,  1.0, 1.0),  # Neutral (low MSE) to more red (high MSE)
-              (1.0,  1.0, 1.0)],
-    'green': [(0.0,  1.0, 1.0),  # Neutral (low MSE) to less green (high MSE)
-              (1.0,  0.0, 0.0)],
-    'blue':  [(0.0,  1.0, 1.0),  # Neutral (low MSE) to less blue (high MSE)
-              (1.0,  0.0, 0.0)]
-}
+        'red':   [(0.0,  0.0, 0.0),  # Low MSE: No red
+                  (1.0,  1.0, 1.0)], # High MSE: Full red
+        'green': [(0.0,  1.0, 1.0),  # Low MSE: Full green
+                  (1.0,  0.0, 0.0)], # High MSE: No green
+        'blue':  [(0.0,  1.0, 1.0),  # Low MSE: Full blue
+                  (1.0,  0.0, 0.0)]  # High MSE: No blue
+    }
 
     custom_cmap = LinearSegmentedColormap('custom_cmap', segmentdata=cdict, N=256)
     
     fig, ax = plt.subplots(figsize=(width/100, 2))
+    # Use the custom colormap in the heatmap generation
     im = ax.imshow(combined_mse, aspect='auto', cmap=custom_cmap, extent=[0, total_frames, 0, 2])
     ax.set_yticks([0.5, 1.5])
     ax.set_yticklabels(['Face', 'Posture'])
@@ -280,4 +282,4 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, outpu
     out.release()
     plt.close(fig)
     
-    return output_path
+    return output_path