Update visualization.py

visualization.py

@@ -9,6 +9,7 @@ import pandas as pd
 import cv2
 from moviepy.editor import VideoFileClip, AudioFileClip, CompositeVideoClip, ImageClip, VideoClip, concatenate_videoclips
 from moviepy.video.fx.all import resize
+from moviepy.video.io.bindings import mplfig_to_npimage
 from PIL import Image, ImageDraw, ImageFont
 from matplotlib.patches import Rectangle
 from utils import seconds_to_timecode
@@ -267,4 +268,45 @@ def plot_stacked_mse_heatmaps(mse_face, mse_posture, mse_voice, df, title="Combi
     plt.suptitle(title)
     plt.tight_layout()
     plt.close()
-    return fig
+    return fig
+
+
+def create_video_with_heatmap(video_path, mse_face, mse_posture, mse_voice, df, output_path):
+    # Load the original video
+    video = VideoFileClip(video_path)
+    
+    # Create the stacked heatmap
+    fig = plot_stacked_mse_heatmaps(mse_face, mse_posture, mse_voice, df)
+    heatmap_img = mplfig_to_npimage(fig)
+    
+    # Resize heatmap to match video width
+    heatmap_height = int(video.h * 0.2)  # 20% of video height
+    heatmap_resized = cv2.resize(heatmap_img, (video.w, heatmap_height))
+    
+    def make_frame(t):
+        # Get the current frame from the original video
+        frame = video.get_frame(t)
+        
+        # Calculate the position of the vertical line
+        line_pos = int(t / video.duration * video.w)
+        
+        # Add the vertical line to the heatmap
+        heatmap_with_line = heatmap_resized.copy()
+        cv2.line(heatmap_with_line, (line_pos, 0), (line_pos, heatmap_height), (0, 0, 0), 2)
+        
+        # Combine the original frame with the heatmap
+        combined_frame = np.vstack((frame, heatmap_with_line))
+        
+        return combined_frame
+
+    # Create a new video clip with the combined frames
+    final_clip = VideoClip(make_frame, duration=video.duration)
+    
+    # Write the final video
+    final_clip.write_videofile(output_path, codec='libx264', fps=video.fps)
+    
+    # Close the clips
+    video.close()
+    final_clip.close()
+
+    return output_path