Update visualization.py

visualization.py

@@ -204,6 +204,12 @@ def plot_posture(df, posture_scores, color='blue', anomaly_threshold=3):
     return fig
 
 
+import matplotlib.pyplot as plt
+from matplotlib.backends.backend_agg import FigureCanvasAgg
+import matplotlib.colors as mcolors
+import numpy as np
+import cv2
+
 def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, output_path, desired_fps):
     cap = cv2.VideoCapture(video_path)
     original_fps = cap.get(cv2.CAP_PROP_FPS)
@@ -215,27 +221,32 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, outpu
     out = cv2.VideoWriter(output_path, fourcc, desired_fps, (width, height + 200))
     
     # Create custom colormap
-    colors = ['navy', 'white', 'purple']  # Changed order: blue for facial, purple for posture
-    n_bins = 100
-    cmap = mcolors.LinearSegmentedColormap.from_list('custom', colors, N=n_bins)
+    cmap = mcolors.LinearSegmentedColormap.from_list("custom", 
+        [(0, 0, 1), (1, 1, 1), (0.5, 0, 0.5)], N=256)
     
-    # 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)
+    # Ensure heatmap data covers all frames
+    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), 
+                            np.arange(len(mse_posture)), mse_posture)
     
-    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))
+    # Normalize MSE values
+    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))
     
-    # Combine MSEs: negative for facial features (blue), positive for body posture (purple)
-    combined_mse = mse_posture_norm - mse_embeddings_norm
+    # Combine MSEs
+    combined_mse = np.zeros((2, total_frames))
+    combined_mse[0] = mse_embeddings_norm
+    combined_mse[1] = mse_posture_norm
     
     fig, ax = plt.subplots(figsize=(width/100, 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([])
+    im = ax.imshow(combined_mse, aspect='auto', cmap=cmap, extent=[0, total_frames, 0, 2], vmin=0, vmax=1)
+    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([-1, 0, 1])
-    cbar.set_ticklabels(['High Facial MSE', 'Neutral', 'High Posture MSE'])
+    cbar.set_ticks([0, 0.5, 1])
+    cbar.set_ticklabels(['Low MSE', 'Medium MSE', 'High MSE'])
     plt.tight_layout()
     
     line = None
@@ -259,7 +270,6 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, outpu
         
         combined_frame = np.vstack((frame, heatmap_img))
         
-        # Use frame_count to get timecode, not df index
         seconds = frame_count / original_fps
         timecode = f"{int(seconds//3600):02d}:{int((seconds%3600)//60):02d}:{int(seconds%60):02d}"
         cv2.putText(combined_frame, f"Time: {timecode}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)