Update visualization.py

visualization.py

@@ -1,5 +1,6 @@
 import matplotlib.pyplot as plt
 from matplotlib.backends.backend_agg import FigureCanvasAgg
+import matplotlib.colors as mcolors
 import seaborn as sns
 import numpy as np
 import pandas as pd
@@ -203,30 +204,45 @@ def plot_posture(df, posture_scores, color='blue', anomaly_threshold=3):
     return fig
 
 
-def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, output_path):
+def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, output_path, desired_fps):
     cap = cv2.VideoCapture(video_path)
-    fps = cap.get(cv2.CAP_PROP_FPS)
+    original_fps = cap.get(cv2.CAP_PROP_FPS)
     width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
     height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
     total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
     
     fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height + 200))
+    out = cv2.VideoWriter(output_path, fourcc, desired_fps, (width, height + 200))
     
-    # Ensure heatmap data covers all frames
-    heatmap_data = np.zeros((2, total_frames))
-    heatmap_data[0, :len(mse_embeddings)] = mse_embeddings
-    heatmap_data[1, :len(mse_posture)] = mse_posture
+    # Create custom colormap
+    colors = ['navy', 'blue', 'white', 'purple', 'magenta']
+    n_bins = 100
+    cmap = mcolors.LinearSegmentedColormap.from_list('custom', colors, N=n_bins)
+    
+    # Ensure heatmap data covers all frames and normalize
+    mse_embeddings = np.pad(mse_embeddings, (0, max(0, total_frames - len(mse_embeddings))), 'constant', constant_values=np.nan)
+    mse_posture = np.pad(mse_posture, (0, max(0, total_frames - len(mse_posture))), 'constant', constant_values=np.nan)
+    
+    mse_embeddings_norm = (mse_embeddings - np.nanmin(mse_embeddings)) / (np.nanmax(mse_embeddings) - np.nanmin(mse_embeddings))
+    mse_posture_norm = (mse_posture - np.nanmin(mse_posture)) / (np.nanmax(mse_posture) - np.nanmin(mse_posture))
+    
+    # Combine MSEs: negative for facial features, positive for body posture
+    combined_mse = mse_posture_norm - mse_embeddings_norm
     
     fig, ax = plt.subplots(figsize=(width/100, 2))
-    im = ax.imshow(heatmap_data, aspect='auto', cmap='YlOrRd', extent=[0, total_frames, 0, 2])
+    im = ax.imshow(combined_mse.reshape(1, -1), aspect='auto', cmap=cmap, extent=[0, total_frames, 0, 1], vmin=-1, vmax=1)
     ax.set_yticks([])
     ax.set_xticks([])
+    cbar = plt.colorbar(im, ax=ax, orientation='horizontal', pad=0.1, aspect=20, shrink=0.5)
+    cbar.set_ticks([-1, 0, 1])
+    cbar.set_ticklabels(['High Facial MSE', 'Neutral', 'High Posture MSE'])
     plt.tight_layout()
     
     line = None
+    frame_interval = max(1, int(original_fps / desired_fps))
     
-    for frame_count in range(total_frames):
+    for frame_count in range(0, total_frames, frame_interval):
+        cap.set(cv2.CAP_PROP_POS_FRAMES, frame_count)
         ret, frame = cap.read()
         if not ret:
             break