Update visualization.py

visualization.py

@@ -6,7 +6,7 @@ import seaborn as sns
 import numpy as np
 import pandas as pd
 import cv2
-from moviepy.editor import VideoFileClip, AudioFileClip, CompositeVideoClip
+from moviepy.editor import VideoFileClip, AudioFileClip, CompositeVideoClip, ImageClip
 from moviepy.video.fx.all import resize
 from matplotlib.patches import Rectangle
 from utils import seconds_to_timecode
@@ -267,11 +267,9 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
     ax.set_xticks([])
     plt.tight_layout()
 
-    def make_frame(frame):
-        t = frame / video.fps if isinstance(frame, np.ndarray) else frame
+    def create_heatmap(t):
         frame_count = int(t * video.fps)
         
-        # Update heatmap
         if hasattr(ax, 'line'):
             ax.lines.pop(0)
         ax.axvline(x=frame_count, color='r', linewidth=2)
@@ -280,28 +278,33 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
         canvas.draw()
         heatmap_img = np.frombuffer(canvas.tostring_rgb(), dtype='uint8')
         heatmap_img = heatmap_img.reshape(canvas.get_width_height()[::-1] + (3,))
-        heatmap_img = resize(ImageClip(heatmap_img), width=width).get_frame(0)
-        
-        # Combine frame with heatmap
-        combined_frame = np.vstack((frame, heatmap_img))
-        
-        # Add timecode
+        return heatmap_img
+
+    def add_timecode(frame, t):
         seconds = t
         timecode = f"{int(seconds//3600):02d}:{int((seconds%3600)//60):02d}:{int(seconds%60):02d}"
         
-        pil_img = Image.fromarray(combined_frame.astype('uint8'))
+        pil_img = Image.fromarray(frame.astype('uint8'))
         draw = ImageDraw.Draw(pil_img)
         font = ImageFont.load_default()
         draw.text((10, 30), f"Time: {timecode}", font=font, fill=(255, 255, 255))
         
         return np.array(pil_img)
 
-    # Create the final video
-    final_clip = video.fl_image(make_frame)
-    final_clip = final_clip.set_duration(video.duration)
+    heatmap_clip = VideoClip(create_heatmap, duration=video.duration)
+    heatmap_clip = heatmap_clip.resize(height=200)
+
+    def combine_video_and_heatmap(t):
+        video_frame = video.get_frame(t)
+        heatmap_frame = heatmap_clip.get_frame(t)
+        combined_frame = np.vstack((video_frame, heatmap_frame))
+        return add_timecode(combined_frame, t)
+
+    final_clip = VideoClip(combine_video_and_heatmap, duration=video.duration)
+    final_clip = final_clip.set_audio(video.audio)
 
     # Write the final video
-    final_clip.write_videofile(heatmap_video_path, codec='libx264', audio_codec='aac')
+    final_clip.write_videofile(heatmap_video_path, codec='libx264', audio_codec='aac', fps=video.fps)
     
     # Close the video clips
     video.close()