Update visualization.py

visualization.py

@@ -217,6 +217,72 @@ def plot_posture(df, posture_scores, color='blue', anomaly_threshold=3):
     plt.close()
     return fig
 
+
+def filter_mse_for_most_frequent_person(df, mse_embeddings, mse_posture, mse_voice, most_frequent_person_frames):
+    # Create a mask for the most frequent person frames
+    mask = df['Frame'].isin(most_frequent_person_frames)
+
+    # Apply the mask to filter the MSE arrays
+    mse_embeddings_filtered = np.where(mask, mse_embeddings, 0)
+    mse_posture_filtered = np.where(mask, mse_posture, 0)
+    mse_voice_filtered = np.where(mask, mse_voice, 0)
+
+    return mse_embeddings_filtered, mse_posture_filtered, mse_voice_filtered
+
+def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_voice, output_folder, desired_fps, most_frequent_person_frames):
+    print(f"Creating heatmap video. Output folder: {output_folder}")
+    
+    os.makedirs(output_folder, exist_ok=True)
+    
+    output_filename = os.path.basename(video_path).rsplit('.', 1)[0] + '_heatmap.mp4'
+    heatmap_video_path = os.path.join(output_folder, output_filename)
+    
+    print(f"Heatmap video will be saved at: {heatmap_video_path}")
+
+    # Load the original video
+    video = VideoFileClip(video_path)
+    
+    # Get video properties
+    width, height = video.w, video.h
+    total_frames = int(video.duration * video.fps)
+    
+    # Ensure all MSE arrays have the same length as total_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)
+    mse_voice = np.interp(np.linspace(0, len(mse_voice) - 1, total_frames), 
+                          np.arange(len(mse_voice)), mse_voice)
+    
+    # Filter MSE arrays for the most frequent person frames
+    mse_embeddings_filtered, mse_posture_filtered, mse_voice_filtered = filter_mse_for_most_frequent_person(df, mse_embeddings, mse_posture, mse_voice, most_frequent_person_frames)
+    
+    def combine_video_and_heatmap(t):
+        video_frame = video.get_frame(t)
+        heatmap_frame = create_heatmap(t, mse_embeddings_filtered, mse_posture_filtered, mse_voice_filtered, video.fps, total_frames, width)
+        heatmap_frame_resized = cv2.resize(heatmap_frame, (width, heatmap_frame.shape[0]))
+        combined_frame = np.vstack((video_frame, heatmap_frame_resized))
+        return combined_frame
+
+    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', fps=video.fps)
+    
+    # Close the video clips
+    video.close()
+    final_clip.close()
+    
+    if os.path.exists(heatmap_video_path):
+        print(f"Heatmap video created at: {heatmap_video_path}")
+        print(f"Heatmap video size: {os.path.getsize(heatmap_video_path)} bytes")
+        return heatmap_video_path
+    else:
+        print(f"Failed to create heatmap video at: {heatmap_video_path}")
+        return None
+
+# Define the create_heatmap function
 def create_heatmap(t, mse_embeddings, mse_posture, mse_voice, video_fps, total_frames, video_width):
     frame_count = int(t * video_fps)
     
@@ -246,31 +312,6 @@ def create_heatmap(t, mse_embeddings, mse_posture, mse_voice, video_fps, total_f
     heatmap_img = heatmap_img.reshape(canvas.get_width_height()[::-1] + (3,))
     plt.close(fig)
     return heatmap_img
-
-def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_voice, output_folder, desired_fps, largest_cluster):
-    print(f"Creating heatmap video. Output folder: {output_folder}")
-    
-    os.makedirs(output_folder, exist_ok=True)
-    
-    output_filename = os.path.basename(video_path).rsplit('.', 1)[0] + '_heatmap.mp4'
-    heatmap_video_path = os.path.join(output_folder, output_filename)
-    
-    print(f"Heatmap video will be saved at: {heatmap_video_path}")
-
-    # Load the original video
-    video = VideoFileClip(video_path)
-    
-    # Get video properties
-    width, height = video.w, video.h
-    total_frames = int(video.duration * video.fps)
-    
-    # Ensure all MSE arrays have the same length as total_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)
-    mse_voice = np.interp(np.linspace(0, len(mse_voice) - 1, total_frames), 
-                          np.arange(len(mse_voice)), mse_voice)
     
     def combine_video_and_heatmap(t):
         video_frame = video.get_frame(t)