Update visualization.py

visualization.py

@@ -210,7 +210,7 @@ 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):
+def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, output_path, desired_fps, largest_cluster):
     cap = cv2.VideoCapture(video_path)
     original_fps = cap.get(cv2.CAP_PROP_FPS)
     width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
@@ -222,35 +222,37 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, outpu
     
     # Create custom colormap
     cmap = mcolors.LinearSegmentedColormap.from_list("custom", 
-        [(0, 0, 1), (1, 1, 1), (0.5, 0, 0.5)], N=256)
+        [(1, 1, 1), (0, 0, 1), (0.5, 0, 0.5)], N=256)
+    
+    # Filter data for the most frequent person
+    df_most_frequent = df[df['Cluster'] == largest_cluster]
     
     # 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)
+    full_range = np.arange(total_frames)
+    mse_embeddings_full = np.zeros(total_frames)
+    mse_posture_full = np.zeros(total_frames)
+    
+    mse_embeddings_full[df_most_frequent['Frame']] = mse_embeddings
+    mse_posture_full[df_most_frequent['Frame']] = 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))
+    mse_embeddings_norm = (mse_embeddings_full - np.min(mse_embeddings_full)) / (np.max(mse_embeddings_full) - np.min(mse_embeddings_full))
+    mse_posture_norm = (mse_posture_full - np.min(mse_posture_full)) / (np.max(mse_posture_full) - np.min(mse_posture_full))
     
     # Combine MSEs
-    combined_mse = np.zeros((2, total_frames))
-    combined_mse[0] = mse_embeddings_norm
-    combined_mse[1] = mse_posture_norm
+    combined_mse = np.zeros((2, total_frames, 3))
+    combined_mse[0] = np.array([1 - mse_embeddings_norm, 1 - mse_embeddings_norm, mse_embeddings_norm]).T  # RGB for facial
+    combined_mse[1] = np.array([1 - mse_posture_norm, mse_posture_norm, 1 - mse_posture_norm]).T  # RGB for posture
     
     fig, ax = plt.subplots(figsize=(width/100, 2))
-    im = ax.imshow(combined_mse, aspect='auto', cmap=cmap, extent=[0, total_frames, 0, 2], vmin=0, vmax=1)
+    im = ax.imshow(combined_mse, aspect='auto', extent=[0, total_frames, 0, 2])
     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([0, 0.5, 1])
-    cbar.set_ticklabels(['Low MSE', 'Medium MSE', 'High MSE'])
     plt.tight_layout()
     
     line = None
-    frame_interval = max(1, int(original_fps / desired_fps))
+    frame_interval = int(original_fps / desired_fps)
     
     for frame_count in range(0, total_frames, frame_interval):
         cap.set(cv2.CAP_PROP_POS_FRAMES, frame_count)