Update visualization.py

visualization.py

@@ -217,12 +217,18 @@ 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):
     # Ensure most_frequent_person_frames is a list
     if not isinstance(most_frequent_person_frames, (list, np.ndarray)):
         most_frequent_person_frames = [most_frequent_person_frames]
 
+    # Ensure df and mse arrays have the same length
+    min_length = min(len(df), len(mse_embeddings), len(mse_posture), len(mse_voice))
+    df = df.iloc[:min_length].copy()
+    mse_embeddings = mse_embeddings[:min_length]
+    mse_posture = mse_posture[:min_length]
+    mse_voice = mse_voice[:min_length]
+
     # Create a mask for the most frequent person frames
     mask = df['Frame'].isin(most_frequent_person_frames)
 
@@ -233,7 +239,6 @@ def filter_mse_for_most_frequent_person(df, mse_embeddings, mse_posture, mse_voi
 
     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}")
     
@@ -287,7 +292,7 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
         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)
     
@@ -317,34 +322,7 @@ 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 combine_video_and_heatmap(t):
-        video_frame = video.get_frame(t)
-        heatmap_frame = create_heatmap(t, mse_embeddings, mse_posture, mse_voice, 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
-
 
-# Function to create the correlation heatmap
 def plot_correlation_heatmap(mse_embeddings, mse_posture, mse_voice):
     data = np.vstack((mse_embeddings, mse_posture, mse_voice)).T
     df = pd.DataFrame(data, columns=["Facial Features", "Body Posture", "Voice"])