Update visualization.py

visualization.py

@@ -217,10 +217,10 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, outpu
     total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
     
     fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-    out = cv2.VideoWriter(output_path, fourcc, desired_fps, (width, height + 200))
+    out = cv2.VideoWriter(output_path, fourcc, original_fps, (width, height + 200))
     
     cmap = mcolors.LinearSegmentedColormap.from_list("custom", 
-        [(1, 1, 1), (0, 0, 1), (0.5, 0, 0.5)], N=256)
+        [(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)
@@ -231,8 +231,8 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, outpu
     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, mse_embeddings_norm]).T
-    combined_mse[1] = np.array([1 - mse_posture_norm, mse_posture_norm, 1 - mse_posture_norm]).T
+    combined_mse[0] = np.array([1, 1 - mse_embeddings_norm, 1 - mse_embeddings_norm]).T  # More reddish for facial
+    combined_mse[1] = np.array([1, 1 - mse_posture_norm, 1 - mse_posture_norm]).T  # More reddish for posture
     
     fig, ax = plt.subplots(figsize=(width/100, 2))
     im = ax.imshow(combined_mse, aspect='auto', extent=[0, total_frames, 0, 2])
@@ -242,9 +242,7 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, outpu
     plt.tight_layout()
     
     line = None
-    frame_interval = int(original_fps / desired_fps)
-    
-    for frame_count in range(0, total_frames, frame_interval):
+    for frame_count in range(total_frames):
         cap.set(cv2.CAP_PROP_POS_FRAMES, frame_count)
         ret, frame = cap.read()
         if not ret:
@@ -273,4 +271,3 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, outpu
     plt.close(fig)
     
     return output_path
-