Update visualization.py

visualization.py

@@ -1,6 +1,7 @@
 import matplotlib.pyplot as plt
 from matplotlib.backends.backend_agg import FigureCanvasAgg
 import matplotlib.colors as mcolors
+from matplotlib.colors import LinearSegmentedColormap
 import seaborn as sns
 import numpy as np
 import pandas as pd
@@ -138,6 +139,7 @@ def plot_mse_histogram(mse_values, title, anomaly_threshold, color='blue'):
     plt.close()
     return fig
 
+
 def plot_mse_heatmap(mse_values, title, df):
     plt.figure(figsize=(20, 3), dpi=300)
     fig, ax = plt.subplots(figsize=(20, 3))
@@ -219,9 +221,6 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, outpu
     fourcc = cv2.VideoWriter_fourcc(*'mp4v')
     out = cv2.VideoWriter(output_path, fourcc, original_fps, (width, height + 200))
     
-    cmap = mcolors.LinearSegmentedColormap.from_list("custom", 
-        [(1, 1, 1), (1, 0, 0)], N=256)  # More reddish for higher MSE values
-    
     mse_embeddings = np.interp(np.linspace(0, len(mse_embeddings) - 1, total_frames), 
                                np.arange(len(mse_embeddings)), mse_embeddings)
     mse_posture = np.interp(np.linspace(0, len(mse_posture) - 1, total_frames), 
@@ -230,12 +229,23 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, outpu
     mse_embeddings_norm = (mse_embeddings - np.min(mse_embeddings)) / (np.max(mse_embeddings) - np.min(mse_embeddings))
     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 - 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
+    combined_mse = np.zeros((2, total_frames))
+    combined_mse[0] = mse_embeddings_norm  # Use normalized MSE values for facial
+    combined_mse[1] = mse_posture_norm  # Use normalized MSE values for posture
+
+    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)]
+}
+
+    custom_cmap = LinearSegmentedColormap('custom_cmap', segmentdata=cdict, N=256)
     
     fig, ax = plt.subplots(figsize=(width/100, 2))
-    im = ax.imshow(combined_mse, aspect='auto', extent=[0, total_frames, 0, 2])
+    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'])
     ax.set_xticks([])