Update visualization.py

visualization.py

@@ -211,7 +211,7 @@ def fig_to_img(fig):
     plt.close(fig)
     return img
 
-def create_video_with_heatmap(video_path, mse_heatmap_embeddings_img, mse_heatmap_posture_img, output_path, desired_fps):
+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)
     width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
@@ -219,30 +219,63 @@ def create_video_with_heatmap(video_path, mse_heatmap_embeddings_img, mse_heatma
     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 + 400))
+    out = cv2.VideoWriter(output_path, fourcc, desired_fps, (width, height + 200))
     
-    # Resize heatmap images to match the width of the video frames
-    mse_heatmap_embeddings_img = cv2.resize(mse_heatmap_embeddings_img, (width, 200))
-    mse_heatmap_posture_img = cv2.resize(mse_heatmap_posture_img, (width, 200))
+    # Create custom colormap
+    cmap = mcolors.LinearSegmentedColormap.from_list("custom", 
+        [(1, 1, 1), (0, 0, 1), (0.5, 0, 0.5)], N=256)
     
-    for frame_count in range(0, total_frames):
+    # 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)
+    
+    # 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
+    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', extent=[0, total_frames, 0, 2])
+    ax.set_yticks([0.5, 1.5])
+    ax.set_yticklabels(['Face', 'Posture'])
+    ax.set_xticks([])
+    plt.tight_layout()
+    
+    line = None
+    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)
         ret, frame = cap.read()
         if not ret:
             break
         
-        heatmap_combined_img = np.vstack((mse_heatmap_embeddings_img, mse_heatmap_posture_img))
+        if line:
+            line.remove()
+        line = ax.axvline(x=frame_count, color='r', linewidth=2)
         
-        combined_frame = np.vstack((frame, heatmap_combined_img))
+        canvas = FigureCanvasAgg(fig)
+        canvas.draw()
+        heatmap_img = np.frombuffer(canvas.tostring_rgb(), dtype='uint8')
+        heatmap_img = heatmap_img.reshape(canvas.get_width_height()[::-1] + (3,))
+        heatmap_img = cv2.resize(heatmap_img, (width, 200))
+        
+        combined_frame = np.vstack((frame, heatmap_img))
         
         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, height + 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)
+        cv2.putText(combined_frame, f"Time: {timecode}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)
         
         out.write(combined_frame)
     
     cap.release()
     out.release()
-    plt.close('all')
+    plt.close(fig)
     
     return output_path