Update visualization.py

visualization.py

@@ -215,7 +215,7 @@ 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', 'blue', 'white', 'purple', 'magenta']
+    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)
     
@@ -226,7 +226,7 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, outpu
     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
+    # Combine MSEs: negative for facial features (blue), positive for body posture (purple)
     combined_mse = mse_posture_norm - mse_embeddings_norm
     
     fig, ax = plt.subplots(figsize=(width/100, 2))
@@ -259,7 +259,9 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, outpu
         
         combined_frame = np.vstack((frame, heatmap_img))
         
-        timecode = df['Timecode'][frame_count] if frame_count < len(df) else "End"
+        # 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)
         
         out.write(combined_frame)